From 32e3fb8d4910c41a5dea983bd45e9387fdd61b9e Mon Sep 17 00:00:00 2001 From: coco <1228759711@qq.com> Date: Thu, 31 Oct 2024 20:53:01 +0000 Subject: [PATCH 1/5] [feat]: index_fill_ frontend pytorch op --- src/frontends/pytorch/src/op/index_fill_.cpp | 69 ++++++++++++++ src/frontends/pytorch/src/op_table.cpp | 2 + .../pytorch_tests/test_index_fill_.py | 89 +++++++++++++++++++ 3 files changed, 160 insertions(+) create mode 100644 src/frontends/pytorch/src/op/index_fill_.cpp create mode 100644 tests/layer_tests/pytorch_tests/test_index_fill_.py diff --git a/src/frontends/pytorch/src/op/index_fill_.cpp b/src/frontends/pytorch/src/op/index_fill_.cpp new file mode 100644 index 00000000000000..4d24416722d279 --- /dev/null +++ b/src/frontends/pytorch/src/op/index_fill_.cpp @@ -0,0 +1,69 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include + +#include "openvino/core/node.hpp" +#include "openvino/core/node_output.hpp" +#include "openvino/core/type/element_type.hpp" +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/broadcast.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/convert_like.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/scatter_elements_update.hpp" +#include "openvino/op/shape_of.hpp" +#include "openvino/op/slice.hpp" +#include "utils.hpp" + +using namespace ov::op; + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +OutputVector translate_index_fill_(const NodeContext& context) { + // aten::index_fill_(self, dim, index, value) --> Tensor + num_inputs_check(context, 4, 4); + auto input = context.get_input(0); + auto dim = context.get_input(1); + auto index = context.get_input(2); + auto value = context.get_input(3); + + auto const_1_vec = v0::Constant::create(element::i32, Shape{1}, {1}); + + auto tensor_rank = std::get<1>(get_shape_rank(context, input, true)); + auto tensor_rank_correct_type = context.mark_node(std::make_shared(tensor_rank, dim)); + auto positive_dim = normalize_axis(context, dim, tensor_rank_correct_type); + + // scalar to vec + auto dim_vec = context.mark_node(std::make_shared(positive_dim, const_1_vec, false)); + auto value_vec = context.mark_node(std::make_shared(value, const_1_vec, false)); + + auto input_shape = std::get<0>(get_shape_rank(context, input, false)); + + auto index_shape = std::get<0>(get_shape_rank(context, index, false)); + auto const_0 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0})); + auto const_1 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1})); + auto index_len = context.mark_node(std::make_shared(index_shape, const_0, const_1, const_1)); + + // [A, B, ..., T, ..., K] --> [A, B, ..., len(index), ..., K] + auto target_shape = std::make_shared(input_shape, + dim_vec, + index_len, + v0::Constant::create(element::i32, Shape{}, {0})); + + // broadcast && index fill + auto broadcasted_value = context.mark_node(std::make_shared(value_vec, target_shape, dim_vec)); + auto broadcasted_index = context.mark_node(std::make_shared(index, target_shape, dim_vec)); + auto result = context.mark_node(std::make_shared(input, broadcasted_index, broadcasted_value, dim)); + + return {result}; +}; + +} // 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 607f0bd32db80d..7061b3e6ea241c 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -113,6 +113,7 @@ OP_CONVERTER(translate_im2col); OP_CONVERTER(translate_index); OP_CONVERTER(translate_index_add); OP_CONVERTER(translate_index_copy_); +OP_CONVERTER(translate_index_fill_); OP_CONVERTER(translate_index_put_); OP_CONVERTER(translate_index_select); OP_CONVERTER(translate_instance_norm); @@ -496,6 +497,7 @@ const std::unordered_map get_supported_ops_ts() { // aten::imag - Supported in limited set of patterns // aten::index - Supported in limited set of patterns {"aten::index_copy_", op::inplace_op}, + {"aten::index_fill_", op::inplace_op}, {"aten::index_put_", op::inplace_op}, {"aten::index_add", op::translate_index_add}, {"aten::index_select", op::translate_index_select}, diff --git a/tests/layer_tests/pytorch_tests/test_index_fill_.py b/tests/layer_tests/pytorch_tests/test_index_fill_.py new file mode 100644 index 00000000000000..a2994036dcecab --- /dev/null +++ b/tests/layer_tests/pytorch_tests/test_index_fill_.py @@ -0,0 +1,89 @@ +# Copyright (C) 2018-2023 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + + +import numpy as np +import pytest +import torch + +from pytorch_layer_test_class import PytorchLayerTest + + +class TestIndexFill(PytorchLayerTest): + def _prepare_input(self): + return (self.input_tensor,) + + def create_model(self, dim, index, values): + class aten_index_fill_(torch.nn.Module): + def __init__(self, dim, index, values): + super().__init__() + self.dim = dim + self.index = index + self.values = values + + def forward(self, input_tensor): + input_tensor.index_fill_(self.dim, self.index, self.values) + return input_tensor + + ref_net = None + + return aten_index_fill_(dim, index, values), ref_net, "aten::index_fill_" + + @pytest.mark.parametrize( + "input_data", + ( + # { + # "input_shape": [3, 3], + # "dim": 1, + # "input_value": 2.0, + # "index": [2], + # }, + { + "input_shape": [10], + "dim": 0, + "input_value": 5.6, + "index": [5, 6, 7] + }, + { + "input_shape": [3, 3], + "dim": 0, + "input_value": 10.1, + "index": [1, 0] + }, + { + "input_shape": [4, 3, 5], + "dim": 1, + "input_value": 1234.5, + "index": [2, 0] + }, + { + "input_shape": [5, 6, 7, 8], + "dim": -2, + "input_value": 0.1234, + "index": [6, 4, 2, 0] + }, + { + "input_shape": [5, 6, 7, 8], + "dim": -3, + "input_value": -4321234.5678765, + "index": [5, 4, 3, 1] + }, + { + "input_shape": [5, 6, 7, 8], + "dim": 3, + "input_value": -1234.54321, + "index": [6, 4, 7, 2, 1] + }, + ), + ) + @pytest.mark.nightly + @pytest.mark.precommit + def test_index_fill_single_index(self, ie_device, precision, ir_version, input_data): + self.input_tensor = np.random.randn(*input_data["input_shape"]).astype(np.float32) + values = torch.tensor(np.float32(input_data["input_value"])) + dim = input_data["dim"] + shape = self.input_tensor.shape + max_idx = shape[dim] + n_select = np.random.randint(1, max_idx + 1) + index = torch.from_numpy(np.random.choice(np.arange(0, max_idx), n_select, replace=False)) + self._test(*self.create_model(dim, index, values), ie_device, precision, ir_version) \ No newline at end of file From 3cf588c591eae50a7b360fc4a40967490216391c Mon Sep 17 00:00:00 2001 From: coco <1228759711@qq.com> Date: Thu, 31 Oct 2024 21:14:17 +0000 Subject: [PATCH 2/5] [fix]: test case --- tests/layer_tests/pytorch_tests/test_index_fill_.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/layer_tests/pytorch_tests/test_index_fill_.py b/tests/layer_tests/pytorch_tests/test_index_fill_.py index a2994036dcecab..2ab5e5e2df9605 100644 --- a/tests/layer_tests/pytorch_tests/test_index_fill_.py +++ b/tests/layer_tests/pytorch_tests/test_index_fill_.py @@ -32,12 +32,6 @@ def forward(self, input_tensor): @pytest.mark.parametrize( "input_data", ( - # { - # "input_shape": [3, 3], - # "dim": 1, - # "input_value": 2.0, - # "index": [2], - # }, { "input_shape": [10], "dim": 0, From 9e7eb9106dc666f8609d43bc2582aaf73567943a Mon Sep 17 00:00:00 2001 From: coco <69197635+cocoshe@users.noreply.github.com> Date: Wed, 20 Nov 2024 17:47:43 +0800 Subject: [PATCH 3/5] skip Reshape Co-authored-by: Maxim Vafin --- src/frontends/pytorch/src/op/index_fill_.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/frontends/pytorch/src/op/index_fill_.cpp b/src/frontends/pytorch/src/op/index_fill_.cpp index 4d24416722d279..40760d2f358100 100644 --- a/src/frontends/pytorch/src/op/index_fill_.cpp +++ b/src/frontends/pytorch/src/op/index_fill_.cpp @@ -34,12 +34,11 @@ OutputVector translate_index_fill_(const NodeContext& context) { auto const_1_vec = v0::Constant::create(element::i32, Shape{1}, {1}); - auto tensor_rank = std::get<1>(get_shape_rank(context, input, true)); + auto tensor_rank = std::get<1>(get_shape_rank(context, input, false)); auto tensor_rank_correct_type = context.mark_node(std::make_shared(tensor_rank, dim)); - auto positive_dim = normalize_axis(context, dim, tensor_rank_correct_type); + auto dim_vec = normalize_axis(context, dim, tensor_rank_correct_type); // scalar to vec - auto dim_vec = context.mark_node(std::make_shared(positive_dim, const_1_vec, false)); auto value_vec = context.mark_node(std::make_shared(value, const_1_vec, false)); auto input_shape = std::get<0>(get_shape_rank(context, input, false)); From 04ddd033241a91c18af714c744afae2fbb417942 Mon Sep 17 00:00:00 2001 From: coco <1228759711@qq.com> Date: Wed, 20 Nov 2024 16:25:23 +0000 Subject: [PATCH 4/5] [fix] code style --- src/frontends/pytorch/src/op/index_fill_.cpp | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/frontends/pytorch/src/op/index_fill_.cpp b/src/frontends/pytorch/src/op/index_fill_.cpp index 40760d2f358100..a24f3fa2f5b1c7 100644 --- a/src/frontends/pytorch/src/op/index_fill_.cpp +++ b/src/frontends/pytorch/src/op/index_fill_.cpp @@ -49,15 +49,16 @@ OutputVector translate_index_fill_(const NodeContext& context) { auto index_len = context.mark_node(std::make_shared(index_shape, const_0, const_1, const_1)); // [A, B, ..., T, ..., K] --> [A, B, ..., len(index), ..., K] - auto target_shape = std::make_shared(input_shape, - dim_vec, - index_len, - v0::Constant::create(element::i32, Shape{}, {0})); + auto target_shape = std::make_shared(input_shape, + dim_vec, + index_len, + v0::Constant::create(element::i32, Shape{}, {0})); // broadcast && index fill auto broadcasted_value = context.mark_node(std::make_shared(value_vec, target_shape, dim_vec)); auto broadcasted_index = context.mark_node(std::make_shared(index, target_shape, dim_vec)); - auto result = context.mark_node(std::make_shared(input, broadcasted_index, broadcasted_value, dim)); + auto result = context.mark_node( + std::make_shared(input, broadcasted_index, broadcasted_value, dim)); return {result}; }; From 9637797cdde749dc0911022398850cede275a06c Mon Sep 17 00:00:00 2001 From: coco <1228759711@qq.com> Date: Fri, 22 Nov 2024 06:15:38 +0000 Subject: [PATCH 5/5] [fix] to 'long' for win --- tests/layer_tests/pytorch_tests/test_index_fill_.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/layer_tests/pytorch_tests/test_index_fill_.py b/tests/layer_tests/pytorch_tests/test_index_fill_.py index 2ab5e5e2df9605..878dda7ab3bd7e 100644 --- a/tests/layer_tests/pytorch_tests/test_index_fill_.py +++ b/tests/layer_tests/pytorch_tests/test_index_fill_.py @@ -79,5 +79,5 @@ def test_index_fill_single_index(self, ie_device, precision, ir_version, input_d shape = self.input_tensor.shape max_idx = shape[dim] n_select = np.random.randint(1, max_idx + 1) - index = torch.from_numpy(np.random.choice(np.arange(0, max_idx), n_select, replace=False)) + index = torch.from_numpy(np.random.choice(np.arange(0, max_idx), n_select, replace=False)).to(torch.long) self._test(*self.create_model(dim, index, values), ie_device, precision, ir_version) \ No newline at end of file