feat: Add tunable parameters for Model Garden model training to the "AutoMLImageTrainingJob" in SDK.

PiperOrigin-RevId: 538999389

Author: vertex-sdk-bot

google/cloud/aiplatform/constants/base.py

@@ -79,12 +79,14 @@
 
 MODEL_GARDEN_ICN_MODEL_TYPES = {
     "EFFICIENTNET",
+    "MAXVIT",
     "VIT",
     "COCA",
 }
 
 MODEL_GARDEN_IOD_MODEL_TYPES = {
     "SPINENET",
+    "YOLO",
 }
 
 # TODO(b/177079208): Use EPCL Enums for validating Model Types

google/cloud/aiplatform/hyperparameter_tuning.py

@@ -22,6 +22,20 @@
 
 from google.cloud.aiplatform.compat.types import study as gca_study_compat
 
+SEARCH_ALGORITHM_TO_PROTO_VALUE = {
+    "random": gca_study_compat.StudySpec.Algorithm.RANDOM_SEARCH,
+    "grid": gca_study_compat.StudySpec.Algorithm.GRID_SEARCH,
+    None: gca_study_compat.StudySpec.Algorithm.ALGORITHM_UNSPECIFIED,
+}
+
+MEASUREMENT_SELECTION_TO_PROTO_VALUE = {
+    "best": (gca_study_compat.StudySpec.MeasurementSelectionType.BEST_MEASUREMENT),
+    "last": (gca_study_compat.StudySpec.MeasurementSelectionType.LAST_MEASUREMENT),
+    None: (
+        gca_study_compat.StudySpec.MeasurementSelectionType.MEASUREMENT_SELECTION_TYPE_UNSPECIFIED
+    ),
+}
+
 _SCALE_TYPE_MAP = {
     "linear": gca_study_compat.StudySpec.ParameterSpec.ScaleType.UNIT_LINEAR_SCALE,
     "log": gca_study_compat.StudySpec.ParameterSpec.ScaleType.UNIT_LOG_SCALE,

google/cloud/aiplatform/jobs.py

@@ -1998,18 +1998,6 @@ def job_spec(self):
         return self._gca_resource.job_spec
 
 
-_SEARCH_ALGORITHM_TO_PROTO_VALUE = {
-    "random": gca_study_compat.StudySpec.Algorithm.RANDOM_SEARCH,
-    "grid": gca_study_compat.StudySpec.Algorithm.GRID_SEARCH,
-    None: gca_study_compat.StudySpec.Algorithm.ALGORITHM_UNSPECIFIED,
-}
-
-_MEASUREMENT_SELECTION_TO_PROTO_VALUE = {
-    "best": gca_study_compat.StudySpec.MeasurementSelectionType.BEST_MEASUREMENT,
-    "last": gca_study_compat.StudySpec.MeasurementSelectionType.LAST_MEASUREMENT,
-}
-
-
 class HyperparameterTuningJob(_RunnableJob):
     """Vertex AI Hyperparameter Tuning Job."""
 
@@ -2215,8 +2203,10 @@ def __init__(
         study_spec = gca_study_compat.StudySpec(
             metrics=metrics,
             parameters=parameters,
-            algorithm=_SEARCH_ALGORITHM_TO_PROTO_VALUE[search_algorithm],
-            measurement_selection_type=_MEASUREMENT_SELECTION_TO_PROTO_VALUE[
+            algorithm=hyperparameter_tuning.SEARCH_ALGORITHM_TO_PROTO_VALUE[
+                search_algorithm
+            ],
+            measurement_selection_type=hyperparameter_tuning.MEASUREMENT_SELECTION_TO_PROTO_VALUE[
                 measurement_selection
             ],
         )