From 1783aad48d2ea61606a973fa9af54049d6885caf Mon Sep 17 00:00:00 2001 From: Mateusz Mikolajczyk Date: Mon, 18 Dec 2023 09:54:06 +0100 Subject: [PATCH] [PT FE][Opset13] Add translation for torchscript aten::multinomial operator (#20397) * Enable aten::multinomial in PT FE * Add aten::multinomial test * Fix typo * Update tests/layer_tests/pytorch_tests/test_multinomial.py Co-authored-by: Roman Kazantsev * Update multinomial to changes in specification * rm ; * Add temporary convert for multinomial i64 support * Revert convert * Fix torch test --------- Co-authored-by: Roman Kazantsev --- src/frontends/pytorch/src/op/multinomial.cpp | 55 +++++++++++ src/frontends/pytorch/src/op_table.cpp | 2 + .../pytorch_tests/test_multinomial.py | 99 +++++++++++++++++++ 3 files changed, 156 insertions(+) create mode 100644 src/frontends/pytorch/src/op/multinomial.cpp create mode 100644 tests/layer_tests/pytorch_tests/test_multinomial.py diff --git a/src/frontends/pytorch/src/op/multinomial.cpp b/src/frontends/pytorch/src/op/multinomial.cpp new file mode 100644 index 00000000000..c359c34bbcb --- /dev/null +++ b/src/frontends/pytorch/src/op/multinomial.cpp @@ -0,0 +1,55 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/op/multinomial.hpp" + +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/concat.hpp" +#include "openvino/op/gather.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/scatter_elements_update.hpp" +#include "openvino/op/shape_of.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +using namespace ov::op; + +OutputVector translate_multinomial(const NodeContext& context) { + num_inputs_check(context, 3, 5); + auto const_neg_1 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1})); + 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 input = context.get_input(0); + auto num_samples = context.mark_node(std::make_shared(context.get_input(1), const_1, false)); + auto replacement = context.const_input(2); + FRONT_END_OP_CONVERSION_CHECK(context.input_is_none(3), + "aten::multinomial conversion with generator is not supported"); + + // Torch multinomial accept input of [class_probs] or [bs, class_probs], convert always to [bs, class_probs] for OV. + auto input_shape = context.mark_node(std::make_shared(input, element::i32)); + auto class_probs_shape = context.mark_node(std::make_shared(input_shape, const_neg_1, const_0)); + auto inp_shape = context.mark_node(std::make_shared(OutputVector{const_neg_1, class_probs_shape}, 0)); + input = context.mark_node(std::make_shared(input, inp_shape, false)); + + auto multinomial = + context.mark_node(std::make_shared(input, num_samples, element::i64, replacement, false)); + + // Torch multinomial can return [num_samples] or [bs, num_samples] based on input dim, reshape to correct shape. + auto out_shape = context.mark_node( + std::make_shared(input_shape, const_neg_1, num_samples, const_neg_1)); + multinomial = context.mark_node(std::make_shared(multinomial, out_shape, false)); + + if (!context.input_is_none(5)) { + context.mutate_input(5, multinomial); + } + return {multinomial}; +}; +} // 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 4ddf3a5b68a..9a4c1e3f69a 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -126,6 +126,7 @@ OP_CONVERTER(translate_mean); OP_CONVERTER(translate_meshgrid); OP_CONVERTER(translate_min); OP_CONVERTER(translate_minimum); +OP_CONVERTER(translate_multinomial); OP_CONVERTER(translate_narrow); OP_CONVERTER(translate_native_multi_head_attention); OP_CONVERTER(translate_neg); @@ -448,6 +449,7 @@ const std::map get_supported_ops_ts() { {"aten::mul_", op::inplace_translate_1to1_match_2_inputs_align_types}, {"aten::multiply", op::translate_1to1_match_2_inputs_align_types}, {"aten::multiply_", op::inplace_translate_1to1_match_2_inputs_align_types}, + {"aten::multinomial", op::translate_multinomial}, {"aten::narrow", op::translate_narrow}, {"aten::ne", op::translate_1to1_match_2_inputs_align_types}, {"aten::neg", op::translate_neg}, diff --git a/tests/layer_tests/pytorch_tests/test_multinomial.py b/tests/layer_tests/pytorch_tests/test_multinomial.py new file mode 100644 index 00000000000..1ddca3a2258 --- /dev/null +++ b/tests/layer_tests/pytorch_tests/test_multinomial.py @@ -0,0 +1,99 @@ +# 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 TestMultinomial(PytorchLayerTest): + def _prepare_input(self, input, num_samples, out): + model_inputs = [input, np.array(num_samples)] + if out: + model_inputs.append(np.zeros_like(input).astype(np.int64)) + return model_inputs + + def create_model(self, replacement, out, test_type): + class aten_multinomial(torch.nn.Module): + def __init__(self, replacement, out, test_type) -> None: + super().__init__() + self.replacement = replacement + self.test_type = test_type + if not out: + self.forward = self.multinomial + else: + self.forward = self.multinomial_out + + def multinomial(self, input, num_samples): + multinomial = torch.multinomial(input, num_samples, self.replacement) + return self.mode(multinomial) + + def multinomial_out(self, input, num_samples, out): + multinomial = torch.multinomial(input, num_samples, self.replacement, out=out) + return self.mode(multinomial) + + def mode(self, op): + if self.test_type == "shape": + return torch.ones_like(op) + if self.test_type == "sorted": + return torch.sort(op)[0] + return op + + ref_net = None + + return aten_multinomial(replacement, out, test_type), ref_net, "aten::multinomial" + + @pytest.mark.parametrize( + ("input", "num_samples", "replacement", "test_type"), + [ + ( + np.array([[0, 1, 0, 0], [0, 0, 0, 1], [1, 0, 0, 0]], dtype=np.float32), + 1024, + True, + "exact", + ), + ( + np.array([[0.001, 0.001, 0.1, 0.9], [5, 10, 0.1, 256], [1, 1e-5, 1e-5, 1e-5]], dtype=np.float32), + 4, + False, + "sorted", + ), + ( + np.array([[0.001, 0, 0.1, 0.9], [5, 10, 0, 256], [0.9, 0.001, 1e-5, 0]], dtype=np.float64), + 3, + False, + "sorted", + ), + ( + np.array([[1, 2, 3, 4], [4, 3, 2, 1], [1, 0, 0, 0]], dtype=np.float32), + 1, + True, + "shape", + ), + ( + np.array([1, 2, 3, 4], dtype=np.float32), + 256, + True, + "shape", + ), + ( + np.array([[1, 2, 3, 4]], dtype=np.float32), + 256, + True, + "shape", + ), + ], + ) + @pytest.mark.parametrize("out", [True, False]) + @pytest.mark.nightly + @pytest.mark.precommit + def test_multinomial(self, input, num_samples, replacement, out, test_type, ie_device, precision, ir_version): + self._test( + *self.create_model(replacement, out, test_type), + ie_device, + precision, + ir_version, + kwargs_to_prepare_input={"input": input, "num_samples": num_samples, "out": out} + )