Training Models on Vertex AI Intermediate

Vertex AI offers two primary approaches to model training: AutoML for no-code/low-code training, and Custom Training for full control over your training code. This lesson covers both approaches, along with TPU acceleration and hyperparameter tuning.

AutoML Training

AutoML automatically searches for the best model architecture and hyperparameters for your dataset. It supports image classification, object detection, text classification, tabular classification/regression, video classification, and forecasting.

Python

from google.cloud import aiplatform

aiplatform.init(project="my-project", location="us-central1")

# Create a tabular dataset
dataset = aiplatform.TabularDataset.create(
    display_name="my-dataset",
    gcs_source="gs://my-bucket/data.csv"
)

# Train an AutoML model
job = aiplatform.AutoMLTabularTrainingJob(
    display_name="my-automl-job",
    optimization_prediction_type="classification",
    optimization_objective="maximize-au-roc"
)

model = job.run(
    dataset=dataset,
    target_column="target",
    training_fraction_split=0.8,
    validation_fraction_split=0.1,
    test_fraction_split=0.1,
    budget_milli_node_hours=1000
)

Custom Training

For full control, use custom training with your own code. Vertex AI provides pre-built containers for popular frameworks and supports custom Docker containers for any framework.

Using Pre-built Containers

Python

# Define a custom training job
job = aiplatform.CustomTrainingJob(
    display_name="my-custom-job",
    script_path="trainer/task.py",
    container_uri="us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-12:latest",
    requirements=["pandas", "scikit-learn"],
    model_serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-12:latest"
)

# Run the training job with GPU
model = job.run(
    replica_count=1,
    machine_type="n1-standard-8",
    accelerator_type="NVIDIA_TESLA_V100",
    accelerator_count=1,
    args=["--epochs=50", "--batch-size=32"]
)

Using Custom Containers

Python

# Use your own Docker image
job = aiplatform.CustomContainerTrainingJob(
    display_name="my-container-job",
    container_uri="gcr.io/my-project/my-trainer:latest",
    model_serving_container_image_uri="gcr.io/my-project/my-predictor:latest"
)

model = job.run(
    replica_count=1,
    machine_type="n1-standard-16",
    accelerator_type="NVIDIA_TESLA_T4",
    accelerator_count=2
)

TPU Training

Vertex AI supports Cloud TPUs for accelerated training of large models, especially beneficial for TensorFlow and JAX workloads:

Python

# TPU training job
job = aiplatform.CustomTrainingJob(
    display_name="tpu-training-job",
    script_path="trainer/tpu_task.py",
    container_uri="us-docker.pkg.dev/vertex-ai/training/tf-tpu.2-12:latest"
)

model = job.run(
    replica_count=1,
    machine_type="cloud-tpu",
    accelerator_type="TPU_V3",
    accelerator_count=8
)

Hyperparameter Tuning

Vertex AI provides managed hyperparameter tuning that automatically searches for the best hyperparameters using Bayesian optimization, grid search, or random search:

Python

from google.cloud.aiplatform import hyperparameter_tuning as hpt

# Define the hyperparameter tuning job
hp_job = aiplatform.HyperparameterTuningJob(
    display_name="hp-tuning-job",
    custom_job=my_custom_job,
    metric_spec={"accuracy": "maximize"},
    parameter_spec={
        "learning_rate": hpt.DoubleParameterSpec(min=0.001, max=0.1, scale="log"),
        "batch_size": hpt.DiscreteParameterSpec(values=[16, 32, 64, 128], scale="linear"),
        "num_layers": hpt.IntegerParameterSpec(min=2, max=10, scale="linear"),
    },
    max_trial_count=20,
    parallel_trial_count=5,
    search_algorithm="bayesian"
)

hp_job.run()

Pro Tip: Start with a small number of trials and a wide parameter range, then narrow down the range and increase trials based on initial results. This approach saves both time and compute costs.

Distributed Training

For large datasets and models, Vertex AI supports multi-worker distributed training:

Strategy	Use Case	Framework Support
Data Parallelism	Large datasets, standard model sizes	TensorFlow, PyTorch
Model Parallelism	Very large models that don't fit on one GPU	PyTorch (FSDP), DeepSpeed
Reduction Server	Optimized all-reduce for multi-GPU training	TensorFlow, PyTorch

Models Trained!

Now that you know how to train models on Vertex AI, the next lesson covers deploying those models to production endpoints.

Next: Deployment →

← Setup Deployment →

Ready to Go Deeper?

Live instructor-led courses from our partners. Affiliate disclosure.

🧠

AI & ML Courses - 30% Off

Live instructor-led AI, machine learning, data science, and cloud courses for working professionals. Use code Limited30 at checkout.

Edureka

📊

DataCamp - AI & Data Science

Hands-on Python, machine learning, and AI courses with interactive exercises and real projects.

DataCamp

🎓

edX - Top AI Courses

University-level AI courses from MIT, Harvard, Stanford. Earn certificates that employers recognize.

edX