[PT FE] [23325] Add aten::masked_select support for pytorch models.

[YutingGao7]

Details:

• Add aten::masked_select support for pytorch models.

Tickets:

• ticket-id

[YutingGao7]

Hi openvino team,
I need some help for two issues that I encountered when I ran python -m pytest pytorch_tests/test_masked_select.py.

1. tests with GPU failed with:

E       RuntimeError: Exception from src/inference/src/cpp/core.cpp:109:
E       Exception from src/inference/src/dev/core_impl.cpp:566:
E       Device with "GPU" name is not registered in the OpenVINO Runtime

2. tests with CPU failed with:

______________ TestMaskedSelect.test_masked_select[ ie_device:CPU - precision:FP16 - mask_dtype:<class 'bool'> - input_dtype:<class 'numpy.int32'> - mask_select:ones ] ______________

self = <test_masked_select.TestMaskedSelect object at 0x154a3a590>, mask_select = 'ones', mask_dtype = <class 'bool'>, input_dtype = <class 'numpy.int32'>, ie_device = 'CPU'
precision = 'FP16', ir_version = 11

    @pytest.mark.parametrize(
        "mask_select", ['zeros', 'ones', 'random'])
    @pytest.mark.parametrize("input_dtype", [np.float32, np.float64, int, np.int32])
    @pytest.mark.parametrize("mask_dtype", [bool])  # np.float32 incorrectly casted to bool
    @pytest.mark.nightly
    @pytest.mark.precommit
    def test_masked_select(self, mask_select, mask_dtype, input_dtype, ie_device, precision, ir_version):
>       self._test(*self.create_model(),
                   ie_device, precision, ir_version,
                   kwargs_to_prepare_input={'mask_select': mask_select, 'mask_dtype': mask_dtype, "input_dtype": input_dtype})

pytorch_tests/test_masked_select.py:44: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
pytorch_tests/pytorch_layer_test_class.py:141: in _test
    compiled = core.compile_model(converted_model, ie_device)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <Core: available plugins[CPU]>
model = <Model: 'Model66'
inputs[
<ConstOutput: names[x] shape[?,?] type: i32>,
<ConstOutput: names[mask] shape[?,?] type: boolean>
]
outputs[
<ConstOutput: names[] shape[...] type: i32>
]>
device_name = 'CPU', config = None

    def compile_model(
        self,
        model: Union[Model, str, Path],
        device_name: Optional[str] = None,
        config: Optional[dict] = None,
    ) -> CompiledModel:
        """Creates a compiled model.
    
        Creates a compiled model from a source Model object or
        reads model and creates a compiled model from IR / ONNX / PDPD / TF and TFLite file.
        This can be more efficient than using read_model + compile_model(model_in_memory_object) flow,
        especially for cases when caching is enabled and cached model is available.
        If device_name is not specified, the default OpenVINO device will be selected by AUTO plugin.
        Users can create as many compiled models as they need, and use them simultaneously
        (up to the limitation of the hardware resources).
    
        :param model: Model acquired from read_model function or a path to a model in IR / ONNX / PDPD /
                      TF and TFLite format.
        :type model: Union[openvino.runtime.Model, str, pathlib.Path]
        :param device_name: Optional. Name of the device to load the model to. If not specified,
                            the default OpenVINO device will be selected by AUTO plugin.
        :type device_name: str
        :param config: Optional dict of pairs:
                       (property name, property value) relevant only for this load operation.
        :type config: dict, optional
        :return: A compiled model.
        :rtype: openvino.runtime.CompiledModel
        """
        if device_name is None:
            return CompiledModel(
                super().compile_model(model, {} if config is None else config),
            )
    
        return CompiledModel(
>           super().compile_model(model, device_name, {} if config is None else config),
        )
E       RuntimeError: Exception from src/inference/src/cpp/core.cpp:109:
E       Exception from src/inference/src/dev/plugin.cpp:54:
E       Exception from src/plugins/intel_cpu/src/node.cpp:90:
E       Unexpected: CPU plug-in doesn't support GatherND operation with dynamic rank. Operation name: aten::masked_select/GatherND

[mvafin]

tests with GPU failed with:

Please ignore tests on GPU, you can run tests with TEST_DEVICE=CPU

tests with CPU failed with:

You have dynamic rank of input, that shouldn't happen usually since we set static input rank in tests. You could try adding trace_model=True in your self._test call or if that doesn't work you could use dynamic_shapes=False, but that is less preferred.

[mvafin]

np.float32 is incorrectly casted by torch or openvino?

[YutingGao7]

I was just following test case in test_masked_fill.
https://github.com/openvinotoolkit/openvino/blob/master/tests/layer_tests/pytorch_tests/test_masked_fill.py#L67

Should I keep it or remove it?

[YutingGao7]

Hi @mvafin,
It looks like the failure is about the outputs' shapes. However, masked_select's output shape depends on the data.

https://pytorch.org/docs/stable/generated/torch.masked_select.html

What prevents cpu to support dynamic shapes? From the code, it looks like only the Result can be dynamic.

E       RuntimeError: Exception from src/inference/src/cpp/core.cpp:109:
E       Exception from src/inference/src/dev/plugin.cpp:54:
E       Exception from src/plugins/intel_cpu/src/node.cpp:90:
E       Unexpected: CPU plug-in doesn't support GatherND operation with dynamic rank. Operation name: aten::masked_select/GatherND

In plugins/intel_cpu/src/node.cpp:90, it requires shapes on outputs.

    if (typeStr != "Result" && typeStr != "Assign") {
        if (op->get_output_size() == 0) {
            OPENVINO_THROW("Node with type '", typeStr, "' and name '", name, "' does not have any outputs.");
        }
        for (size_t i = 0; i < op->get_output_size(); i++) {
            const auto &shape = op->get_output_partial_shape(i);
            if (shape.rank().is_dynamic()) {
                OPENVINO_THROW("Unexpected: CPU plug-in doesn't support ",
                               getTypeStr(),
                               " operation with dynamic rank. Operation name: ",
                               getName());
            }

            bool isScalar = shape.rank().get_length() == 0;
            outputShapes.emplace_back(isScalar ? ov::PartialShape{1} : shape);
            originalOutputPrecisions.emplace_back(op->get_output_element_type(i));
        }

        childEdges.reserve(outputShapes.size());
    }

Here is the inputs/outputs shapes printed by test_masked_select:

self = <Core: available plugins[CPU]>
model = <Model: 'Model69'
inputs[
<ConstOutput: names[x] shape[1,10] type: i32>,
<ConstOutput: names[mask] shape[1,10] type: boolean>
]
outputs[
<ConstOutput: names[] shape[...] type: i32>
]>
device_name = 'CPU', config = None

[mlukasze]

Hi @mvafin, It looks like the failure is about the outputs' shapes. However, masked_select's output shape depends on the data.

https://pytorch.org/docs/stable/generated/torch.masked_select.html

What prevents cpu to support dynamic shapes? From the code, it looks like only the Result can be dynamic.

E       RuntimeError: Exception from src/inference/src/cpp/core.cpp:109:
E       Exception from src/inference/src/dev/plugin.cpp:54:
E       Exception from src/plugins/intel_cpu/src/node.cpp:90:
E       Unexpected: CPU plug-in doesn't support GatherND operation with dynamic rank. Operation name: aten::masked_select/GatherND

In plugins/intel_cpu/src/node.cpp:90, it requires shapes on outputs.

    if (typeStr != "Result" && typeStr != "Assign") {
        if (op->get_output_size() == 0) {
            OPENVINO_THROW("Node with type '", typeStr, "' and name '", name, "' does not have any outputs.");
        }
        for (size_t i = 0; i < op->get_output_size(); i++) {
            const auto &shape = op->get_output_partial_shape(i);
            if (shape.rank().is_dynamic()) {
                OPENVINO_THROW("Unexpected: CPU plug-in doesn't support ",
                               getTypeStr(),
                               " operation with dynamic rank. Operation name: ",
                               getName());
            }

            bool isScalar = shape.rank().get_length() == 0;
            outputShapes.emplace_back(isScalar ? ov::PartialShape{1} : shape);
            originalOutputPrecisions.emplace_back(op->get_output_element_type(i));
        }

        childEdges.reserve(outputShapes.size());
    }

Here is the inputs/outputs shapes printed by test_masked_select:

self = <Core: available plugins[CPU]>
model = <Model: 'Model69'
inputs[
<ConstOutput: names[x] shape[1,10] type: i32>,
<ConstOutput: names[mask] shape[1,10] type: boolean>
]
outputs[
<ConstOutput: names[] shape[...] type: i32>
]>
device_name = 'CPU', config = None

@mvafin what do you think about it?

[mvafin]

You need to do Transpose after NonZero, since the format of indices in GatherND is different. Like here:

openvino/src/frontends/pytorch/src/op/index.cpp

Line 241 in c20d232

auto masked_id = context.mark_node(std::make_shared<v1::Transpose>(nonzero, input_order));

[github-actions]

This PR will be closed in a week because of 2 weeks of no activity.

[github-actions]

This PR was closed because it has been stalled for 2 week with no activity.