[AutoML Image] added image search space schema changes from other repo (Azure#26364)

* added image search space schema changes from other repo

* fixed the az ml job show validation error when automl image job is executed inside pipeline

Author: nvijayrania

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/automl/image_vertical/image_model_distribution_settings.py

@@ -6,6 +6,12 @@
 
 from marshmallow import fields, post_dump, post_load, pre_load
 
+from azure.ai.ml._restclient.v2022_06_01_preview.models import (
+    LearningRateScheduler,
+    ModelSize,
+    StochasticOptimizer,
+    ValidationMetricType,
+)
 from azure.ai.ml._schema._sweep.search_space import (
     ChoiceSchema,
     NormalSchema,
@@ -14,12 +20,39 @@
     RandintSchema,
     UniformSchema,
 )
-from azure.ai.ml._schema.core.fields import NestedField, UnionField
+from azure.ai.ml._schema.core.fields import (
+    DumpableIntegerField,
+    DumpableStringField,
+    NestedField,
+    StringTransformedEnum,
+    UnionField,
+)
 from azure.ai.ml._schema.core.schema import PatchedSchemaMeta
+from azure.ai.ml._utils.utils import camel_to_snake
+from azure.ai.ml.constants._job.automl import (
+    ImageClassificationModelNames,
+    ImageInstanceSegmentationModelNames,
+    ImageObjectDetectionModelNames,
+)
+
+
+def get_choice_schema_of_type(cls, **kwargs):
+    class CustomChoiceSchema(ChoiceSchema):
+        values = fields.List(cls(**kwargs))
+
+    return CustomChoiceSchema()
+
 
-SEARCH_SPACE_UNION_FIELD = UnionField(
+def get_choice_and_single_value_schema_of_type(cls, **kwargs):
+    return UnionField([cls(**kwargs), NestedField(get_choice_schema_of_type(cls, **kwargs))])
+
+
+FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD = UnionField(
     [
-        NestedField(ChoiceSchema()),
+        fields.Float(),
+        DumpableIntegerField(strict=True),
+        NestedField(get_choice_schema_of_type(DumpableIntegerField, strict=True)),
+        NestedField(get_choice_schema_of_type(fields.Float)),
         NestedField(UniformSchema()),
         NestedField(QUniformSchema()),
         NestedField(NormalSchema()),
@@ -28,54 +61,100 @@
     ]
 )
 
+INT_SEARCH_SPACE_DISTRIBUTION_FIELD = UnionField(
+    [
+        DumpableIntegerField(strict=True),
+        NestedField(get_choice_schema_of_type(DumpableIntegerField, strict=True)),
+        NestedField(RandintSchema()),
+    ]
+)
+
+STRING_SEARCH_SPACE_DISTRIBUTION_FIELD = get_choice_and_single_value_schema_of_type(DumpableStringField)
+BOOL_SEARCH_SPACE_DISTRIBUTION_FIELD = get_choice_and_single_value_schema_of_type(fields.Bool)
+
+classification_model_name_args = {"allowed_values": [o.value for o in ImageClassificationModelNames]}
+detection_model_name_args = {"allowed_values": [o.value for o in ImageObjectDetectionModelNames]}
+segmentation_model_name_args = {"allowed_values": [o.value for o in ImageInstanceSegmentationModelNames]}
+model_size_enum_args = {"allowed_values": [o.value for o in ModelSize], "casing_transform": camel_to_snake}
+learning_rate_scheduler_enum_args = {
+    "allowed_values": [o.value for o in LearningRateScheduler],
+    "casing_transform": camel_to_snake,
+}
+optimizer_enum_args = {"allowed_values": [o.value for o in StochasticOptimizer], "casing_transform": camel_to_snake}
+validation_metric_enum_args = {
+    "allowed_values": [o.value for o in ValidationMetricType],
+    "casing_transform": camel_to_snake,
+}
+
+CLASSIFICATION_MODEL_NAME_DISTRIBUTION_FIELD = get_choice_and_single_value_schema_of_type(
+    StringTransformedEnum, **classification_model_name_args
+)
+DETECTION_MODEL_NAME_DISTRIBUTION_FIELD = get_choice_and_single_value_schema_of_type(
+    StringTransformedEnum, **detection_model_name_args
+)
+SEGMENTATION_MODEL_NAME_DISTRIBUTION_FIELD = get_choice_and_single_value_schema_of_type(
+    StringTransformedEnum, **segmentation_model_name_args
+)
+MODEL_SIZE_DISTRIBUTION_FIELD = get_choice_and_single_value_schema_of_type(
+    StringTransformedEnum, **model_size_enum_args
+)
+LEARNING_RATE_SCHEDULER_DISTRIBUTION_FIELD = get_choice_and_single_value_schema_of_type(
+    StringTransformedEnum, **learning_rate_scheduler_enum_args
+)
+OPTIMIZER_DISTRIBUTION_FIELD = get_choice_and_single_value_schema_of_type(StringTransformedEnum, **optimizer_enum_args)
+VALIDATION_METRIC_DISTRIBUTION_FIELD = get_choice_and_single_value_schema_of_type(
+    StringTransformedEnum, **validation_metric_enum_args
+)
+
 
 class ImageModelDistributionSettingsSchema(metaclass=PatchedSchemaMeta):
-    ams_gradient = SEARCH_SPACE_UNION_FIELD
-    beta1 = SEARCH_SPACE_UNION_FIELD
-    beta2 = SEARCH_SPACE_UNION_FIELD
-    distributed = SEARCH_SPACE_UNION_FIELD
-    early_stopping = SEARCH_SPACE_UNION_FIELD
-    early_stopping_delay = SEARCH_SPACE_UNION_FIELD
-    early_stopping_patience = SEARCH_SPACE_UNION_FIELD
-    evaluation_frequency = SEARCH_SPACE_UNION_FIELD
-    enable_onnx_normalization = SEARCH_SPACE_UNION_FIELD
-    gradient_accumulation_step = SEARCH_SPACE_UNION_FIELD
-    layers_to_freeze = SEARCH_SPACE_UNION_FIELD
-    learning_rate = SEARCH_SPACE_UNION_FIELD
-    learning_rate_scheduler = SEARCH_SPACE_UNION_FIELD
-    model_name = SEARCH_SPACE_UNION_FIELD
-    momentum = SEARCH_SPACE_UNION_FIELD
-    nesterov = SEARCH_SPACE_UNION_FIELD
-    number_of_epochs = SEARCH_SPACE_UNION_FIELD
-    number_of_workers = SEARCH_SPACE_UNION_FIELD
-    optimizer = SEARCH_SPACE_UNION_FIELD
-    random_seed = SEARCH_SPACE_UNION_FIELD
-    step_lr_gamma = SEARCH_SPACE_UNION_FIELD
-    step_lr_step_size = SEARCH_SPACE_UNION_FIELD
-    training_batch_size = SEARCH_SPACE_UNION_FIELD
-    validation_batch_size = SEARCH_SPACE_UNION_FIELD
-    warmup_cosine_lr_cycles = SEARCH_SPACE_UNION_FIELD
-    warmup_cosine_lr_warmup_epochs = SEARCH_SPACE_UNION_FIELD
-    weight_decay = SEARCH_SPACE_UNION_FIELD
+    ams_gradient = BOOL_SEARCH_SPACE_DISTRIBUTION_FIELD
+    augmentations = STRING_SEARCH_SPACE_DISTRIBUTION_FIELD
+    beta1 = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    beta2 = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    distributed = BOOL_SEARCH_SPACE_DISTRIBUTION_FIELD
+    early_stopping = BOOL_SEARCH_SPACE_DISTRIBUTION_FIELD
+    early_stopping_delay = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    early_stopping_patience = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    evaluation_frequency = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    enable_onnx_normalization = BOOL_SEARCH_SPACE_DISTRIBUTION_FIELD
+    gradient_accumulation_step = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    layers_to_freeze = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    learning_rate = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    learning_rate_scheduler = LEARNING_RATE_SCHEDULER_DISTRIBUTION_FIELD
+    momentum = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    nesterov = BOOL_SEARCH_SPACE_DISTRIBUTION_FIELD
+    number_of_epochs = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    number_of_workers = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    optimizer = OPTIMIZER_DISTRIBUTION_FIELD
+    random_seed = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    step_lr_gamma = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    step_lr_step_size = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    training_batch_size = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    validation_batch_size = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    warmup_cosine_lr_cycles = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    warmup_cosine_lr_warmup_epochs = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    weight_decay = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
 
 
 class ImageModelDistributionSettingsClassificationSchema(ImageModelDistributionSettingsSchema):
-    training_crop_size = SEARCH_SPACE_UNION_FIELD
-    validation_crop_size = SEARCH_SPACE_UNION_FIELD
-    validation_resize_size = SEARCH_SPACE_UNION_FIELD
-    weighted_loss = SEARCH_SPACE_UNION_FIELD
+    model_name = CLASSIFICATION_MODEL_NAME_DISTRIBUTION_FIELD
+    training_crop_size = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    validation_crop_size = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    validation_resize_size = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    weighted_loss = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
 
     @post_dump
     def conversion(self, data, **kwargs):
         if self.context.get("inside_pipeline", False):
             # AutoML job inside pipeline does load(dump) instead of calling to_rest_object explicitly for creating the autoRest Object from sdk job.
             # Hence for pipeline job, we explicitly convert Sweep Distribution dict to str after dump in this method.
             # For standalone automl job, same conversion happens in image_classification_job._to_rest_object()
-            from azure.ai.ml.entities._job.automl.image.image_search_space_utils import _convert_sweep_dist_dict_to_str
+            from azure.ai.ml.entities._job.automl.image.image_search_space_utils import (
+                _convert_sweep_dist_dict_to_str_dict,
+            )
 
-            for k, sweep_dist_dict in data.items():
-                if sweep_dist_dict is not None:
-                    data[k] = _convert_sweep_dist_dict_to_str(sweep_dist_dict)
+            data = _convert_sweep_dist_dict_to_str_dict(data)
         return data
 
     @pre_load
@@ -84,9 +163,7 @@ def before_make(self, data, **kwargs):
             from azure.ai.ml.entities._job.automl.image.image_search_space_utils import _convert_sweep_dist_str_to_dict
 
             # Converting Sweep Distribution str to Sweep Distribution dict for complying with search_space schema.
-            for k, val in data.items():
-                if isinstance(val, str):
-                    data[k] = _convert_sweep_dist_str_to_dict(val)
+            data = _convert_sweep_dist_str_to_dict(data)
         return data
 
     @post_load
@@ -96,20 +173,20 @@ def make(self, data, **kwargs):
         return ImageClassificationSearchSpace(**data)
 
 
-class ImageModelDistributionSettingsObjectDetectionSchema(ImageModelDistributionSettingsSchema):
-    box_detections_per_image = SEARCH_SPACE_UNION_FIELD
-    box_score_threshold = SEARCH_SPACE_UNION_FIELD
-    image_size = SEARCH_SPACE_UNION_FIELD
-    max_size = SEARCH_SPACE_UNION_FIELD
-    min_size = SEARCH_SPACE_UNION_FIELD
-    model_size = SEARCH_SPACE_UNION_FIELD
-    multi_scale = SEARCH_SPACE_UNION_FIELD
-    nms_iou_threshold = SEARCH_SPACE_UNION_FIELD
-    tile_grid_size = SEARCH_SPACE_UNION_FIELD
-    tile_overlap_ratio = SEARCH_SPACE_UNION_FIELD
-    tile_predictions_nms_threshold = SEARCH_SPACE_UNION_FIELD
-    validation_iou_threshold = SEARCH_SPACE_UNION_FIELD
-    validation_metric_type = SEARCH_SPACE_UNION_FIELD
+class ImageModelDistributionSettingsDetectionCommonSchema(ImageModelDistributionSettingsSchema):
+    box_detections_per_image = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    box_score_threshold = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    image_size = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    max_size = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    min_size = INT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    model_size = MODEL_SIZE_DISTRIBUTION_FIELD
+    multi_scale = BOOL_SEARCH_SPACE_DISTRIBUTION_FIELD
+    nms_iou_threshold = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    tile_grid_size = STRING_SEARCH_SPACE_DISTRIBUTION_FIELD
+    tile_overlap_ratio = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    tile_predictions_nms_threshold = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    validation_iou_threshold = FLOAT_SEARCH_SPACE_DISTRIBUTION_FIELD
+    validation_metric_type = VALIDATION_METRIC_DISTRIBUTION_FIELD
 
     @post_dump
     def conversion(self, data, **kwargs):
@@ -118,11 +195,11 @@ def conversion(self, data, **kwargs):
             # from sdk job object.
             # Hence for pipeline job, we explicitly convert Sweep Distribution dict to str after dump in this method.
             # For standalone automl job, same conversion happens in image_object_detection_job._to_rest_object()
-            from azure.ai.ml.entities._job.automl.image.image_search_space_utils import _convert_sweep_dist_dict_to_str
+            from azure.ai.ml.entities._job.automl.image.image_search_space_utils import (
+                _convert_sweep_dist_dict_to_str_dict,
+            )
 
-            for k, sweep_dist_dict in data.items():
-                if sweep_dist_dict is not None:
-                    data[k] = _convert_sweep_dist_dict_to_str(sweep_dist_dict)
+            data = _convert_sweep_dist_dict_to_str_dict(data)
         return data
 
     @pre_load
@@ -131,13 +208,19 @@ def before_make(self, data, **kwargs):
             from azure.ai.ml.entities._job.automl.image.image_search_space_utils import _convert_sweep_dist_str_to_dict
 
             # Converting Sweep Distribution str to Sweep Distribution dict for complying with search_space schema.
-            for k, val in data.items():
-                if isinstance(val, str):
-                    data[k] = _convert_sweep_dist_str_to_dict(val)
+            data = _convert_sweep_dist_str_to_dict(data)
         return data
 
     @post_load
     def make(self, data, **kwargs):
         from azure.ai.ml.automl import ImageObjectDetectionSearchSpace
 
         return ImageObjectDetectionSearchSpace(**data)
+
+
+class ImageModelDistributionSettingsObjectDetectionSchema(ImageModelDistributionSettingsDetectionCommonSchema):
+    model_name = DETECTION_MODEL_NAME_DISTRIBUTION_FIELD
+
+
+class ImageModelDistributionSettingsInstanceSegmentationSchema(ImageModelDistributionSettingsObjectDetectionSchema):
+    model_name = SEGMENTATION_MODEL_NAME_DISTRIBUTION_FIELD

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/automl/image_vertical/image_object_detection.py

@@ -14,6 +14,7 @@
     TaskType,
 )
 from azure.ai.ml._schema.automl.image_vertical.image_model_distribution_settings import (
+    ImageModelDistributionSettingsInstanceSegmentationSchema,
     ImageModelDistributionSettingsObjectDetectionSchema,
 )
 from azure.ai.ml._schema.automl.image_vertical.image_model_settings import (
@@ -26,11 +27,7 @@
 from azure.ai.ml.constants import AutoMLConstants
 
 
-class ImageObjectDetectionBaseSchema(ImageVerticalSchema):
-    search_space = fields.List(NestedField(ImageModelDistributionSettingsObjectDetectionSchema()))
-
-
-class ImageObjectDetectionSchema(ImageObjectDetectionBaseSchema):
+class ImageObjectDetectionSchema(ImageVerticalSchema):
     task_type = StringTransformedEnum(
         allowed_values=TaskType.IMAGE_OBJECT_DETECTION,
         casing_transform=camel_to_snake,
@@ -43,14 +40,15 @@ class ImageObjectDetectionSchema(ImageObjectDetectionBaseSchema):
         load_default=camel_to_snake(ObjectDetectionPrimaryMetrics.MEAN_AVERAGE_PRECISION),
     )
     training_parameters = NestedField(ImageModelSettingsObjectDetectionSchema())
+    search_space = fields.List(NestedField(ImageModelDistributionSettingsObjectDetectionSchema()))
 
     @post_load
     def make(self, data, **kwargs) -> Dict[str, Any]:
         data.pop("task_type")
         return data
 
 
-class ImageInstanceSegmentationSchema(ImageObjectDetectionBaseSchema):
+class ImageInstanceSegmentationSchema(ImageVerticalSchema):
     task_type = StringTransformedEnum(
         allowed_values=TaskType.IMAGE_INSTANCE_SEGMENTATION,
         casing_transform=camel_to_snake,
@@ -63,6 +61,7 @@ class ImageInstanceSegmentationSchema(ImageObjectDetectionBaseSchema):
         load_default=camel_to_snake(InstanceSegmentationPrimaryMetrics.MEAN_AVERAGE_PRECISION),
     )
     training_parameters = NestedField(ImageModelSettingsInstanceSegmentationSchema())
+    search_space = fields.List(NestedField(ImageModelDistributionSettingsInstanceSegmentationSchema()))
 
     @post_load
     def make(self, data, **kwargs) -> Dict[str, Any]:

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/image_classification_job.py

@@ -172,6 +172,8 @@ def _load_from_dict(
         from azure.ai.ml._schema.pipeline.automl_node import ImageClassificationMulticlassNodeSchema
 
         if inside_pipeline:
+            if context.get("inside_pipeline", None) is None:
+                context["inside_pipeline"] = True
             loaded_data = load_from_dict(
                 ImageClassificationMulticlassNodeSchema,
                 data,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/image_classification_multilabel_job.py

@@ -174,6 +174,8 @@ def _load_from_dict(
         from azure.ai.ml._schema.pipeline.automl_node import ImageClassificationMultilabelNodeSchema
 
         if inside_pipeline:
+            if context.get("inside_pipeline", None) is None:
+                context["inside_pipeline"] = True
             loaded_data = load_from_dict(
                 ImageClassificationMultilabelNodeSchema,
                 data,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/image_instance_segmentation_job.py

@@ -172,6 +172,8 @@ def _load_from_dict(
         from azure.ai.ml._schema.pipeline.automl_node import ImageInstanceSegmentationNodeSchema
 
         if inside_pipeline:
+            if context.get("inside_pipeline", None) is None:
+                context["inside_pipeline"] = True
             loaded_data = load_from_dict(
                 ImageInstanceSegmentationNodeSchema,
                 data,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/image_object_detection_job.py

@@ -171,6 +171,8 @@ def _load_from_dict(
         from azure.ai.ml._schema.pipeline.automl_node import ImageObjectDetectionNodeSchema
 
         if inside_pipeline:
+            if context.get("inside_pipeline", None) is None:
+                context["inside_pipeline"] = True
             loaded_data = load_from_dict(
                 ImageObjectDetectionNodeSchema,
                 data,