feat: Add the new model types we plan to introduce in Model Garden to…

… the "AutoMLImageTrainingJob" in SDK. PiperOrigin-RevId: 524909493
googleapis · Apr 17, 2023 · 4d032d5 · 4d032d5
1 parent b989dbb
commit 4d032d5
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 2 deletions.
diff --git a/google/cloud/aiplatform/constants/base.py b/google/cloud/aiplatform/constants/base.py
@@ -77,16 +77,30 @@
     "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
 #
 # 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
@@ -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.