From 4d032d59840fdc812dbbedb45acdf45c259e50bd Mon Sep 17 00:00:00 2001
From: A Vertex SDK engineer <vertex-sdk-bot@google.com>
Date: Mon, 17 Apr 2023 11:55:19 -0700
Subject: [PATCH] feat: Add the new model types we plan to introduce in Model
 Garden to the "AutoMLImageTrainingJob" in SDK.

PiperOrigin-RevId: 524909493
---
 google/cloud/aiplatform/constants/base.py | 18 ++++++++++++++++--
 google/cloud/aiplatform/training_jobs.py  | 15 +++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/google/cloud/aiplatform/constants/base.py b/google/cloud/aiplatform/constants/base.py
index 008119c9ac..72fb75edcf 100644
--- a/google/cloud/aiplatform/constants/base.py
+++ b/google/cloud/aiplatform/constants/base.py
@@ -77,6 +77,17 @@
     "MOBILE_TF_HIGH_ACCURACY_1",
 }
 
+MODEL_GARDEN_ICN_MODEL_TYPES = {
+    "EFFICIENTNET",
+    "RESNET",
+    "VIT",
+    "COCA",
+}
+
+MODEL_GARDEN_IOD_MODEL_TYPES = {
+    "SPINENET",
+}
+
 # TODO(b/177079208): Use EPCL Enums for validating Model Types
 # Defined by gs://google-cloud-aiplatform/schema/trainingjob/definition/automl_image_*
 # Format: "prediction_type": set() of model_type's
@@ -84,9 +95,12 @@
 # NOTE: When adding a new prediction_type's, ensure it fits the pattern
 #       "automl_image_{prediction_type}_*" used by the YAML schemas on GCS
 AUTOML_IMAGE_PREDICTION_MODEL_TYPES = {
-    "classification": {"CLOUD", "CLOUD_1"} | MOBILE_TF_MODEL_TYPES,
+    "classification": {"CLOUD", "CLOUD_1"}
+    | MOBILE_TF_MODEL_TYPES
+    | MODEL_GARDEN_ICN_MODEL_TYPES,
     "object_detection": {"CLOUD_1", "CLOUD_HIGH_ACCURACY_1", "CLOUD_LOW_LATENCY_1"}
-    | MOBILE_TF_MODEL_TYPES,
+    | MOBILE_TF_MODEL_TYPES
+    | MODEL_GARDEN_IOD_MODEL_TYPES,
 }
 
 AUTOML_VIDEO_PREDICTION_MODEL_TYPES = {
diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py
index 99a4f0e2b9..ab91ee8f0e 100644
--- a/google/cloud/aiplatform/training_jobs.py
+++ b/google/cloud/aiplatform/training_jobs.py
@@ -5338,6 +5338,21 @@ def __init__(
                         or Core ML model and used on a mobile or edge device afterwards.
                         Expected to have a higher latency, but should also have a higher
                         prediction quality than other mobile models.
+                    "EFFICIENTNET" - A model that, available in Vertex Model Garden image
+                        classification training with customizable hyperparameters. Best tailored
+                        to be used within Google Cloud, and cannot be exported externally.
+                    "RESNET" - A model that, available in Vertex Model Garden image
+                        classification training with customizable hyperparameters. Best tailored
+                        to be used within Google Cloud, and cannot be exported externally.
+                    "VIT" - A model that, available in Vertex Model Garden image
+                        classification training with customizable hyperparameters. Best tailored
+                        to be used within Google Cloud, and cannot be exported externally.
+                    "COCA" - A model that, available in Vertex Model Garden image
+                        classification training with customizable hyperparameters. Best tailored
+                        to be used within Google Cloud, and cannot be exported externally.
+                    "SPINENET" - A model that, available in Vertex Model Garden image object
+                        detection training with customizable hyperparameters. Best tailored
+                        to be used within Google Cloud, and cannot be exported externally.
             base_model: Optional[models.Model] = None
                 Optional. Only permitted for Image Classification models.
                 If it is specified, the new model will be trained based on the `base` model.