saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support input dimension swap in Attention op (#5774)

Browse source 
* checkin cpu

* checkin cpu

* add test

* cuda

* update comments

* review comments

* update

* modify var name

* remove unnecessary error msg

* fix comments

Co-authored-by: wangye <wangye@OrtDevTest2v100.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Ye Wang 2020-11-13 18:29:08 -08:00 committed by GitHub
parent dfbf6d78be
commit 262e9ef21d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 115 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
onnxruntime/contrib_ops/cpu/bert/attention.cc
View file

@ -46,6 +46,7 @@ AttentionBase::AttentionBase(const OpKernelInfo& info) {
num_heads_ = static_cast<int>(num_heads);
is_unidirectional_ = info.GetAttrOrDefault<int64_t>("unidirectional", 0) == 1;
is_input_dim_swapped_ = info.GetAttrOrDefault<int64_t>("input_dimension_swapped", 0) == 1;
}
Status AttentionBase::CheckInputs(const TensorShape& input_shape,
@ -54,7 +55,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
const Tensor*& mask_index,
const Tensor* past) const {
// Input shapes:
// input : (batch_size, sequence_length, hidden_size)
// input : (batch_size, sequence_length, hidden_size) or (sequence_length, batch_size, hidden_size)
// weights : (hidden_size, 3 * hidden_size)
// bias : (3 * hidden_size)
// mask_index : nullptr, (batch_size), (2 * batch_size), (batch_size, 1), (1, 1) or (batch_size, past_sequence_length + sequence_length)
@ -65,8 +66,8 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'input' is expected to have 3 dimensions, got ",
dims.size());
}
int batch_size = static_cast<int>(dims[0]);
int sequence_length = static_cast<int>(dims[1]);
int batch_size = is_input_dim_swapped_ ? static_cast<int>(dims[1]) : static_cast<int>(dims[0]);
int sequence_length = is_input_dim_swapped_ ? static_cast<int>(dims[0]) : static_cast<int>(dims[1]);
int hidden_size = static_cast<int>(dims[2]);
if (hidden_size % num_heads_ != 0) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
@ -233,8 +234,8 @@ Status Attention<T>::Compute(OpKernelContext* context) const {
past));
const auto& shape = input->Shape().GetDims();
const int batch_size = static_cast<int>(shape[0]);
const int sequence_length = static_cast<int>(shape[1]);
const int batch_size = is_input_dim_swapped_ ? static_cast<int>(shape[1]) : static_cast<int>(shape[0]);
const int sequence_length = is_input_dim_swapped_ ? static_cast<int>(shape[0]) : static_cast<int>(shape[1]);
const int hidden_size = static_cast<int>(shape[2]);
const int head_size = hidden_size / num_heads_;

5
onnxruntime/contrib_ops/cpu/bert/attention_base.h
View file

@ -26,8 +26,9 @@ class AttentionBase {
int sequence_length,
int& past_sequence_length) const;
int num_heads_; // number of attention heads
bool is_unidirectional_; // whether every token can only attend to previous tokens.
int num_heads_; // number of attention heads
bool is_unidirectional_; // whether every token can only attend to previous tokens.
bool is_input_dim_swapped_; // whether the input_shape is (S, B, NH) instead of (B, S, NH)
};
} // namespace contrib

8
onnxruntime/contrib_ops/cuda/bert/attention.cc
View file

@ -42,11 +42,11 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(), weights->Shape(), bias->Shape(), mask_index, past));
// Input and output shapes:
// Input 0 - input : (batch_size, sequence_length, hidden_size)
// Output 0 - output : (batch_size, sequence_length, hidden_size)
// Input 0 - input : (batch_size, sequence_length, hidden_size) or (sequence_length, batch_size, hidden_size)
// Output 0 - output : (batch_size, sequence_length, hidden_size) or (sequence_length, batch_size, hidden_size)
const auto& shape = input->Shape();
int batch_size = static_cast<int>(shape[0]);
int sequence_length = static_cast<int>(shape[1]);
int batch_size = is_input_dim_swapped_ ? static_cast<int>(shape[1]) : static_cast<int>(shape[0]);
int sequence_length = is_input_dim_swapped_ ? static_cast<int>(shape[0]) : static_cast<int>(shape[1]);
int hidden_size = static_cast<int>(shape[2]);
int head_size = hidden_size / num_heads_;

18
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

@ -406,7 +406,8 @@ we also support other two formats: When input has right-side padding, mask_index
where value of each element is the end position, or valid length of actual sequence excluding padding. When input has
left-side padding, mask_index has shape (2 * batch_size), where the values are the exclusive end positions followed by
the inclusive start positions. When unidirectional is 1, and each token only attend to previous tokens. For GPT-2, both past
and present state are optional. Present state could appear in output even when past state is not in input.
and present state are optional. Present state could appear in output even when past state is not in input. When
input_dimension_swapped is 1, the input shape is (sequence_length, batch_size, hidden_size) which happens in Bart models, etc.
)DOC";
ONNX_CONTRIB_OPERATOR_SCHEMA(Attention)
@ -418,12 +419,16 @@ and present state are optional. Present state could appear in output even when p
"Whether every token can only attend to previous tokens. Default value is 0.",
AttributeProto::INT,
static_cast<int64_t>(0))
.Input(0, "input", "3D input tensor with shape (batch_size, sequence_length, hidden_size), hidden_size = num_heads * head_size", "T")
.Attr("input_dimension_swapped",
"Whether input shape is (sequence_length, batch_size, hidden_size) instead of (batch_size, sequence_length, hidden_size). Default value is 0.",
AttributeProto::INT,
static_cast<int64_t>(0))
.Input(0, "input", "3D input tensor with shape (batch_size, sequence_length, hidden_size) or (sequence_length, batch_size, hidden_size), hidden_size = num_heads * head_size", "T")
.Input(1, "weight", "2D input tensor with shape (hidden_size, 3 * hidden_size)", "T")
.Input(2, "bias", "1D input tensor with shape (3 * hidden_size)", "T")
.Input(3, "mask_index", "Attention mask with shape (batch_size, past_sequence_length + sequence_length), or index with shape (batch_size) or (2 * batch_size).", "M", OpSchema::Optional)
.Input(4, "past", "past state for key and value with shape (2, batch_size, num_heads, past_sequence_length, head_size).", "T", OpSchema::Optional)
.Output(0, "output", "3D output tensor with shape (batch_size, append_length, hidden_size)", "T")
.Output(0, "output", "3D output tensor with shape (batch_size, append_length, hidden_size) or (sequence_length, batch_size, hidden_size)", "T")
.Output(1, "present", "present state for key and value with shape (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size)", "T", OpSchema::Optional)
.TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float tensors.")
.TypeConstraint("M", {"tensor(int32)"}, "Constrain mask index to integer types")
@ -451,6 +456,9 @@ and present state are optional. Present state could appear in output even when p
}
if (past_dims[3].has_dim_value() && input_dims[1].has_dim_value()) {
if (ctx.getAttribute("input_dimension_swapped")->i() != 0) {
fail_shape_inference("Past shall be work with input_dimension_swapped=0. aka when input shape equals to (B,S,NH)");
}
auto all_sequence_length = past_shape.dim(3).dim_value() + input_shape.dim(1).dim_value();
ONNX_NAMESPACE::TensorShapeProto present_shape;
@ -1879,7 +1887,7 @@ Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-
});
static const char* TransposeMatMul_doc = R"DOC(
Duplicate of FusedMatMul. Going forward FusedMatMul should be used. This OP will be supported for backward compatibility.
Duplicate of FusedMatMul. Going forward FusedMatMul should be used. This OP will be supported for backward compatibility.
Matrix product that behaves like numpy.matmul: https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.matmul.html
)DOC";
@ -2305,7 +2313,7 @@ and produces one output data (Tensor<T>) where the function `f(x) = quantize(alp
.TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput);
const char* QLinearSigmoidDoc_ver1 = R"DOC(
QLinearSigmoid takes quantized input data (Tensor), and quantize parameter for output, and produces one output data
QLinearSigmoid takes quantized input data (Tensor), and quantize parameter for output, and produces one output data
(Tensor<T>) where the function `f(x) = quantize(Sigmoid(dequantize(x)))`, is applied to the data tensor elementwise.
Wwhere the function `Sigmoid(x) = 1 / (1 + exp(-x))` )DOC";

112
onnxruntime/test/contrib_ops/attention_op_test.cc
View file

@ -17,18 +17,19 @@ enum MaskIndexType {
};
static void RunAttentionTest(
const std::vector<float>& input_data, // input: [batch_size, sequence_length, hidden_size]
const std::vector<float>& input_data, // input: [batch_size, sequence_length, hidden_size] or [sequence_length, batch_size, hidden_size]
const std::vector<float>& weights_data, // weights: [hidden_size, 3 * hidden_size]
bool is_weights_constant,
const std::vector<float>& bias_data, // bias: [3 * hidden_size]
const std::vector<int32_t>& mask_index_data, // mask_index: [batch_size] or [batch_size, past_sequence_length + sequence_length] or empty
const std::vector<float>& output_data, // output: [batch_size, sequence_length, hidden_size]
const std::vector<float>& output_data, // output: [batch_size, sequence_length, hidden_size] or [sequence_length, batch_size, hidden_size]
int batch_size,
int sequence_length,
int hidden_size,
int number_of_heads,
bool use_float16 = false,
bool is_unidirectional = false,
bool is_input_dimension_swapped = false,
bool use_past_state = false,
int past_sequence_length = 0,
const std::vector<float>* past_data = nullptr,
@ -43,6 +44,7 @@ static void RunAttentionTest(
OpTester tester("Attention", 1, onnxruntime::kMSDomain);
tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
tester.AddAttribute<int64_t>("unidirectional", static_cast<int64_t>(is_unidirectional ? 1 : 0));
tester.AddAttribute<int64_t>("input_dimension_swapped", static_cast<int64_t>(is_input_dimension_swapped ? 1 : 0));
std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
std::vector<int64_t> weights_dims = {hidden_size, 3 * hidden_size};
@ -122,17 +124,18 @@ static void RunAttentionTest(
}
static void RunAttentionTest(
const std::vector<float>& input_data, // input: [batch_size, sequence_length, hidden_size]
const std::vector<float>& input_data, // input: [batch_size, sequence_length, hidden_size] or [sequence_length, batch_size, hidden_size]
const std::vector<float>& weights_data, // weights: [hidden_size, 3 * hidden_size]
const std::vector<float>& bias_data, // bias: [3 * hidden_size]
const std::vector<int32_t>& mask_index_data, // mask_index: [batch_size] or [batch_size, past_sequence_length + sequence_length] or empty
const std::vector<float>& output_data, // output: [batch_size, sequence_length, hidden_size]
const std::vector<float>& output_data, // output: [batch_size, sequence_length, hidden_size] or [sequence_length, batch_size, hidden_size]
int batch_size,
int sequence_length,
int hidden_size,
int number_of_heads,
bool use_float16 = false,
bool is_unidirectional = false,
bool is_input_dimension_swapped = false,
bool use_past_state = false,
int past_sequence_length = 0,
const std::vector<float>* past_data = nullptr,
@ -140,12 +143,12 @@ static void RunAttentionTest(
MaskIndexType mask_index_type = kMaskIndexEnd) {
RunAttentionTest(input_data, weights_data, false, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length,
past_data, present_data, mask_index_type);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, mask_index_type);
RunAttentionTest(input_data, weights_data, true, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length,
past_data, present_data, mask_index_type);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, mask_index_type);
}
TEST(AttentionTest, AttentionBatch1) {
@ -329,6 +332,41 @@ TEST(AttentionTest, AttentionNoMaskIndex) {
batch_size, sequence_length, hidden_size, number_of_heads);
}
TEST(AttentionTest, AttentionNoMaskInputShapeSwapped) {
int batch_size_as_sequence_length = 1;
int sequence_length_as_batch_size = 2;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> weight_data = {
0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
std::vector<float> bias_data = {
-0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
// No mask_index
std::vector<int32_t> mask_index_data = {};
std::vector<float> output_data = {
8.6899995803833008f, -0.13000002503395081f, 4.25f, 5.6499996185302734f,
-4.0900001525878906f, 0.42000001668930054f, -0.10999995470046997f, 0.56999993324279785f};
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = true;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size_as_sequence_length, sequence_length_as_batch_size, hidden_size,
number_of_heads, use_float16, is_unidirectional, is_input_dimension_swapped);
}
TEST(AttentionTest, AttentionUnidirectional) {
int batch_size = 1;
int sequence_length = 2;
@ -506,12 +544,13 @@ TEST(AttentionTest, AttentionEmptyPastState) {
0.053175069391727448f, 0.12795503437519073f, 0.11125634610652924f, -0.0510881207883358f, -0.55345797538757324f, -0.3045809268951416f, -0.36920222640037537f, 0.060108467936515808f, 0.28109729290008545f, 0.069518551230430603f, 0.45718482136726379f, -0.010400654748082161f, 0.0038009658455848694f, 0.29213353991508484f, -0.17697516083717346f, 0.27086889743804932f};
bool is_unidirectional = true;
bool is_input_dimension_swapped = false;
bool use_past_state = true;
int past_sequence_length = 0;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional,
use_past_state, past_sequence_length, &past_data, &present_data);
is_input_dimension_swapped, use_past_state, past_sequence_length, &past_data, &present_data);
}
TEST(AttentionTest, AttentionPastStateBatch1) {
@ -605,12 +644,13 @@ TEST(AttentionTest, AttentionPastStateBatch1) {
0.45075402f, 0.85365993f, 0.097346395f, 0.28859729f, 0.26926181f, 0.65922296f, -0.027254611f, -0.096526355f, 0.8177433f, 0.4212271f, 0.34352475f, 0.059609573f, 0.46556228f, 0.7226882f, -0.025281552f, -0.25482416f};
bool is_unidirectional = true;
bool is_input_dimension_swapped = false;
bool use_past_state = true;
int past_sequence_length = 3;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional,
use_past_state, past_sequence_length, &past_data, &present_data);
is_input_dimension_swapped, use_past_state, past_sequence_length, &past_data, &present_data);
}
TEST(AttentionTest, AttentionPastStateBatch2) {
@ -708,12 +748,13 @@ TEST(AttentionTest, AttentionPastStateBatch2) {
0.3113814f, 0.63999802f, 0.28603253f, 0.98899829f, 0.044405211f, 0.95105386f, -0.033968594f, -0.034833729f, 0.81278932f, 0.63969064f, 0.14494057f, 0.11349615f, 0.87086016f, 0.20983537f, 0.045759238f, -0.26863033f, 0.35107401f, 0.90144604f, 0.68950737f, 0.18928574f, 0.18029204f, 0.074517399f, -0.033201858f, -0.10592631f, 0.70763874f, 0.48440042f, 0.58114725f, 0.1048766f, 0.73694098f, 0.17766342f, -0.054369561f, -0.24562015f};
bool is_unidirectional = true;
bool is_input_dimension_swapped = false;
bool use_past_state = true;
int past_sequence_length = 3;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional,
use_past_state, past_sequence_length, &past_data, &present_data);
is_input_dimension_swapped, use_past_state, past_sequence_length, &past_data, &present_data);
}
TEST(AttentionTest, AttentionPastStateBatch2WithPadding) {
@ -811,11 +852,13 @@ TEST(AttentionTest, AttentionPastStateBatch2WithPadding) {
0.3113814f, 0.63999802f, 0.28603253f, 0.98899829f, 0.044405211f, 0.95105386f, -0.033968594f, -0.034833729f, 0.81278932f, 0.63969064f, 0.14494057f, 0.11349615f, 0.87086016f, 0.20983537f, 0.045759238f, -0.26863033f, 0.35107401f, 0.90144604f, 0.68950737f, 0.18928574f, 0.18029204f, 0.074517399f, -0.033201858f, -0.10592631f, 0.70763874f, 0.48440042f, 0.58114725f, 0.1048766f, 0.73694098f, 0.17766342f, -0.054369561f, -0.24562015f};
bool is_unidirectional = true;
bool is_input_dimension_swapped = false;
bool use_past_state = true;
int past_sequence_length = 3;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional,
use_past_state, past_sequence_length, &past_data, &present_data, kMaskIndexEndAndStart);
is_input_dimension_swapped, use_past_state, past_sequence_length, &past_data, &present_data,
kMaskIndexEndAndStart);
}
TEST(AttentionTest, AttentionBatch2MaskIndex2) {
@ -849,13 +892,15 @@ TEST(AttentionTest, AttentionBatch2MaskIndex2) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
const std::vector<float>* present_data = nullptr;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
}
TEST(AttentionTest, AttentionRightPaddingMaskIndex2) {
@ -886,13 +931,15 @@ TEST(AttentionTest, AttentionRightPaddingMaskIndex2) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
const std::vector<float>* present_data = nullptr;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
}
TEST(AttentionTest, AttentionLeftPaddingMaskIndex2) {
@ -923,13 +970,15 @@ TEST(AttentionTest, AttentionLeftPaddingMaskIndex2) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
const std::vector<float>* present_data = nullptr;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
}
TEST(AttentionTest, AttentionBatch2LeftPaddingMaskIndex2) {
@ -964,13 +1013,15 @@ TEST(AttentionTest, AttentionBatch2LeftPaddingMaskIndex2) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
const std::vector<float>* present_data = nullptr;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
}
TEST(AttentionTest, AttentionBatch2AttentionMask) {
@ -1005,13 +1056,15 @@ TEST(AttentionTest, AttentionBatch2AttentionMask) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
const std::vector<float>* present_data = nullptr;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskRaw);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskRaw);
}
TEST(AttentionTest, AttentionUnidirectionalAttentionMask) {
@ -1046,13 +1099,15 @@ TEST(AttentionTest, AttentionUnidirectionalAttentionMask) {
bool use_float16 = false;
bool is_unidirectional = true;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
const std::vector<float>* present_data = nullptr;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskRaw);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskRaw);
}
TEST(AttentionTest, AttentionMask1DEndNoWord) {
@ -1087,6 +1142,7 @@ TEST(AttentionTest, AttentionMask1DEndNoWord) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
@ -1094,7 +1150,8 @@ TEST(AttentionTest, AttentionMask1DEndNoWord) {
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskIndexEnd);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskIndexEnd);
}
TEST(AttentionTest, AttentionMask1DNoWord) {
@ -1129,6 +1186,7 @@ TEST(AttentionTest, AttentionMask1DNoWord) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
@ -1136,7 +1194,8 @@ TEST(AttentionTest, AttentionMask1DNoWord) {
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
}
TEST(AttentionTest, AttentionMask2DNoWord) {
@ -1171,6 +1230,7 @@ TEST(AttentionTest, AttentionMask2DNoWord) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
@ -1178,7 +1238,8 @@ TEST(AttentionTest, AttentionMask2DNoWord) {
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskRaw);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskRaw);
}
@ -1213,6 +1274,7 @@ TEST(AttentionTest, AttentionDummyMask2D) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
@ -1220,7 +1282,8 @@ TEST(AttentionTest, AttentionDummyMask2D) {
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskDummy);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskDummy);
}
TEST(AttentionTest, AttentionMaskIndexOutOfRange) {
@ -1255,13 +1318,15 @@ TEST(AttentionTest, AttentionMaskIndexOutOfRange) {
bool use_float16 = false;
bool is_unidirectional = false;
bool is_input_dimension_swapped = false;
bool use_past_state = false;
int past_sequence_length = 0;
const std::vector<float>* past_data = nullptr;
const std::vector<float>* present_data = nullptr;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads,
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
use_float16, is_unidirectional, is_input_dimension_swapped, use_past_state,
past_sequence_length, past_data, present_data, kMaskIndexEndAndStart);
}
TEST(AttentionTest, AttentionPastState_dynamic) {
@ -1283,6 +1348,7 @@ TEST(AttentionTest, AttentionPastState_dynamic) {
OpTester test("Attention", 1, onnxruntime::kMSDomain);
test.AddAttribute<int64_t>("num_heads", 12);
test.AddAttribute<int64_t>("unidirectional", 1);
test.AddAttribute<int64_t>("input_dimension_swapped", 0);
test.AddInput<float>("input", input_dims, input_data);
test.AddInput<float>("weight", weight_dims, weight_data);
test.AddInput<float>("bias", bias_dims, bias_data);