diff --git a/src/frontends/pytorch/src/op/index_select.cpp b/src/frontends/pytorch/src/op/index_select.cpp new file mode 100644 index 00000000000..bf1e13bd39b --- /dev/null +++ b/src/frontends/pytorch/src/op/index_select.cpp @@ -0,0 +1,31 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/gather.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +OutputVector translate_index_select(const NodeContext& context) { + // aten::index_select(Tensor self, int dim, Tensor index) -> Tensor + // aten::index_select.out(Tensor self, int dim, Tensor index, *, Tensor(a!) out) -> Tensor(a!) + num_inputs_check(context, 3, 4); + auto x = context.get_input(0); + auto dim = context.get_input(1); + auto indicies = context.get_input(2); + auto gather = context.mark_node(std::make_shared(x, indicies, dim)); + if (!context.input_is_none(3)) { + context.mutate_input(3, gather); + } + return {gather}; +}; + +} // namespace op +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp index 590c8e6706a..ac9455b1e9b 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -59,6 +59,7 @@ OP_CONVERTER(translate_hardtanh); OP_CONVERTER(translate_if); OP_CONVERTER(translate_im2col); OP_CONVERTER(translate_index_put_); +OP_CONVERTER(translate_index_select); OP_CONVERTER(translate_instance_norm); OP_CONVERTER(translate_int); OP_CONVERTER(translate_layer_norm); @@ -230,6 +231,7 @@ const std::map get_supported_ops() { {"aten::hardtanh_", op::inplace_op}, {"aten::im2col", op::translate_im2col}, {"aten::index_put_", op::inplace_op}, + {"aten::index_select", op::translate_index_select}, {"aten::instance_norm", op::translate_instance_norm}, {"aten::Int", op::translate_int}, {"aten::IntImplicit", op::translate_int}, diff --git a/src/frontends/pytorch/src/transforms/aten_index_replacer.cpp b/src/frontends/pytorch/src/transforms/aten_index_replacer.cpp index cf9bd5c9fa2..e423606112d 100644 --- a/src/frontends/pytorch/src/transforms/aten_index_replacer.cpp +++ b/src/frontends/pytorch/src/transforms/aten_index_replacer.cpp @@ -213,7 +213,7 @@ AtenIndexToSelect::AtenIndexToSelect() { } concat_dims.push_back(cum_adv_index_shape_tensor); for (auto i : non_used_dims) { - if (i < advanced_ids[i]) { + if (i < advanced_ids[0]) { continue; } concat_dims.push_back(input_dims->output(i)); diff --git a/tests/layer_tests/pytorch_tests/test_index.py b/tests/layer_tests/pytorch_tests/test_index.py index c4e303d244f..f3f681d1c3c 100644 --- a/tests/layer_tests/pytorch_tests/test_index.py +++ b/tests/layer_tests/pytorch_tests/test_index.py @@ -82,7 +82,7 @@ class TestIndexRange(PytorchLayerTest): def create_model(self): import torch - class aten_index_unsqueeze(torch.nn.Module): + class aten_index_arange(torch.nn.Module): def forward(self, x, y): x = x.reshape(x.shape[0], -1) @@ -90,7 +90,20 @@ class TestIndexRange(PytorchLayerTest): ref_net = None - return aten_index_unsqueeze(), ref_net, "aten::index" + return aten_index_arange(), ref_net, "aten::index" + + def create_model2(self): + import torch + + class aten_index_arange(torch.nn.Module): + + def forward(self, x, y): + x = x.reshape(x.shape[0], x.shape[1], -1, 1) + return x[torch.arange(x.shape[0]), y] + + ref_net = None + + return aten_index_arange(), ref_net, "aten::index" @pytest.mark.nightly @pytest.mark.precommit @@ -102,3 +115,13 @@ class TestIndexRange(PytorchLayerTest): def test_index_range(self, input_shape, idx, ie_device, precision, ir_version): self._test(*self.create_model(), ie_device, precision, ir_version, kwargs_to_prepare_input={ "input_shape": input_shape, "idx": idx}, trace_model=True, dynamic_shapes=False) + @pytest.mark.nightly + @pytest.mark.precommit + @pytest.mark.parametrize(("input_shape", "idx"), ( + ((1, 1), [0]), + ([2, 3], [1, 2]), + ([7, 8, 9], [1]), + ([2, 2, 3, 4], [0]))) + def test_index_range_free_dims(self, input_shape, idx, ie_device, precision, ir_version): + self._test(*self.create_model2(), ie_device, precision, ir_version, kwargs_to_prepare_input={ + "input_shape": input_shape, "idx": idx}, trace_model=True, dynamic_shapes=False) diff --git a/tests/layer_tests/pytorch_tests/test_index_select.py b/tests/layer_tests/pytorch_tests/test_index_select.py new file mode 100644 index 00000000000..1c8c2c91eb3 --- /dev/null +++ b/tests/layer_tests/pytorch_tests/test_index_select.py @@ -0,0 +1,46 @@ +# Copyright (C) 2018-2023 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from pytorch_layer_test_class import PytorchLayerTest + + +class TestIndexSelect(PytorchLayerTest): + def _prepare_input(self, index, out=False, dim=0): + import numpy as np + index = np.array(index).astype(np.int32) + input_data = np.random.randn(2, 3, 10, 10).astype(np.float32) + if not out: + return (input_data, index) + out = np.zeros_like(np.take(input_data, axis=dim, indices=index)) + return (input_data, index, out) + + def create_model(self, dim, out=False): + import torch + + class aten_index_select(torch.nn.Module): + def __init__(self, dim, out=False): + super(aten_index_select, self).__init__() + self.dim = dim + if out: + self.forward = self.forward_out + + def forward(self, x, indices): + return torch.index_select(x, self.dim, indices) + + def forward_out(self, x, indices, out): + return out, torch.index_select(x, self.dim, indices, out=out) + + ref_net = None + + return aten_index_select(dim, out), ref_net, "aten::index_select" + + @pytest.mark.parametrize("dim", [0, 1, 2, 3, -1, -2, -3]) + @pytest.mark.parametrize("indices", [[0, 1], [0], [1, 0]]) + @pytest.mark.parametrize("out", [False, True]) + @pytest.mark.nightly + @pytest.mark.precommit + def test_index_select(self, dim, out, indices, ie_device, precision, ir_version): + self._test(*self.create_model(dim, out), ie_device, precision, ir_version, + kwargs_to_prepare_input={"index": indices, "out": out, "dim": dim})