[PT FE] Support elu, celu, mish, shuffle ops (#23332)

### Details:
- *Support `aten::elu_`, `aten::celu`, `aten::celu_`, `aten::mish`,
`aten::mish_` TorchScript ops*
- *Support `aten.celu.default`, `aten.channel_shuffle.default`,
`aten.elu.default`, `aten.mish.default`, `aten.pixel_unshuffle.default`,
`aten.pixel_shuffle.default` FX ops*

### Tickets:
 - *ticket-id*

---------

Co-authored-by: Ilya Lavrenov <ilya.lavrenov@intel.com>
This commit is contained in:
Maxim Vafin 2024-03-13 22:41:08 +01:00 committed by GitHub
parent 02fe45a061
commit 0afddfe02b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 163 additions and 9 deletions

View File

@ -0,0 +1,40 @@
// Copyright (C) 2018-2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "openvino/frontend/pytorch/node_context.hpp"
#include "openvino/op/divide.hpp"
#include "openvino/op/elu.hpp"
#include "openvino/op/multiply.hpp"
#include "utils.hpp"
namespace ov {
namespace frontend {
namespace pytorch {
namespace op {
using namespace ov::op;
OutputVector translate_celu(const NodeContext& context) {
// aten::celu(%x_copy.1, %self.alpha)
num_inputs_check(context, 1, 2);
auto x = context.get_input(0);
Output<Node> alpha;
if (context.input_is_none(1)) {
alpha = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1.}));
} else {
alpha = context.get_input(1);
}
alpha = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, x));
auto divide_node = context.mark_node(std::make_shared<v1::Divide>(x, alpha));
auto elu_node = context.mark_node(std::make_shared<v0::Elu>(divide_node, 1.));
auto elu = context.mark_node(std::make_shared<v1::Multiply>(alpha, elu_node));
return {elu};
};
} // namespace op
} // namespace pytorch
} // namespace frontend
} // namespace ov

View File

@ -5,6 +5,7 @@
#include "openvino/op/elu.hpp"
#include "openvino/frontend/pytorch/node_context.hpp"
#include "openvino/op/multiply.hpp"
#include "utils.hpp"
namespace ov {
@ -12,17 +13,22 @@ namespace frontend {
namespace pytorch {
namespace op {
using namespace ov::op;
OutputVector translate_elu(const NodeContext& context) {
// aten::elu(Tensor self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> Tensor
num_inputs_check(context, 2, 4);
auto x = context.get_input(0);
auto alpha = context.const_input<float>(1);
// TODO: Figure out what scale and input_scale do
PYTORCH_OP_CONVERSION_CHECK(context.input_is_none(2) || context.const_input<int64_t>(2) == 1,
"Unexpected value of scale input for elu operation");
auto elu = context.mark_node(std::make_shared<v0::Elu>(x, alpha));
if (!context.input_is_none(2)) {
auto scale = context.get_input(2);
scale = context.mark_node(std::make_shared<v1::ConvertLike>(scale, elu));
elu = context.mark_node(std::make_shared<v1::Multiply>(elu, scale));
}
PYTORCH_OP_CONVERSION_CHECK(context.input_is_none(3) || context.const_input<int64_t>(3) == 1,
"Unexpected value of input_scale input for elu operation");
return {context.mark_node(std::make_shared<ov::op::v0::Elu>(x, alpha))};
return {elu};
};
} // namespace op

View File

@ -49,6 +49,7 @@ OP_CONVERTER(translate_bitwise_or);
OP_CONVERTER(translate_bitwise_xor);
OP_CONVERTER(translate_cat);
OP_CONVERTER(translate_cdist);
OP_CONVERTER(translate_celu);
OP_CONVERTER(translate_channel_shuffle);
OP_CONVERTER(translate_clamp);
OP_CONVERTER(translate_constant);
@ -361,6 +362,8 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
{"aten::cdist", op::translate_cdist},
{"aten::ceil", op::optional_out<op::translate_1to1_match_1_inputs<opset10::Ceiling>, 1>},
{"aten::ceil_", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Ceiling>>},
{"aten::celu", op::translate_celu},
{"aten::celu_", op::inplace_op<op::translate_celu>},
{"aten::channel_shuffle", op::translate_channel_shuffle},
// aten::chunk - Supported in limited set of patterns
{"aten::clamp", op::translate_clamp},
@ -399,6 +402,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
{"aten::dropout_", op::skip_node},
// aten::einsum - Supported in limited set of patterns
{"aten::elu", op::translate_elu},
{"aten::elu_", op::inplace_op<op::translate_elu>},
{"aten::embedding", op::translate_embedding},
{"aten::embedding_bag", op::translate_embedding_bag},
{"aten::empty", op::translate_empty},
@ -512,6 +516,8 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
{"aten::meshgrid", op::translate_meshgrid},
{"aten::min", op::translate_min},
{"aten::minimum", op::translate_minimum},
{"aten::mish", op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Mish>},
{"aten::mish_", op::inplace_op<op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Mish>>},
{"aten::mm", op::translate_1to1_match_2_inputs<opset10::MatMul>},
{"aten::mul", op::translate_mul},
{"aten::mul_", op::translate_mul_},
@ -739,6 +745,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_fx() {
{"aten.bmm.default", op::translate_1to1_match_2_inputs_align_types<opset10::MatMul>},
{"aten.cat.default", op::translate_cat_fx},
{"aten.ceil.default", op::translate_1to1_match_1_inputs<opset10::Ceiling>},
{"aten.celu.default", op::translate_celu},
{"aten.clamp.default", op::translate_clamp},
{"aten.clamp_max.default", op::translate_1to1_match_2_inputs_align_types<opset10::Minimum>},
{"aten.clamp_max.Tensor", op::translate_1to1_match_2_inputs_align_types<opset10::Minimum>},
@ -752,10 +759,12 @@ const std::map<std::string, CreatorFunction> get_supported_ops_fx() {
{"aten.cos.default", op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Cos>},
{"aten.cosh.default", op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Cosh>},
{"aten.cumsum.default", op::translate_cumsum_fx},
{"aten.channel_shuffle.default", op::translate_channel_shuffle},
{"aten.detach.default", op::skip_node},
{"aten.div.Scalar", op::translate_div_fx},
{"aten.div.Tensor", op::translate_div_fx},
{"aten.div.Tensor_mode", op::translate_div_fx},
{"aten.elu.default", op::translate_elu},
{"aten.embedding.default", op::translate_embedding},
{"aten.empty.memory_format", op::translate_empty},
{"aten.eq.Scalar", op::translate_1to1_match_2_inputs_align_types<opset10::Equal>},
@ -816,6 +825,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_fx() {
{"aten.min.default", op::translate_min},
{"aten.min.dim", op::translate_min_dim_fx},
{"aten.minimum.default", op::translate_minimum},
{"aten.mish.default", op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Mish>},
{"aten.mm.default", op::translate_1to1_match_2_inputs<opset10::MatMul>},
{"aten.mul.Scalar", op::translate_mul},
{"aten.mul.Tensor", op::translate_mul},
@ -832,6 +842,8 @@ const std::map<std::string, CreatorFunction> get_supported_ops_fx() {
{"aten.pow.Scalar", op::translate_pow},
{"aten.pow.Tensor_Scalar", op::translate_pow},
{"aten.pow.Tensor_Tensor", op::translate_pow},
{"aten.pixel_shuffle.default", op::translate_pixel_shuffle},
{"aten.pixel_unshuffle.default", op::translate_pixel_unshuffle},
{"aten.reciprocal.default", op::translate_reciprocal},
{"aten.relu.default", op::translate_1to1_match_1_inputs<opset10::Relu>},
{"aten.relu_.default", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Relu>>},

View File

@ -51,7 +51,6 @@ AtenStackListConstructReplacer::AtenStackListConstructReplacer() {
node = compression;
} else {
OutputVector node_vector;
auto zero = v0::Constant::create(element::i32, Shape{}, {0});
// Iterate over values in ListConstruct
for (const auto& list_input : list_inputs) {
auto node = concat_list_construct(list_input);

View File

@ -152,9 +152,10 @@ const std::unordered_map<int64_t, element::Type> TORCH_TO_OV_TYPE{
{6, element::f32},
{7, element::f64},
{11, element::boolean},
{12, element::i8}, // quantized i8
{13, element::u8}, // quantized u8
{14, element::i32} // quantized i32
{12, element::i8}, // quantized i8
{13, element::u8}, // quantized u8
{14, element::i32}, // quantized i32
{15, element::bf16},
};
const std::unordered_map<std::string, ov::op::PadType> TORCH_AUTO_PAD_TO_OV{{"valid", ov::op::PadType::VALID},

View File

@ -250,7 +250,7 @@ def test_so_extension():
converted_model = fe.convert(input_model)
assert converted_model
assert [n.get_type_name() for n in converted_model.get_ordered_ops()] == [
"Parameter", "Elu", "Result"]
'Parameter', 'Elu', 'Constant', 'ConvertLike', 'Multiply', 'Result']
fe.add_extension(get_builtin_extensions_path())
converted_model = fe.convert(input_model)

View File

@ -0,0 +1,45 @@
# Copyright (C) 2018-2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import pytest
import torch
import torch.nn.functional as F
from pytorch_layer_test_class import PytorchLayerTest, skip_if_export
class aten_celu(torch.nn.Module):
def __init__(self, alpha, dtype, inplace):
super(aten_celu, self).__init__()
self.alpha = alpha
self.dtype = dtype
self.inplace = inplace
if alpha is None:
self.forward = self.forward_no_alpha
def forward(self, x):
x_copy = x.to(self.dtype)
return F.celu(x_copy, alpha=self.alpha, inplace=self.inplace), x_copy
def forward_no_alpha(self, x):
x_copy = x.to(self.dtype)
return F.celu(x_copy, inplace=self.inplace), x_copy
class TestCelu(PytorchLayerTest):
def _prepare_input(self):
import numpy as np
return (np.random.randn(2, 4, 224, 224).astype(np.float32),)
@pytest.mark.nightly
@pytest.mark.precommit
@pytest.mark.precommit_torch_export
@pytest.mark.parametrize("alpha", [None, 0.5, 2.])
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64])
@pytest.mark.parametrize("inplace", [skip_if_export(True), False])
def test_celu(self, alpha, dtype, inplace, ie_device, precision, ir_version):
kwargs = {}
if dtype == torch.float16:
kwargs["custom_eps"] = 1e-2
self._test(aten_celu(alpha, dtype, inplace), None,
"aten::celu_" if inplace else "aten::celu", ie_device, precision, ir_version, **kwargs)

View File

@ -0,0 +1,45 @@
# Copyright (C) 2018-2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import pytest
import torch
import torch.nn.functional as F
from pytorch_layer_test_class import PytorchLayerTest, skip_if_export
class aten_elu(torch.nn.Module):
def __init__(self, alpha, dtype, inplace):
super(aten_elu, self).__init__()
self.alpha = alpha
self.dtype = dtype
self.inplace = inplace
if alpha is None:
self.forward = self.forward_no_alpha
def forward(self, x):
x_copy = x.to(self.dtype)
return F.elu(x_copy, alpha=self.alpha, inplace=self.inplace), x_copy
def forward_no_alpha(self, x):
x_copy = x.to(self.dtype)
return F.elu(x_copy, inplace=self.inplace), x_copy
class TestElu(PytorchLayerTest):
def _prepare_input(self):
import numpy as np
return (np.random.randn(2, 4, 224, 224).astype(np.float32),)
@pytest.mark.nightly
@pytest.mark.precommit
@pytest.mark.precommit_torch_export
@pytest.mark.parametrize("alpha", [None, 0.5, 2.])
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64])
@pytest.mark.parametrize("inplace", [skip_if_export(True), False])
def test_elu(self, alpha, dtype, inplace, ie_device, precision, ir_version):
kwargs = {}
if dtype == torch.float16:
kwargs["custom_eps"] = 1e-2
self._test(aten_elu(alpha, dtype, inplace), None,
"aten::elu_" if inplace else "aten::elu", ie_device, precision, ir_version, **kwargs)

View File

@ -29,6 +29,7 @@ class TestPixelShuffle(PytorchLayerTest):
(2, [1, 2, 3, 8, 4, 4]),])
@pytest.mark.nightly
@pytest.mark.precommit
@pytest.mark.precommit_torch_export
def test_pixel_shuffle(self, upscale_factor, shape, ie_device, precision, ir_version):
self.shape = shape
self._test(*self.create_model(upscale_factor),
@ -57,6 +58,7 @@ class TestPixelUnshuffle(PytorchLayerTest):
(2, [1, 2, 3, 2, 8, 8]),])
@pytest.mark.nightly
@pytest.mark.precommit
@pytest.mark.precommit_torch_export
def test_pixel_unshuffle(self, upscale_factor, shape, ie_device, precision, ir_version):
self.shape = shape
self._test(*self.create_model(upscale_factor),
@ -90,6 +92,7 @@ class TestChannelShuffle(PytorchLayerTest):
])
@pytest.mark.nightly
@pytest.mark.precommit
@pytest.mark.precommit_torch_export
def test_channel_shuffle(self, groups, shape, ie_device, precision, ir_version):
self.shape = shape
self._test(*self.create_model(groups),

View File

@ -42,6 +42,7 @@ OPS = {
"aten::log1p": torch.log1p,
"aten::log1p_": torch.log1p_,
"aten::log_sigmoid": F.logsigmoid,
"aten::mish": F.mish,
"aten::cos": torch.cos,
"aten::cos_": torch.cos_,
"aten::sin": torch.sin,
@ -178,6 +179,7 @@ class TestUnaryOp(PytorchLayerTest):
"aten::log2_",
"aten::log10_",
"aten::log1p_",
"aten::mish",
# trigonometry
"aten::cos_",
"aten::sin_",
@ -245,6 +247,7 @@ class TestUnaryOp(PytorchLayerTest):
"aten::relu6",
"aten::selu",
"aten::silu",
"aten::mish",
])
def test_unary_func_op_inplace(self, op_type, dtype, ie_device, precision, ir_version):
self.dtype = dtype