Add tests of Attention and QAttention for pruned model (#6914)

This commit is contained in:
Tianlei Wu 2021-03-08 11:56:31 -08:00 committed by GitHub
parent f2f60eed59
commit b89f52c277
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 127 additions and 25 deletions

View file

@ -34,19 +34,22 @@ static void RunAttentionTest(
int past_sequence_length = 0,
const std::vector<float>* past_data = nullptr,
const std::vector<float>* 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<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
tester.AddAttribute<int64_t>("unidirectional", static_cast<int64_t>(is_unidirectional ? 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};
std::vector<int64_t> input_dims = {batch_size, sequence_length, input_hidden_size};
std::vector<int64_t> weights_dims = {input_hidden_size, 3 * hidden_size};
std::vector<int64_t> bias_dims = {3 * hidden_size};
std::vector<int64_t> mask_index_dims_1 = {batch_size};
@ -78,7 +81,7 @@ static void RunAttentionTest(
std::vector<int64_t> past_dims = {2, batch_size, number_of_heads, past_sequence_length, head_size};
std::vector<int64_t> present_dims = {2, batch_size, number_of_heads, past_sequence_length + sequence_length, head_size};
std::vector<int64_t> output_dims = input_dims;
std::vector<int64_t> output_dims = {batch_size, sequence_length, hidden_size};
if (use_float16) {
tester.AddInput<MLFloat16>("input", input_dims, ToFloat16(input_data));
@ -142,15 +145,16 @@ static void RunAttentionTest(
int past_sequence_length = 0,
const std::vector<float>* past_data = nullptr,
const std::vector<float>* 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<float> 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<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,
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<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};
std::vector<int32_t> mask_data = {1, 1, 1, 1};
std::vector<float> 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<float>* past_data = nullptr;
const std::vector<float>* 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

View file

@ -85,18 +85,21 @@ void RunQAttention(const std::vector<float>& 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<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
if (is_unidirectional) {
tester.AddAttribute<int64_t>("unidirectional", 1);
}
std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
std::vector<int64_t> weights_dims = {hidden_size, 3 * hidden_size};
std::vector<int64_t> input_dims = {batch_size, sequence_length, input_hidden_size};
std::vector<int64_t> weights_dims = {input_hidden_size, 3 * hidden_size};
std::vector<int64_t> bias_dims = {3 * hidden_size};
std::vector<int64_t> mask_index_dims = {batch_size};
std::vector<int64_t> output_dims = input_dims;
std::vector<int64_t> 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<float>& input_data, // input:
tester.AddOutput<float>("output", output_dims, output_data);
}
if (mask_index_data.size() > 0) { // mask index is optional.
if (mask_index_data.size() > 0) {
tester.AddInput<int32_t>("mask_index", mask_index_dims, mask_index_data);
} else { // mask index is optional.
tester.AddInput<int32_t>("", mask_index_dims, {1});
} else {
// mask index is optional.
tester.AddMissingOptionalInput<int32_t>();
}
tester.AddInput<QInput>("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<int8_t, int8_t, EP::CUDA>(
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<uint8_t, uint8_t> 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<uint8_t, uint8_t, EP::CPU>(
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<uint8_t, int8_t> 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<uint8_t, int8_t, EP::CPU>(
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<float> 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<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,
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<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};
std::vector<int32_t> mask_data = {};
std::vector<float> 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