Training classes

The Vertex AI SDK includes several classes that you use when youtrain your model. Most of the training classes are used to create, train, andreturn your model. Use theHyperparameterTuningJob to tune thetraining job's hyperparameters. Use thePipelineJob manage your machine learning (ML)workflow so you can automate and monitor your ML systems.

The following topics provide a high-level description of each training-relatedclass in the Vertex AI SDK.

AutoML training classes for structured data

Vertex AI SDK includes the following classes that are used to traina structured AutoML model.

AutoMLForecastingTrainingJob

TheAutoMLForecastingTrainingJobclass uses theAutoML training method to train and run a forecasting model.TheAutoML training method is a good choice for most forecasting use cases. Ifyour use case doesn't benefit from theSeq2seq or theTemporal fusion transformer training method that theSequenceToSequencePlusForecastingTrainingJobandTemporalFusionTransformerForecastingTrainingJobclasses offer respectively, thenAutoML is likely the best training method for your forecastingpredictions.

For sample code that shows you how to useAutoMLForecastingTrainingJob, see theCreate a training pipeline forecasting sample on GitHub.

AutoMLTabularTrainingJob

TheAutoMLTabularTrainingJob classrepresents a job that creates, trains, and returns anAutoML tabular model.For more information about training tabular models and Vertex AI, seeTabular data andTabular dataoverview.

The following sample code snippet shows how you might use theVertex AI SDK to create and run anAutoML tabular model:

dataset=aiplatform.TabularDataset('projects/my-project/location/us-central1/datasets/{DATASET_ID}')job=aiplatform.AutoMLTabularTrainingJob(display_name="train-automl",optimization_prediction_type="regression",optimization_objective="minimize-rmse",)model=job.run(dataset=dataset,target_column="target_column_name",training_fraction_split=0.6,validation_fraction_split=0.2,test_fraction_split=0.2,budget_milli_node_hours=1000,model_display_name="my-automl-model",disable_early_stopping=False,)

SequenceToSequencePlusForecastingTrainingJob

TheSequenceToSequencePlusForecastingTrainingJobclass uses theSeq2seq+ training method to train and run a forecasting model.TheSeq2seq+ training method is a good choice for experimentation. Itsalgorithm is simpler and uses a smaller search space than theAutoML option.Seq2seq+ is a good option if you want fast results and your datasets aresmaller than 1 GB.

For sample code that shows you how to useSequenceToSequencePlusForecastingTrainingJob, see theCreate a training pipeline forecasting Seq2seq sample on GitHub.

TemporalFusionTransformerForecastingTrainingJob

TheTemporalFusionTransformerForecastingTrainingJobclass uses the Temporal Fusion Transformer (TFT) training method to train andrun a forecasting model. The TFT training method implements an attention-baseddeep neural network (DNN) model that uses a multi-horizon forecasting task toproduce predictions.

For sample code that shows you how to useTemporalFusionTransformerForecastingTrainingJob,see theCreate a training pipeline forecasting temporal fusion transformersampleon GitHub.

TimeSeriesDenseEncoderForecastingTrainingJob

TheTimeSeriesDenseEncoderForecastingTrainingJobclass uses the Time-series Dense Encoder (TiDE) training method to train and runa forecasting model. TiDE uses amulti-layer perceptron (MLP) to provide thespeed of forecasting linear models with covariates and non-linear dependencies.For more information about TiDE, seeRecent advances in deep long-horizon forecastingand thisTiDE blog post.

AutoML training classes for unstructured data

The Vertex AI SDK includes the following classes to trainunstructured image models:

AutoMLImageTrainingJob

Use theAutoMLImageTrainingJob class tocreate, train, and return an image model. For more information about workingwith image data models in Vertex AI, seeImage data.

For an example of how to use theAutoMLImageTrainingJob class, see thetutorial in theAutoML imageclassificationnotebook.

Custom data training classes

You can use the Vertex AI SDK to automate a custom trainingworkflow. For information about using Vertex AI to run custom trainingapplications, seeCustom training overview.

The Vertex AI SDK includes three classes that create a customtraining pipeline. A training pipeline accepts an input Vertex AI manageddataset that it uses to train a model. Next, it returns the model after thetraining job completes. Each of the three custom training pipeline classescreates a training pipeline differently.CustomTrainingJob uses a Python script,CustomContainerTrainingJob uses acustom container, andCustomPythonPackageTrainingJobuses a Python package and a prebuilt container.

TheCustomJob class creates a custom training jobbut is not a pipeline. Unlike a custom training pipeline, theCustomJob class can use a dataset that's not aVertex AI managed dataset to train a model, and it doesn't return thetrained model. Because the class accepts different types of datasets and doesn'treturn a trained model, it's less automated and more flexible than a customtraining pipeline.

CustomContainerTrainingJob

Use theCustomContainerTrainingJobclass to use a container to launch a custom training pipeline in Vertex AI.

For an example of how to use theCustomContainerTrainingJob class,see the tutorial in thePyTorch Image Classification Multi-Node DistributedData Parallel Training on GPU using Vertex AI Training with Custom Containernotebook.

CustomJob

Use theCustomJob class to use a script to launch acustom training job in Vertex AI.

A training job is more flexible than a training pipeline because you aren'trestricted to loading your data in a Vertex AI managed dataset and areference to your model isn't registered after the training job completes. Forexample, you might want to use theCustomJob class, itsfrom_local_scriptmethod, and a script to load a dataset fromscikit-learn orTensorFlow. Or, you might want to analyze or testyour trained model before you register it to Vertex AI.

For more information about custom training jobs, including requirements beforesubmitting a custom training job, what a custom job includes, and a Python codesample, seeCreate custom training jobs.

Because theCustomJob.rundoesn't return the trained model, you need to use a script to write the modelartifact to a location, such as a Cloud Storage bucket. For more information,seeExport a trained ML model.

The following sample code demonstrates how to create and run a custom job usinga sample worker pool specification. The code writes the trained model to aCloud Storage bucket namedartifact-bucket.

#CreateaworkerpoolspecthatspecifiesaTensorFlowcassavadatasetand#includesthemachinetypeandDockerimage.The Google CloudprojectID#is'project-id'.worker_pool_specs=[{"replica_count":1,"machine_spec":{"machine_type":"n1-standard-8","accelerator_type":"NVIDIA_TESLA_V100","accelerator_count":1},"container_spec":{"image_uri":"gcr.io/{project-id}/multiworker:cassava"}},{"replica_count":1,"machine_spec":{"machine_type":"n1-standard-8","accelerator_type":"NVIDIA_TESLA_V100","accelerator_count":1},"container_spec":{"image_uri":"gcr.io/{project-id}/multiworker:cassava"}}]#Usetheworkerpoolspectocreateacustomtrainingjob.Thecustomtraining#jobartifactsarestoredintheCloudStoragebucket#named'artifact-bucket'.your_custom_training_job=aiplatform.CustomJob(display_name='multiworker-cassava-sdk',worker_pool_specs=worker_pool_specs,staging_bucket='gs://{artifact-bucket}')#Runthetrainingjob.Thismethoddoesn'treturnthetrainedmodel.my_multiworker_job.run()

CustomPythonPackageTrainingJob

Use theCustomPythonPackageTrainingJobclass to use a Python package to launch a custom training pipeline inVertex AI.

For an example of how to use theCustomPythonPackageTrainingJobclass, see the tutorial in theCustom training using Python package, managedtext dataset, and TensorFlow servingcontainernotebook.

CustomTrainingJob

Use theCustomTrainingJob class to launch acustom training pipeline in Vertex AI with a script.

For an example of how to use theCustomTrainingJob class, see the tutorial intheCustom training image classification model for online prediction withexplainabilitynotebook.

Hyperparameter training class

The Vertex AI SDK includes a class for hyperparameter tuning.Hyperparameter tuning maximizes your model's predictive accuracy by optimizingvariables (known ashyperparameters) that govern the training process. Formore information, seeOverview of hyperparametertuning.

HyperparameterTuningJob

Use theHyperparameterTuningJob classto automate hyperparameter tuning on a training application.

To learn how to use theHyperparameterTuningJob class to create and tune a custom trained model, see theHyperparameter tuning tutorial on GitHub.

To learn how to use theHyperparameterTuningJob class to run a Vertex AI hyperparameter tuning job for a TensorFlow model, see theRun hyperparameter tuning for a TensorFlow model tutorial on GitHub.

Pipeline training class

A pipeline orchestrates your ML workflow in Vertex AI. You can use a pipeline to automate, monitor, and govern your machine learning systems. To learn more about pipelines in Vertex AI, seeIntroduction to Vertex AI pipelines.

PipelineJob

An instance of thePipelineJob class represents a Vertex AI pipeline.

There are several tutorial notebooks that demonstrate how to use thePipelineJob class:

For more tutorial notebooks, seeVertex AI notebook tutorials.

What's next

Except as otherwise noted, the content of this page is licensed under theCreative Commons Attribution 4.0 License, and code samples are licensed under theApache 2.0 License. For details, see theGoogle Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates.

Last updated 2025-12-15 UTC.