Prototype to Production: Training custom models with Vertex AI

1. Overview

In this lab, you'll use Vertex AI to run a custom training job.

This lab is part of the Prototype to Production video series. You'll build an image classification model using the Flowers dataset. You can watch the accompanying video to learn more:


What you learn

You'll learn how to:

  • Create a Vertex AI Workbench managed notebook
  • Configure and launch a custom training job from the Vertex AI UI
  • Configure and launch a custom training job with the Vertex AI Python SDK

The total cost to run this lab on Google Cloud is about $1.

2. Intro to Vertex AI

This lab uses the newest AI product offering available on Google Cloud. Vertex AI integrates the ML offerings across Google Cloud into a seamless development experience. Previously, models trained with AutoML and custom models were accessible via separate services. The new offering combines both into a single API, along with other new products. You can also migrate existing projects to Vertex AI.

Vertex AI includes many different products to support end-to-end ML workflows. This lab will focus on the products highlighted below: Training and Workbench

Vertex product overview

3. Set up your environment

You'll need a Google Cloud Platform project with billing enabled to run this codelab. To create a project, follow the instructions here.

Step 1: Enable the Compute Engine API

Navigate to Compute Engine and select Enable if it isn't already enabled.

Step 2: Enable the Artifact Registry API

Navigate to Artifact Registry and select Enable if it isn't already. You'll use this to create a container for your custom training job.

Step 3: Enable the Vertex AI API

Navigate to the Vertex AI section of your Cloud Console and click Enable Vertex AI API.

Vertex AI dashboard

Step 4: Create a Vertex AI Workbench instance

From the Vertex AI section of your Cloud Console, click on Workbench:

Vertex AI menu

Enable the Notebooks API if it isn't already.


Once enabled, click MANAGED NOTEBOOKS:


Then select NEW NOTEBOOK.


Give your notebook a name, and under Permission select Service account


Select Advanced Settings.

Under Security select "Enable terminal" if it is not already enabled.


You can leave all of the other advanced settings as is.

Next, click Create. The instance will take a couple minutes to be provisioned.

Once the instance has been created, select OPEN JUPYTERLAB.


4. Containerize training application code

You'll submit this training job to Vertex AI by putting your training application code in a Docker container and pushing this container to Google Artifact Registry. Using this approach, you can train a model built with any framework.

To start, from the Launcher menu, open a Terminal window in your notebook instance:

Open terminal in notebook

Step 1: Create a Cloud Storage bucket

In this training job, you'll export the trained TensorFlow model to a Cloud Storage Bucket. You'll also store the data for training in a Cloud Storage bucket.

From your Terminal, run the following to define an env variable for your project, making sure to replace your-cloud-project with the ID of your project:


Next, run the following in your Terminal to create a new bucket in your project.

gsutil mb -l us-central1 $BUCKET

Step 2: Copy data to Cloud Storage bucket

We need to get our flowers dataset into Cloud Storage. For demonstration purposes, you'll first download the dataset to this Workbench instance, and then copy it to a bucket.

Download and untar the data.

tar xvzf flower_photos.tgz

Then copy it to the bucket you just created. We add the -r because we want to copy the entire directory, and -m to perform a multi-processing copy, which will speed things up.

gsutil -m cp -r flower_photos $BUCKET

Step 3: Write training code

Create a new directory called flowers and cd into it:

mkdir flowers
cd flowers

Run the following to create a directory for the training code and a Python file where you'll add the code.

mkdir trainer
touch trainer/

You should now have the following in your flowers/ directory:

+ trainer/

For more details on how to structure your training application code, check out the docs.

Next, open the file you just created and copy the code below.

You'll need to replace {your-gcs-bucket} with the name of the Cloud Storage bucket you just created.

Through the Cloud Storage FUSE tool, training jobs on Vertex AI Training can access data on Cloud Storage as files in the local file system. When you start a custom training job, the job sees a directory /gcs which contains all your Cloud Storage buckets as subdirectories. That's why the data paths in the training code start with /gcs.

import tensorflow as tf
import numpy as np
import os

## Replace {your-gcs-bucket} !!

# Define variables


DATA_DIR = f'{BUCKET_ROOT}/flower_photos'

def create_datasets(data_dir, batch_size):
  '''Creates train and validation datasets.'''
  train_dataset = tf.keras.utils.image_dataset_from_directory(
    image_size=(IMG_HEIGHT, IMG_WIDTH),

  validation_dataset = tf.keras.utils.image_dataset_from_directory(
    image_size=(IMG_HEIGHT, IMG_WIDTH),

  train_dataset = train_dataset.cache().shuffle(1000).prefetch(
  validation_dataset = validation_dataset.cache().prefetch(

  return train_dataset, validation_dataset

def create_model():
  '''Creates model.'''

  model = tf.keras.Sequential([
    tf.keras.layers.Resizing(IMG_HEIGHT, IMG_WIDTH),
    tf.keras.layers.Rescaling(1./255, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3)),
    tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')
  return model

train_dataset, validation_dataset = create_datasets(DATA_DIR, BATCH_SIZE)

model = create_model()

history =

# SAVE MODEL'{BUCKET_ROOT}/model_output')

Step 4: Create a Dockerfile

To containerize your code, you'll need to create a Dockerfile. In the Dockerfile you'll include all the commands needed to run the image. It'll install all the necessary libraries and set up the entry point for the training code.

From your Terminal, create an empty Dockerfile in the root of your flowers directory:

touch Dockerfile

You should now have the following in your flowers/ directory:

+ Dockerfile
+ trainer/

Open the Dockerfile and copy the following into it:



# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

Let's review the commands in this file.

The FROM command specifies the base image, which is the parent image that the image you create will be built on. As the base image, you'll use the use the Deep Learning Container TensorFlow Enterprise 2.8 GPU Docker image. The Deep Learning Containers on Google Cloud come with many common ML and data science frameworks pre-installed.

The WORKDIR command specifies the directory on the image where subsequent instructions are run.

The COPY command copies the trainer code to the Docker image. Note that in this example we only have one python file in our trainer directory, but for a more realistic example you would likely have additional files. Maybe one called, which handles data preprocessing, and one called, which just has the model code, etc. For more complex training code, check out the Python docs on packaging Python projects.

If you wanted to add any additional libraries, you could use the RUN command to pip install (ex: RUN pip install -r requirements.txt). But we don't need anything additional for our example here.

Lastly, the ENTRYPOINT command sets up the entry point to invoke the trainer. This is what will run when we start our training job. In our case that is executing the file.

You can learn more about writing Dockerfiles for Vertex AI Training here.

Step 4: Build the container

From the terminal of your Workbench notebook, run the following to define an env variable for your project, making sure to replace your-cloud-project with the ID of your project:


Create a repo in Artifact Registry


gcloud artifacts repositories create $REPO_NAME --repository-format=docker \
--location=us-central1 --description="Docker repository"

Define a variable with the URI of your container image in Google Artifact Registry:$PROJECT_ID/$REPO_NAME/flower_image:latest

Configure docker

gcloud auth configure-docker \

Then, build the container by running the following from the root of your flower directory:

docker build ./ -t $IMAGE_URI

Lastly, push it to Artifact Registry:

docker push $IMAGE_URI

With the container pushed to Artifact Registry, you're now ready to kick off a the training job.

5. Run a custom training job on Vertex AI

This lab uses custom training via a custom container on Google Artifact Registry, but you can also run a training job with the Pre-built containers.

To start, navigate to the Training section in the Vertex section of your Cloud console:

train menu

Step 1: Configure training job

Click Create to enter the parameters for your training job.


  • Under Dataset, select No managed dataset
  • Then select Custom training (advanced) as your training method and click Continue.
  • Select Train new model then enter flowers-model (or whatever you'd like to call your model) for Model name
  • Click Continue

In the Container settings step, select Custom container:

Custom container option

In the first box (Container image), enter the value of your IMAGE_URI variable from the previous section. It should be:{PROJECT_ID}/flower-app/flower_image:latest, with your own project ID. Leave the rest of the fields blank and click Continue.

Skip the Hyperparameters step by clicking Continue again.

Step 2: Configure compute cluster

Configure Worker pool 0 as follows:


You'll skip step 6 for now and configure the prediction container in the next lab in this series.

Click START TRAINING to kick off the training job. In the Training section of your console under the TRAINING PIPELINES tab you'll see your newly launched job:

Training jobs

🎉 Congratulations! 🎉

You've learned how to use Vertex AI to:

  • Launch a custom training job for training code provided in a custom container. You used a TensorFlow model in this example, but you can train a model built with any framework using custom or built-in containers.

To learn more about different parts of Vertex, check out the documentation.

6. [Optional] Use the Vertex AI Python SDK

The previous section showed how to launch the training job via the UI. In this section, you'll see an alternative way to submit the training job by using the Vertex AI Python SDK.

Return to your notebook instance, and create a create a TensorFlow 2 notebook from the Launcher:


Import the Vertex AI SDK.

from import aiplatform

Then create a CustomContainerTrainingJob. You'll need to replace {PROJECT_ID} in the container_uri with the name of your project, and you'll need to replace {BUCKET} in staging_bucket with the bucket you created earlier.

my_job = aiplatform.CustomContainerTrainingJob(display_name='flower-sdk-job',

Then, run the job.,

For demonstration purposes, this job has been configured to run on a larger machine than compared to the previous section. Additionally, we are running with a GPU. If you don't specify the machine-type, accelerator_type, or accelerator_count, the job will run by default on an n1-standard-4.

In the Training section of your console under the CUSTOM JOBS tab you'll see your training job.

7. Cleanup

Because Vertex AI Workbench managed notebooks have an idle shutdown feature, we don't need to worry about shutting the instance down. If you would like to manually shut down the instance, click the Stop button on the Vertex AI Workbench section of the console. If you'd like to delete the notebook entirely, click the Delete button.

Stop instance

To delete the Storage Bucket, using the Navigation menu in your Cloud Console, browse to Storage, select your bucket, and click Delete:

Delete storage