Support aten::lstm, aten::gru and aten::rnn_{relu,tanh} operations (#21579)
* [PT FE] Initial LSTM impl
* Support aten::lstm operation
* Support aten::gru and aten::rnn_{relu,tanh}
* Apply suggestions from code review
* Fix build
* Simplify implementation
* More simplification
* Update src/frontends/pytorch/src/op/lstm.cpp
* Simplification
* Use default LSTM helpers
---------
Co-authored-by: Roboreptile <piotr.krzeminski@intel.com>
This commit is contained in:
Maxim Vafin
GitHub
parent
df029eb9c7
commit
0f6b9abee8
|
|
@ -0,0 +1,366 @@
|
|||
// Copyright (C) 2018-2023 Intel Corporation
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
//
|
||||
#include "openvino/frontend/pytorch/node_context.hpp"
|
||||
#include "openvino/op/add.hpp"
|
||||
#include "openvino/op/broadcast.hpp"
|
||||
#include "openvino/op/concat.hpp"
|
||||
#include "openvino/op/constant.hpp"
|
||||
#include "openvino/op/convert_like.hpp"
|
||||
#include "openvino/op/gather.hpp"
|
||||
#include "openvino/op/gru_sequence.hpp"
|
||||
#include "openvino/op/lstm_sequence.hpp"
|
||||
#include "openvino/op/multiply.hpp"
|
||||
#include "openvino/op/reshape.hpp"
|
||||
#include "openvino/op/rnn_sequence.hpp"
|
||||
#include "openvino/op/shape_of.hpp"
|
||||
#include "openvino/op/split.hpp"
|
||||
#include "openvino/op/squeeze.hpp"
|
||||
#include "openvino/op/transpose.hpp"
|
||||
#include "openvino/op/unsqueeze.hpp"
|
||||
#include "utils.hpp"
|
||||
|
||||
namespace ov {
|
||||
namespace frontend {
|
||||
namespace pytorch {
|
||||
namespace op {
|
||||
|
||||
using namespace ov::op;
|
||||
|
||||
namespace {
|
||||
enum RnnVariant { LSTM, GRU, RNN, RNN_RELU, RNN_TANH };
|
||||
|
||||
Output<Node> convert_data_format(ov::pass::NodeRegistry& rg, RnnVariant variant, const Output<Node>& node) {
|
||||
Output<Node> res;
|
||||
switch (variant) {
|
||||
case RnnVariant::LSTM:
|
||||
res = ov::op::util::convert_lstm_node_format(node, ov::op::util::LSTMWeightsFormat::IFCO);
|
||||
break;
|
||||
case RnnVariant::GRU:
|
||||
res = ov::op::util::convert_lstm_peepholes_format(node, ov::op::util::LSTMPeepholesFormat::IFO);
|
||||
break;
|
||||
default:
|
||||
res = node;
|
||||
break;
|
||||
}
|
||||
const auto axis_const = rg.make<v0::Constant>(element::i32, ov::Shape{}, 0);
|
||||
return rg.make<v0::Unsqueeze>(res, axis_const);
|
||||
}
|
||||
|
||||
Output<Node> format_bias(ov::pass::NodeRegistry& rg,
|
||||
RnnVariant variant,
|
||||
const Output<Node>& b_ih,
|
||||
const Output<Node>& b_hh) {
|
||||
Output<Node> res;
|
||||
if (variant == RnnVariant::GRU) {
|
||||
const auto one = v0::Constant::create(element::i32, Shape{}, {1});
|
||||
const auto bias_ih = convert_data_format(rg, variant, b_ih);
|
||||
const auto bias_hh = convert_data_format(rg, variant, b_hh);
|
||||
const auto split_bias_ih = rg.make<v1::Split>(bias_ih, one, 3);
|
||||
const auto split_bias_hh = rg.make<v1::Split>(bias_hh, one, 3);
|
||||
const auto wr_z_bias = rg.make<v1::Add>(split_bias_ih->output(0), split_bias_hh->output(0));
|
||||
const auto wr_r_bias = rg.make<v1::Add>(split_bias_ih->output(1), split_bias_hh->output(1));
|
||||
// The result has shape: [num_directions, 4 * hidden_size]
|
||||
// and data layout: [ [Wb_z + Rb_z], [Wb_r + Rb_r], [Wb_h], [Rb_h], ]
|
||||
res =
|
||||
rg.make<v0::Concat>(OutputVector{wr_z_bias, wr_r_bias, split_bias_ih->output(2), split_bias_hh->output(2)},
|
||||
1);
|
||||
} else {
|
||||
res = rg.make<v1::Add>(b_ih, b_hh);
|
||||
res = convert_data_format(rg, variant, res);
|
||||
}
|
||||
return res;
|
||||
}
|
||||
|
||||
OutputVector generic_rnn(ov::pass::NodeRegistry& rg,
|
||||
RnnVariant variant,
|
||||
const Output<Node>& input,
|
||||
const std::deque<Output<Node>>& initial_states,
|
||||
const std::deque<Output<Node>>& all_weights,
|
||||
bool has_biases,
|
||||
int64_t num_layers,
|
||||
bool bidirectional,
|
||||
bool batch_first,
|
||||
const Output<Node>& batch_sizes = {}) {
|
||||
std::string rnn_activation;
|
||||
if (variant == RnnVariant::RNN_RELU) {
|
||||
variant = RnnVariant::RNN;
|
||||
rnn_activation = "relu";
|
||||
} else if (variant == RnnVariant::RNN_TANH) {
|
||||
variant = RnnVariant::RNN;
|
||||
rnn_activation = "tanh";
|
||||
}
|
||||
const auto direction =
|
||||
bidirectional ? RecurrentSequenceDirection::BIDIRECTIONAL : RecurrentSequenceDirection::FORWARD;
|
||||
int64_t weights_per_layer = has_biases ? 4 : 2;
|
||||
int64_t mult = bidirectional ? 2 : 1;
|
||||
FRONT_END_OP_CONVERSION_CHECK(static_cast<int64_t>(all_weights.size()) == num_layers * weights_per_layer * mult,
|
||||
"Unexpected length of list with weights for rnn operation.");
|
||||
|
||||
const auto w_hh = all_weights[1];
|
||||
const auto w_hh_pshape = w_hh.get_partial_shape();
|
||||
FRONT_END_OP_CONVERSION_CHECK(w_hh_pshape.rank().is_static() && w_hh_pshape[1].is_static(), "");
|
||||
const auto hidden_size = w_hh_pshape[1].get_length();
|
||||
|
||||
const auto zero = v0::Constant::create(element::i32, Shape{}, {0});
|
||||
const auto zero_1d = v0::Constant::create(element::i32, Shape{1}, {0});
|
||||
const auto one = v0::Constant::create(element::i32, Shape{}, {1});
|
||||
const auto one_1d = v0::Constant::create(element::i32, Shape{1}, {1});
|
||||
const auto order_102 = v0::Constant::create(element::i32, Shape{3}, {1, 0, 2});
|
||||
|
||||
OutputVector h_outs;
|
||||
OutputVector c_outs;
|
||||
Output<Node> h0;
|
||||
Output<Node> c0;
|
||||
if (variant == RnnVariant::RNN || variant == RnnVariant::GRU) {
|
||||
h0 = initial_states[0];
|
||||
} else if (variant == RnnVariant::LSTM) {
|
||||
h0 = initial_states[0];
|
||||
c0 = initial_states[1];
|
||||
} else {
|
||||
FRONT_END_OP_CONVERSION_CHECK(false, "Unsupported rnn variant.");
|
||||
}
|
||||
|
||||
Output<Node> prev_output = input;
|
||||
if (!batch_first)
|
||||
prev_output = rg.make<v1::Transpose>(prev_output, order_102);
|
||||
Output<Node> sequence_lens = batch_sizes;
|
||||
|
||||
const auto h_states = rg.make<v1::Split>(h0, zero, num_layers)->outputs();
|
||||
OutputVector c_states;
|
||||
if (variant == RnnVariant::LSTM) {
|
||||
c_states = rg.make<v1::Split>(c0, zero, num_layers)->outputs();
|
||||
}
|
||||
|
||||
Output<Node> bias_concat;
|
||||
Shape::value_type num_directions = bidirectional ? 2 : 1;
|
||||
if (!has_biases) {
|
||||
Shape::value_type gates_count = variant == RnnVariant::RNN ? 1 : 4;
|
||||
Shape::value_type gates_hidden = gates_count * static_cast<Shape::value_type>(hidden_size);
|
||||
bias_concat = rg.make<v0::Constant>(element::i32, Shape{num_directions, gates_hidden}, 0);
|
||||
bias_concat = rg.make<v1::ConvertLike>(bias_concat, input);
|
||||
}
|
||||
|
||||
for (int64_t i = 0; i < num_layers; i++) {
|
||||
Output<Node> weight_ih;
|
||||
Output<Node> weight_hh;
|
||||
|
||||
int64_t idx = i * weights_per_layer;
|
||||
if (!bidirectional) {
|
||||
weight_ih = convert_data_format(rg, variant, all_weights[idx]);
|
||||
weight_hh = convert_data_format(rg, variant, all_weights[idx + 1]);
|
||||
if (has_biases) {
|
||||
const auto bias_ih = all_weights[idx + 2];
|
||||
const auto bias_hh = all_weights[idx + 3];
|
||||
bias_concat = format_bias(rg, variant, bias_ih, bias_hh);
|
||||
}
|
||||
} else {
|
||||
Output<Node> weight_ih_f;
|
||||
Output<Node> weight_hh_f;
|
||||
Output<Node> weight_ih_b;
|
||||
Output<Node> weight_hh_b;
|
||||
if (has_biases) {
|
||||
weight_ih_f = all_weights[2 * idx];
|
||||
weight_hh_f = all_weights[2 * idx + 1];
|
||||
const auto bias_ih_f = all_weights[2 * idx + 2];
|
||||
const auto bias_hh_f = all_weights[2 * idx + 3];
|
||||
weight_ih_b = all_weights[2 * idx + 4];
|
||||
weight_hh_b = all_weights[2 * idx + 5];
|
||||
const auto bias_ih_b = all_weights[2 * idx + 6];
|
||||
const auto bias_hh_b = all_weights[2 * idx + 7];
|
||||
const auto bias_f = format_bias(rg, variant, bias_ih_f, bias_hh_f);
|
||||
const auto bias_b = format_bias(rg, variant, bias_ih_b, bias_hh_b);
|
||||
bias_concat = rg.make<v0::Concat>(OutputVector{bias_f, bias_b}, 0);
|
||||
} else {
|
||||
weight_ih_f = all_weights[2 * idx];
|
||||
weight_hh_f = all_weights[2 * idx + 1];
|
||||
weight_ih_b = all_weights[2 * idx + 2];
|
||||
weight_hh_b = all_weights[2 * idx + 3];
|
||||
}
|
||||
weight_ih_f = convert_data_format(rg, variant, weight_ih_f);
|
||||
weight_hh_f = convert_data_format(rg, variant, weight_hh_f);
|
||||
weight_ih_b = convert_data_format(rg, variant, weight_ih_b);
|
||||
weight_hh_b = convert_data_format(rg, variant, weight_hh_b);
|
||||
weight_ih = rg.make<v0::Concat>(OutputVector{weight_ih_f, weight_ih_b}, 0);
|
||||
weight_hh = rg.make<v0::Concat>(OutputVector{weight_hh_f, weight_hh_b}, 0);
|
||||
}
|
||||
|
||||
const auto shape_of_x = rg.make<v3::ShapeOf>(prev_output, element::i32);
|
||||
const auto axes = v0::Constant::create(element::i32, Shape{1}, {0});
|
||||
const auto batch_size_node = rg.make<v8::Gather>(shape_of_x, zero_1d, axes);
|
||||
if (!sequence_lens.get_node_shared_ptr()) {
|
||||
const auto seq_length_node = rg.make<v8::Gather>(shape_of_x, one_1d, axes);
|
||||
sequence_lens = rg.make<v3::Broadcast>(seq_length_node, batch_size_node);
|
||||
}
|
||||
|
||||
const auto h_state = rg.make<v1::Transpose>(h_states[i], order_102);
|
||||
std::shared_ptr<Node> rnn_node;
|
||||
if (variant == RnnVariant::GRU) {
|
||||
rnn_node = rg.make<v5::GRUSequence>(prev_output,
|
||||
h_state,
|
||||
sequence_lens,
|
||||
weight_ih,
|
||||
weight_hh,
|
||||
bias_concat,
|
||||
hidden_size,
|
||||
direction,
|
||||
std::vector<std::string>{"sigmoid", "tanh"},
|
||||
std::vector<float>{},
|
||||
std::vector<float>{},
|
||||
0.f,
|
||||
true);
|
||||
} else if (variant == RnnVariant::LSTM) {
|
||||
Output<Node> c_state = rg.make<v1::Transpose>(c_states[i], order_102);
|
||||
rnn_node = rg.make<v5::LSTMSequence>(prev_output,
|
||||
h_state,
|
||||
c_state,
|
||||
sequence_lens,
|
||||
weight_ih,
|
||||
weight_hh,
|
||||
bias_concat,
|
||||
hidden_size,
|
||||
direction);
|
||||
} else if (variant == RnnVariant::RNN) {
|
||||
rnn_node = rg.make<v5::RNNSequence>(prev_output,
|
||||
h_state,
|
||||
sequence_lens,
|
||||
weight_ih,
|
||||
weight_hh,
|
||||
bias_concat,
|
||||
hidden_size,
|
||||
direction,
|
||||
std::vector<std::string>{rnn_activation});
|
||||
}
|
||||
prev_output = rnn_node->output(0);
|
||||
|
||||
if (bidirectional) {
|
||||
const auto order = v0::Constant::create(element::i32, Shape{4}, {0, 2, 1, 3});
|
||||
prev_output = rg.make<v1::Transpose>(prev_output, order);
|
||||
const auto new_shape = v0::Constant::create(element::i32, Shape{3}, {0, 0, -1});
|
||||
prev_output = rg.make<v1::Reshape>(prev_output, new_shape, true);
|
||||
} else {
|
||||
prev_output = rg.make<v0::Squeeze>(prev_output, one);
|
||||
}
|
||||
|
||||
h_outs.push_back(rnn_node->output(1));
|
||||
if (variant == RnnVariant::LSTM)
|
||||
c_outs.push_back(rnn_node->output(2));
|
||||
}
|
||||
if (!batch_first)
|
||||
prev_output = rg.make<v1::Transpose>(prev_output, order_102);
|
||||
Output<Node> h_res = rg.make<v0::Concat>(h_outs, 1);
|
||||
h_res = rg.make<v1::Transpose>(h_res, order_102);
|
||||
if (variant == RnnVariant::RNN || variant == RnnVariant::GRU) {
|
||||
return {prev_output, h_res};
|
||||
} else if (variant == RnnVariant::LSTM) {
|
||||
Output<Node> c_res = rg.make<v0::Concat>(c_outs, 1);
|
||||
c_res = rg.make<v1::Transpose>(c_res, order_102);
|
||||
return {prev_output, h_res, c_res};
|
||||
}
|
||||
FRONT_END_OP_CONVERSION_CHECK(false, "Unsupported rnn variant.");
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
OutputVector translate_lstm(const NodeContext& context) {
|
||||
num_inputs_check(context, 9, 9);
|
||||
ov::pass::NodeRegistry rg;
|
||||
if (context.get_input_type(3).is<type::List>()) {
|
||||
// lstm packed
|
||||
FRONT_END_OP_CONVERSION_CHECK(false, "Unsupported lstm variant.");
|
||||
} else {
|
||||
// aten::lstm.input(Tensor input, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout,
|
||||
// bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor, Tensor)
|
||||
const auto data = context.get_input(0);
|
||||
const auto hx = context.get_input(1);
|
||||
const auto params = context.get_input(2);
|
||||
const auto has_bias = context.const_input<bool>(3);
|
||||
const auto num_layers = context.const_input<int64_t>(4);
|
||||
// const auto dropout = context.const_input<float>(5); - skip
|
||||
const auto train = context.const_input<bool>(6);
|
||||
FRONT_END_OP_CONVERSION_CHECK(!train, "LSTM in train mode is not supported.");
|
||||
const auto bidirectional = context.const_input<bool>(7);
|
||||
const auto batch_first = context.const_input<bool>(8);
|
||||
|
||||
const auto initial_states = get_list_as_outputs(hx);
|
||||
const auto all_weights = get_list_as_outputs(params);
|
||||
const auto res = generic_rnn(rg,
|
||||
RnnVariant::LSTM,
|
||||
data,
|
||||
initial_states,
|
||||
all_weights,
|
||||
has_bias,
|
||||
num_layers,
|
||||
bidirectional,
|
||||
batch_first);
|
||||
context.mark_nodes(rg.get());
|
||||
return res;
|
||||
}
|
||||
};
|
||||
|
||||
OutputVector translate_gru(const NodeContext& context) {
|
||||
num_inputs_check(context, 9, 9);
|
||||
ov::pass::NodeRegistry rg;
|
||||
// aten::gru.input(Tensor input, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout,
|
||||
// bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor)
|
||||
const auto input = context.get_input(0);
|
||||
const auto hidden = context.get_input(1);
|
||||
const auto weight_v = context.get_input(2);
|
||||
const auto has_biases = context.const_input<bool>(3);
|
||||
const auto num_layers = context.const_input<int64_t>(4);
|
||||
// const auto dropout = context.const_input<float>(5); - skip
|
||||
const auto train = context.const_input<bool>(6);
|
||||
FRONT_END_OP_CONVERSION_CHECK(!train, "GRU in train mode is not supported.");
|
||||
const auto bidirectional = context.const_input<bool>(7);
|
||||
const auto batch_first = context.const_input<bool>(8);
|
||||
|
||||
const auto weight = get_list_as_outputs(weight_v);
|
||||
const auto res =
|
||||
generic_rnn(rg, RnnVariant::GRU, input, {hidden}, weight, has_biases, num_layers, bidirectional, batch_first);
|
||||
context.mark_nodes(rg.get());
|
||||
return res;
|
||||
};
|
||||
|
||||
namespace {
|
||||
std::map<std::string, RnnVariant> RNN_VARIANT_MAP = {
|
||||
{"aten::rnn_tanh", RnnVariant::RNN_TANH},
|
||||
{"aten::rnn_relu", RnnVariant::RNN_RELU},
|
||||
};
|
||||
}
|
||||
|
||||
OutputVector translate_rnn(const NodeContext& context) {
|
||||
num_inputs_check(context, 9, 9);
|
||||
ov::pass::NodeRegistry rg;
|
||||
// aten::rnn_relu.input(Tensor input, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout,
|
||||
// bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor)
|
||||
const auto input = context.get_input(0);
|
||||
const auto hidden = context.get_input(1);
|
||||
const auto weight_v = context.get_input(2);
|
||||
const auto has_biases = context.const_input<bool>(3);
|
||||
const auto num_layers = context.const_input<int64_t>(4);
|
||||
// const auto dropout = context.const_input<float>(5); - skip
|
||||
const auto train = context.const_input<bool>(6);
|
||||
FRONT_END_OP_CONVERSION_CHECK(!train, "RNN in train mode is not supported.");
|
||||
const auto bidirectional = context.const_input<bool>(7);
|
||||
const auto batch_first = context.const_input<bool>(8);
|
||||
|
||||
const auto weight = get_list_as_outputs(weight_v);
|
||||
const auto variant_it = RNN_VARIANT_MAP.find(context.get_op_type());
|
||||
FRONT_END_OP_CONVERSION_CHECK(variant_it != RNN_VARIANT_MAP.end(), "Unsupported RNN variant.");
|
||||
const auto res = generic_rnn(rg,
|
||||
variant_it->second,
|
||||
input,
|
||||
{hidden},
|
||||
weight,
|
||||
has_biases,
|
||||
num_layers,
|
||||
bidirectional,
|
||||
batch_first);
|
||||
context.mark_nodes(rg.get());
|
||||
return res;
|
||||
};
|
||||
|
||||
} // namespace op
|
||||
} // namespace pytorch
|
||||
} // namespace frontend
|
||||
} // namespace ov
|
||||
|
|
@ -89,6 +89,7 @@ OP_CONVERTER(translate_getitem);
|
|||
OP_CONVERTER(translate_glu);
|
||||
OP_CONVERTER(translate_grid_sampler);
|
||||
OP_CONVERTER(translate_group_norm);
|
||||
OP_CONVERTER(translate_gru);
|
||||
OP_CONVERTER(translate_hardtanh);
|
||||
OP_CONVERTER(translate_if);
|
||||
OP_CONVERTER(translate_im2col);
|
||||
|
|
@ -114,6 +115,7 @@ OP_CONVERTER(translate_log2);
|
|||
OP_CONVERTER(translate_log10);
|
||||
OP_CONVERTER(translate_logsumexp);
|
||||
OP_CONVERTER(translate_loop);
|
||||
OP_CONVERTER(translate_lstm);
|
||||
OP_CONVERTER(translate_masked_fill);
|
||||
OP_CONVERTER(translate_masked_scatter);
|
||||
OP_CONVERTER(translate_max);
|
||||
|
|
@ -167,6 +169,7 @@ OP_CONVERTER(translate_remainder);
|
|||
OP_CONVERTER(translate_repeat_interleave);
|
||||
OP_CONVERTER(translate_reshape);
|
||||
OP_CONVERTER(translate_reshape_as);
|
||||
OP_CONVERTER(translate_rnn);
|
||||
OP_CONVERTER(translate_roi_align);
|
||||
OP_CONVERTER(translate_roll);
|
||||
OP_CONVERTER(translate_round);
|
||||
|
|
@ -372,6 +375,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
|
|||
{"aten::glu", op::translate_glu},
|
||||
{"aten::grid_sampler", op::translate_grid_sampler},
|
||||
{"aten::group_norm", op::translate_group_norm},
|
||||
{"aten::gru", op::translate_gru},
|
||||
{"aten::gt", op::translate_1to1_match_2_inputs_align_types<opset10::Greater>},
|
||||
{"aten::hardsigmoid", op::quantizable_op<op::translate_1to1_match_1_inputs<opset10::HSigmoid>>},
|
||||
{"aten::hardswish", op::quantizable_op<op::translate_1to1_match_1_inputs<opset10::HSwish>>},
|
||||
|
|
@ -415,6 +419,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
|
|||
{"aten::log2_", op::inplace_op<op::translate_log2>},
|
||||
{"aten::log10", op::translate_log10},
|
||||
{"aten::log10_", op::inplace_op<op::translate_log10>},
|
||||
{"aten::lstm", op::translate_lstm},
|
||||
{"aten::lt", op::translate_1to1_match_2_inputs_align_types<opset10::Less>},
|
||||
{"aten::masked_fill", op::translate_masked_fill},
|
||||
{"aten::masked_fill_", op::inplace_op<op::translate_masked_fill>},
|
||||
|
|
@ -484,6 +489,8 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
|
|||
// for real dtypes, these operations return input tensor without changes and can be skipped
|
||||
{"aten::resolve_conj", op::skip_node},
|
||||
{"aten::resolve_neg", op::skip_node},
|
||||
{"aten::rnn_relu", op::translate_rnn},
|
||||
{"aten::rnn_tanh", op::translate_rnn},
|
||||
{"aten::roll", op::translate_roll},
|
||||
{"aten::round", op::translate_round},
|
||||
{"aten::rsqrt", op::translate_rsqrt},
|
||||
|
|
|
|||
|
|
@ -0,0 +1,140 @@
|
|||
# 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 aten_lstm(torch.nn.Module):
|
||||
def __init__(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first):
|
||||
torch.nn.Module.__init__(self)
|
||||
self.lstm = torch.nn.LSTM(input_size,
|
||||
hidden_size,
|
||||
num_layers,
|
||||
has_bias,
|
||||
batch_first,
|
||||
bidirectional=bidirectional)
|
||||
|
||||
def forward(self, input_tensor, h0, c0):
|
||||
return self.lstm(input_tensor, (h0, c0))
|
||||
|
||||
|
||||
class aten_gru(torch.nn.Module):
|
||||
def __init__(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first):
|
||||
torch.nn.Module.__init__(self)
|
||||
self.gru = torch.nn.GRU(input_size,
|
||||
hidden_size,
|
||||
num_layers,
|
||||
has_bias,
|
||||
batch_first,
|
||||
bidirectional=bidirectional)
|
||||
|
||||
def forward(self, input_tensor, h0):
|
||||
return self.gru(input_tensor, h0)
|
||||
|
||||
|
||||
class aten_rnn(torch.nn.Module):
|
||||
def __init__(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, nonlinearity):
|
||||
torch.nn.Module.__init__(self)
|
||||
self.rnn = torch.nn.RNN(input_size,
|
||||
hidden_size,
|
||||
num_layers,
|
||||
nonlinearity=nonlinearity,
|
||||
bias=has_bias,
|
||||
batch_first=batch_first,
|
||||
bidirectional=bidirectional)
|
||||
|
||||
def forward(self, input_tensor, h0):
|
||||
return self.rnn(input_tensor, h0)
|
||||
|
||||
|
||||
class TestLSTM(PytorchLayerTest):
|
||||
def _prepare_input(self):
|
||||
n = self.num_layers
|
||||
if self.bidirectional:
|
||||
n *= 2
|
||||
if self.batch_first:
|
||||
input = np.random.randn(3, 5, self.input_size).astype(np.float32)
|
||||
else:
|
||||
input = np.random.randn(5, 3, self.input_size).astype(np.float32)
|
||||
h0 = np.random.randn(n, 3, self.hidden_size).astype(np.float32)
|
||||
c0 = np.random.randn(n, 3, self.hidden_size).astype(np.float32)
|
||||
return (input, h0, c0)
|
||||
|
||||
@pytest.mark.parametrize("input_size,hidden_size", [(10, 20),])
|
||||
@pytest.mark.parametrize("num_layers", [1, 2, 7])
|
||||
@pytest.mark.parametrize("has_bias", [True, False])
|
||||
@pytest.mark.parametrize("bidirectional", [True, False])
|
||||
@pytest.mark.parametrize("batch_first", [True, False])
|
||||
@pytest.mark.nightly
|
||||
@pytest.mark.precommit
|
||||
def test_lstm(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, ie_device, precision, ir_version):
|
||||
self.input_size = input_size
|
||||
self.hidden_size = hidden_size
|
||||
self.num_layers = num_layers
|
||||
self.bidirectional = bidirectional
|
||||
self.batch_first = batch_first
|
||||
self._test(aten_lstm(input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first), None, "aten::lstm",
|
||||
ie_device, precision, ir_version, trace_model=True)
|
||||
|
||||
|
||||
class TestGRU(PytorchLayerTest):
|
||||
def _prepare_input(self):
|
||||
n = self.num_layers
|
||||
if self.bidirectional:
|
||||
n *= 2
|
||||
if self.batch_first:
|
||||
input = np.random.randn(3, 5, self.input_size).astype(np.float32)
|
||||
else:
|
||||
input = np.random.randn(5, 3, self.input_size).astype(np.float32)
|
||||
h0 = np.random.randn(n, 3, self.hidden_size).astype(np.float32)
|
||||
return (input, h0)
|
||||
|
||||
@pytest.mark.parametrize("input_size,hidden_size", [(10, 20),])
|
||||
@pytest.mark.parametrize("num_layers", [1, 2, 7])
|
||||
@pytest.mark.parametrize("has_bias", [True, False])
|
||||
@pytest.mark.parametrize("bidirectional", [True, False])
|
||||
@pytest.mark.parametrize("batch_first", [True, False])
|
||||
@pytest.mark.nightly
|
||||
@pytest.mark.precommit
|
||||
def test_gru(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, ie_device, precision, ir_version):
|
||||
self.input_size = input_size
|
||||
self.hidden_size = hidden_size
|
||||
self.num_layers = num_layers
|
||||
self.bidirectional = bidirectional
|
||||
self.batch_first = batch_first
|
||||
self._test(aten_gru(input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first), None, "aten::gru",
|
||||
ie_device, precision, ir_version, trace_model=True)
|
||||
|
||||
|
||||
class TestRNN(PytorchLayerTest):
|
||||
def _prepare_input(self):
|
||||
n = self.num_layers
|
||||
if self.bidirectional:
|
||||
n *= 2
|
||||
if self.batch_first:
|
||||
input = np.random.randn(3, 5, self.input_size).astype(np.float32)
|
||||
else:
|
||||
input = np.random.randn(5, 3, self.input_size).astype(np.float32)
|
||||
h0 = np.random.randn(n, 3, self.hidden_size).astype(np.float32)
|
||||
return (input, h0)
|
||||
|
||||
@pytest.mark.parametrize("input_size,hidden_size", [(10, 20),])
|
||||
@pytest.mark.parametrize("num_layers", [1, 2, 7])
|
||||
@pytest.mark.parametrize("has_bias", [True, False])
|
||||
@pytest.mark.parametrize("bidirectional", [True, False])
|
||||
@pytest.mark.parametrize("batch_first", [True, False])
|
||||
@pytest.mark.parametrize("nonlinearity", ["tanh", "relu"])
|
||||
@pytest.mark.nightly
|
||||
@pytest.mark.precommit
|
||||
def test_rnn(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, nonlinearity, ie_device, precision, ir_version):
|
||||
self.input_size = input_size
|
||||
self.hidden_size = hidden_size
|
||||
self.num_layers = num_layers
|
||||
self.bidirectional = bidirectional
|
||||
self.batch_first = batch_first
|
||||
self._test(aten_rnn(input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, nonlinearity), None, f"aten::rnn_{nonlinearity}",
|
||||
ie_device, precision, ir_version, trace_model=True)
|
||||
|
|
@ -67,7 +67,7 @@ facebook/convnextv2-tiny-22k-384,convnextv2
|
|||
facebook/detr-resnet-50,detr
|
||||
facebook/dinov2-base,dinov2,xfail,Tracing error: Please check correctness of provided example_input (but eval was correct)
|
||||
facebook/dpr-question_encoder-single-nq-base,dpr
|
||||
facebook/encodec_24khz,encodec,xfail,Unsupported op aten::lstm
|
||||
facebook/encodec_24khz,encodec
|
||||
facebook/esm2_t6_8M_UR50D,esm
|
||||
facebook/flava-full,flava,xfail,Tracing problem
|
||||
facebook/flava-image-codebook,flava_image_codebook,skip,Load problem
|
||||
|
|
|
|||
|
|
@ -68,7 +68,8 @@ class TestTimmConvertModel(TestTorchConvertModel):
|
|||
@pytest.mark.parametrize("name", ["mobilevitv2_050.cvnets_in1k",
|
||||
"poolformerv2_s12.sail_in1k",
|
||||
"vit_base_patch8_224.augreg_in21k",
|
||||
"beit_base_patch16_224.in22k_ft_in22k"])
|
||||
"beit_base_patch16_224.in22k_ft_in22k",
|
||||
"sequencer2d_l.in1k"])
|
||||
@pytest.mark.precommit
|
||||
def test_convert_model_precommit(self, name, ie_device):
|
||||
self.run(name, None, ie_device)
|
||||
|
|
|
|||
|
|
@ -385,7 +385,7 @@ selecsls60.in1k,None
|
|||
selecsls60b.in1k,None
|
||||
semnasnet_075.rmsp_in1k,None
|
||||
senet154.gluon_in1k,None
|
||||
sequencer2d_l.in1k,None,xfail,Unsupported aten::lstm
|
||||
sequencer2d_l.in1k,None
|
||||
seresnet152d.ra2_in1k,None
|
||||
seresnet33ts.ra2_in1k,None
|
||||
seresnet50.a1_in1k,None
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ basic_gnn_gin,None,xfail,Unsupported op aten::scatter_add_
|
|||
basic_gnn_sage,None,xfail,Unsupported op aten::scatter_add_
|
||||
#cm3leon_generate,None,skip,No install.py is found
|
||||
dcgan,None
|
||||
demucs,None,xfail,Unsupported op aten::lstm
|
||||
demucs,None
|
||||
#densenet121,None - Already tested by torchvision tests
|
||||
#detectron2_fasterrcnn_r_101_c4,None - Already tested by det2 tests
|
||||
#detectron2_fasterrcnn_r_101_dc5,None - Already tested by det2 tests
|
||||
|
|
@ -93,7 +93,7 @@ tacotron2,None,skip,Can't be loaded without CUDA
|
|||
#timm_vision_transformer_large,None - Already tested by timm tests
|
||||
#timm_vovnet,None - Already tested by timm tests
|
||||
torch_multimodal_clip,None,skip,Can't be traced
|
||||
tts_angular,None,xfail,Unsupported op aten::lstm
|
||||
tts_angular,None
|
||||
#vgg16,None - Already tested by torchvision tests
|
||||
#vision_maskrcnn,None,skip,Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions
|
||||
yolov3,None
|
||||
Loading…
Reference in New Issue