From b89f52c2773e27e1f6999f2b2c6aaf33f3f96380 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 8 Mar 2021 11:56:31 -0800 Subject: [PATCH] Add tests of Attention and QAttention for pruned model (#6914) --- .../test/contrib_ops/attention_op_test.cc | 66 ++++++++++++-- .../contrib_ops/quantize_attention_op_test.cc | 86 +++++++++++++++---- 2 files changed, 127 insertions(+), 25 deletions(-) diff --git a/onnxruntime/test/contrib_ops/attention_op_test.cc b/onnxruntime/test/contrib_ops/attention_op_test.cc index 6abe23599b..b4fec55736 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test.cc @@ -34,19 +34,22 @@ static void RunAttentionTest( int past_sequence_length = 0, const std::vector* past_data = nullptr, const std::vector* present_data = nullptr, - MaskIndexType mask_index_type = kMaskIndexEnd) { - int min_cuda_architecture = use_float16 ? 530 : 0; + MaskIndexType mask_index_type = kMaskIndexEnd, + int input_hidden_size = 0) { + input_hidden_size = (input_hidden_size == 0 ? hidden_size : input_hidden_size); // By default, no pruning. + int min_cuda_architecture = use_float16 ? 530 : 0; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture) && !is_weights_constant; bool enable_cpu = (nullptr != DefaultCpuExecutionProvider().get()) && !use_float16; + int head_size = hidden_size / number_of_heads; if (enable_cpu || enable_cuda) { OpTester tester("Attention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(number_of_heads)); tester.AddAttribute("unidirectional", static_cast(is_unidirectional ? 1 : 0)); - std::vector input_dims = {batch_size, sequence_length, hidden_size}; - std::vector weights_dims = {hidden_size, 3 * hidden_size}; + std::vector input_dims = {batch_size, sequence_length, input_hidden_size}; + std::vector weights_dims = {input_hidden_size, 3 * hidden_size}; std::vector bias_dims = {3 * hidden_size}; std::vector mask_index_dims_1 = {batch_size}; @@ -78,7 +81,7 @@ static void RunAttentionTest( std::vector past_dims = {2, batch_size, number_of_heads, past_sequence_length, head_size}; std::vector present_dims = {2, batch_size, number_of_heads, past_sequence_length + sequence_length, head_size}; - std::vector output_dims = input_dims; + std::vector output_dims = {batch_size, sequence_length, hidden_size}; if (use_float16) { tester.AddInput("input", input_dims, ToFloat16(input_data)); @@ -142,15 +145,16 @@ static void RunAttentionTest( int past_sequence_length = 0, const std::vector* past_data = nullptr, const std::vector* present_data = nullptr, - MaskIndexType mask_index_type = kMaskIndexEnd) { + MaskIndexType mask_index_type = kMaskIndexEnd, + int input_hidden_size = 0) { 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); + past_data, present_data, mask_index_type, input_hidden_size); 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); + past_data, present_data, mask_index_type, input_hidden_size); } TEST(AttentionTest, AttentionBatch1) { @@ -1429,5 +1433,51 @@ TEST(AttentionTest, AttentionPastState_dynamic) { test.Run(); } + +TEST(AttentionTest, AttentionPrunedModel) { + int batch_size = 2; + int sequence_length = 2; + // test input_hidden_size > hidden_size + int input_hidden_size = 6; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.5f, 0.2f, 0.3f, -0.6f, 0.0f, 1.0f, + 0.8f, -0.5f, 0.0f, 1.f, 2.0f, 3.0f, + 0.8f, -0.5f, 0.0f, 1.f, 4.0f, 5.0f, + 0.5f, 0.2f, 0.3f, -0.6f, 6.0f, 7.0f}; + + std::vector 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, + 0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 1.0f, 1.1f, 1.2f, + 1.2f, 1.1f, 1.0f, 0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f, 0.1f, + }; + + std::vector 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}; + + std::vector mask_data = {1, 1, 1, 1}; + + std::vector output_data = { + 11.689527f, 2.769937f, 7.05f, 8.350000f, + 11.690000f, 2.770000f, 7.05f, 8.350000f, + 14.276558f, 5.374159f, 9.650001f, 10.95f, + 14.289073f, 5.370287f, 9.650001f, 10.95f + }; + + bool use_float16 = false; + bool is_unidirectional = false; + bool use_past_state = false; + int past_sequence_length = 0; + const std::vector* past_data = nullptr; + const std::vector* present_data = nullptr; + RunAttentionTest(input_data, weight_data, bias_data, mask_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, input_hidden_size); +} } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc b/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc index 4f4d522fb1..6b273bd87b 100644 --- a/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc @@ -85,18 +85,21 @@ void RunQAttention(const std::vector& input_data, // input: int hidden_size, int number_of_heads, bool is_unidirectional = false, - bool use_float16 = false) { + bool use_float16 = false, + int input_hidden_size = 0) { + input_hidden_size = (input_hidden_size == 0) ? hidden_size : input_hidden_size; + OpTester tester("QAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(number_of_heads)); if (is_unidirectional) { tester.AddAttribute("unidirectional", 1); } - std::vector input_dims = {batch_size, sequence_length, hidden_size}; - std::vector weights_dims = {hidden_size, 3 * hidden_size}; + std::vector input_dims = {batch_size, sequence_length, input_hidden_size}; + std::vector weights_dims = {input_hidden_size, 3 * hidden_size}; std::vector bias_dims = {3 * hidden_size}; std::vector mask_index_dims = {batch_size}; - std::vector output_dims = input_dims; + std::vector output_dims = {batch_size, sequence_length, hidden_size}; float input_scale = quantize_parameters.input_scale; float weight_scale = quantize_parameters.weight_scale; @@ -121,10 +124,11 @@ void RunQAttention(const std::vector& input_data, // input: tester.AddOutput("output", output_dims, output_data); } - if (mask_index_data.size() > 0) { // mask index is optional. + if (mask_index_data.size() > 0) { tester.AddInput("mask_index", mask_index_dims, mask_index_data); - } else { // mask index is optional. - tester.AddInput("", mask_index_dims, {1}); + } else { + // mask index is optional. + tester.AddMissingOptionalInput(); } tester.AddInput("input_zero_point", {1}, {input_zero_point}); @@ -153,7 +157,8 @@ static void RunQAttentionCUDA( int number_of_heads, bool use_special_quantize_parameter = true, bool is_unidirectional = false, - bool use_float16 = false) { + bool use_float16 = false, + int input_hidden_size = 0) { int min_cuda_architecture = 530; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); @@ -165,7 +170,7 @@ static void RunQAttentionCUDA( } RunQAttention( input_data, weights_data, bias_data, mask_index_data, output_data, qp, - batch_size, sequence_length, hidden_size, number_of_heads, is_unidirectional, use_float16); + batch_size, sequence_length, hidden_size, number_of_heads, is_unidirectional, use_float16, input_hidden_size); } } @@ -180,7 +185,8 @@ static void RunQAttentionU8U8( int hidden_size, int number_of_heads, bool use_special_quantize_parameter = true, - bool is_unidirectional = false) { + bool is_unidirectional = false, + int input_hidden_size = 0) { QuantizeParameters qp_uint8{0.0f, 0.0f, 0, 0}; if (use_special_quantize_parameter) { qp_uint8.input_scale = 0.1f; @@ -191,7 +197,7 @@ static void RunQAttentionU8U8( RunQAttention( input_data, weights_data, bias_data, mask_index_data, output_data, qp_uint8, - batch_size, sequence_length, hidden_size, number_of_heads, is_unidirectional); + batch_size, sequence_length, hidden_size, number_of_heads, is_unidirectional, false, input_hidden_size); } static void RunQAttentionU8S8( @@ -205,7 +211,8 @@ static void RunQAttentionU8S8( int hidden_size, int number_of_heads, bool use_special_quantize_parameter = true, - bool is_unidirectional = false) { + bool is_unidirectional = false, + int input_hidden_size = 0) { QuantizeParameters qp_int8{0.0f, 0.0f, 0, 0}; if (use_special_quantize_parameter) { qp_int8.input_scale = 0.1f; @@ -216,7 +223,7 @@ static void RunQAttentionU8S8( RunQAttention( input_data, weights_data, bias_data, mask_index_data, output_data, qp_int8, - batch_size, sequence_length, hidden_size, number_of_heads, is_unidirectional); + batch_size, sequence_length, hidden_size, number_of_heads, is_unidirectional, false, input_hidden_size); } static void RunQAttentionAll( @@ -231,16 +238,17 @@ static void RunQAttentionAll( int number_of_heads, bool use_special_quantize_parameter = true, bool is_unidirectional = false, - bool use_float16 = false) { + bool use_float16 = false, + int input_hidden_size = 0) { RunQAttentionU8U8(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, - use_special_quantize_parameter, is_unidirectional); + use_special_quantize_parameter, is_unidirectional, input_hidden_size); RunQAttentionU8S8(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, - use_special_quantize_parameter, is_unidirectional); + use_special_quantize_parameter, is_unidirectional, input_hidden_size); RunQAttentionCUDA(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, - use_special_quantize_parameter, is_unidirectional, use_float16); + use_special_quantize_parameter, is_unidirectional, use_float16, input_hidden_size); } TEST(QAttentionTest, QAttentionBatch1) { @@ -781,5 +789,49 @@ TEST(QAttentionTest, QAttentionPastState_u8s8) { true /*is_weight_constant*/); } +TEST(QAttentionTest, QAttentionPrunedModel) { + int batch_size = 2; + int sequence_length = 2; + // test input_hidden_size > hidden_size + int input_hidden_size = 6; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.5f, 0.2f, 0.3f, -0.6f, 0.0f, 1.0f, + 0.8f, -0.5f, 0.0f, 1.f, 2.0f, 3.0f, + 0.8f, -0.5f, 0.0f, 1.f, 4.0f, 5.0f, + 0.5f, 0.2f, 0.3f, -0.6f, 6.0f, 7.0f}; + + std::vector 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, + 0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 1.0f, 1.1f, 1.2f, + 1.2f, 1.1f, 1.0f, 0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f, 0.1f, + }; + + std::vector 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}; + + std::vector mask_data = {}; + + std::vector output_data = { + 11.689527f, 2.769937f, 7.05f, 8.35f, + 11.69f, 2.77f, 7.05f, 8.35f, + 14.276558f, 5.374159f, 9.65f, 10.95f, + 14.289073f, 5.370287f, 9.65f, 10.95f + }; + + bool use_special_quantize_parameter = true; + bool is_unidirectional = false; + bool use_float16 = false; + RunQAttentionAll(input_data, weight_data, bias_data, mask_data, output_data, + batch_size, sequence_length, hidden_size, number_of_heads, + use_special_quantize_parameter, is_unidirectional, use_float16, + input_hidden_size); +} + } // namespace test } // namespace onnxruntime