feat: Adds the Time series Dense Encoder (TiDE) forecasting job.

vertex-sdk-bot · copybara-github · commit d8e67446dedd · 2023-04-13T12:14:42.000-07:00
PiperOrigin-RevId: 524068121
diff --git a/google/cloud/aiplatform/__init__.py b/google/cloud/aiplatform/__init__.py
@@ -70,6 +70,7 @@
     AutoMLForecastingTrainingJob,
     SequenceToSequencePlusForecastingTrainingJob,
     TemporalFusionTransformerForecastingTrainingJob,
+    TimeSeriesDenseEncoderForecastingTrainingJob,
     AutoMLImageTrainingJob,
     AutoMLTextTrainingJob,
     AutoMLVideoTrainingJob,
@@ -178,5 +179,6 @@
     "TextDataset",
     "TemporalFusionTransformerForecastingTrainingJob",
     "TimeSeriesDataset",
+    "TimeSeriesDenseEncoderForecastingTrainingJob",
     "VideoDataset",
 )
diff --git a/google/cloud/aiplatform/schema.py b/google/cloud/aiplatform/schema.py
@@ -25,6 +25,7 @@ class definition:
         automl_forecasting = "gs://google-cloud-aiplatform/schema/trainingjob/definition/automl_time_series_forecasting_1.0.0.yaml"
         seq2seq_plus_forecasting = "gs://google-cloud-aiplatform/schema/trainingjob/definition/seq2seq_plus_time_series_forecasting_1.0.0.yaml"
         tft_forecasting = "gs://google-cloud-aiplatform/schema/trainingjob/definition/temporal_fusion_transformer_time_series_forecasting_1.0.0.yaml"
+        tide_forecasting = "gs://google-cloud-aiplatform/schema/trainingjob/definition/time_series_dense_encoder_forecasting_1.0.0.yaml"
         automl_image_classification = "gs://google-cloud-aiplatform/schema/trainingjob/definition/automl_image_classification_1.0.0.yaml"
         automl_image_object_detection = "gs://google-cloud-aiplatform/schema/trainingjob/definition/automl_image_object_detection_1.0.0.yaml"
         automl_text_classification = "gs://google-cloud-aiplatform/schema/trainingjob/definition/automl_text_classification_1.0.0.yaml"
diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py
@@ -5257,6 +5257,14 @@ class TemporalFusionTransformerForecastingTrainingJob(_ForecastingTrainingJob):
     _supported_training_schemas = (schema.training_job.definition.tft_forecasting,)
 
 
+class TimeSeriesDenseEncoderForecastingTrainingJob(_ForecastingTrainingJob):
+    """Class to train Time series Dense Encoder (TiDE) forecasting models."""
+
+    _model_type = "TiDE"
+    _training_task_definition = schema.training_job.definition.tide_forecasting
+    _supported_training_schemas = (schema.training_job.definition.tide_forecasting,)
+
+
 class AutoMLImageTrainingJob(_TrainingJob):
     _supported_training_schemas = (
         schema.training_job.definition.automl_image_classification,
diff --git a/tests/system/aiplatform/test_e2e_forecasting.py b/tests/system/aiplatform/test_e2e_forecasting.py
@@ -42,10 +42,8 @@ class TestEndToEndForecasting(e2e_base.TestEndToEnd):
         [
             training_jobs.AutoMLForecastingTrainingJob,
             training_jobs.SequenceToSequencePlusForecastingTrainingJob,
-            pytest.param(
-                training_jobs.TemporalFusionTransformerForecastingTrainingJob,
-                marks=pytest.mark.skip(reason="TFT not yet released."),
-            ),
+            training_jobs.TemporalFusionTransformerForecastingTrainingJob,
+            training_jobs.TimeSeriesDenseEncoderForecastingTrainingJob,
         ],
     )
     def test_end_to_end_forecasting(self, shared_state, training_job):
diff --git a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py
@@ -187,6 +187,7 @@
     training_jobs.AutoMLForecastingTrainingJob,
     training_jobs.SequenceToSequencePlusForecastingTrainingJob,
     training_jobs.TemporalFusionTransformerForecastingTrainingJob,
+    training_jobs.TimeSeriesDenseEncoderForecastingTrainingJob,
 ]
 
 

Original file line number	Diff line number	Diff line change
`@@ -42,10 +42,8 @@ class TestEndToEndForecasting(e2e_base.TestEndToEnd):`
`42`	`42`	`[`
`43`	`43`	`training_jobs.AutoMLForecastingTrainingJob,`
`44`	`44`	`training_jobs.SequenceToSequencePlusForecastingTrainingJob,`
`45`		`- pytest.param(`
`46`		`- training_jobs.TemporalFusionTransformerForecastingTrainingJob,`
`47`		`- marks=pytest.mark.skip(reason="TFT not yet released."),`
`48`		`- ),`
	`45`	`+ training_jobs.TemporalFusionTransformerForecastingTrainingJob,`
	`46`	`+ training_jobs.TimeSeriesDenseEncoderForecastingTrainingJob,`
`49`	`47`	`],`
`50`	`48`	`)`
`51`	`49`	`def test_end_to_end_forecasting(self, shared_state, training_job):`
Original file line number	Diff line number	Diff line change
`@@ -187,6 +187,7 @@`
`187`	`187`	`training_jobs.AutoMLForecastingTrainingJob,`
`188`	`188`	`training_jobs.SequenceToSequencePlusForecastingTrainingJob,`
`189`	`189`	`training_jobs.TemporalFusionTransformerForecastingTrainingJob,`
	`190`	`+ training_jobs.TimeSeriesDenseEncoderForecastingTrainingJob,`
`190`	`191`	`]`
`191`	`192`
`192`	`193`