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

Handle dummy mask in Attention operators (#5108)

Browse source 
* Handle dummy mask with shape (1, 1) or (batch_size, 1).
This commit is contained in:
Tianlei Wu 2020-09-11 09:31:03 -07:00 committed by GitHub
parent c794c88ae0
commit ccfbc56388
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 79 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -33,13 +33,13 @@ AttentionBase::AttentionBase(const OpKernelInfo& info) {
Status AttentionBase::CheckInputs(const TensorShape& input_shape,
const TensorShape& weights_shape,
const TensorShape& bias_shape,
const Tensor* mask_index,
const Tensor*& mask_index,
const Tensor* past) const {
// Input shapes:
// input : (batch_size, sequence_length, hidden_size)
// weights : (hidden_size, 3 * hidden_size)
// bias : (3 * hidden_size)
// mask_index : nullptr, (batch_size), (2 * batch_size), or (batch_size, past_sequence_length + sequence_length)
// mask_index : nullptr, (batch_size), (2 * batch_size), (batch_size, 1), (1, 1) or (batch_size, past_sequence_length + sequence_length)
// past : (2, batch_size, num_heads, past_sequence_length, head_size)
const auto& dims = input_shape.GetDims();
@ -112,7 +112,14 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
}
} else if (mask_dims.size() == 2) {
if (static_cast<int>(mask_dims[0]) != batch_size || static_cast<int>(mask_dims[1]) != past_sequence_length + sequence_length) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'mask_index' with raw attention mask shall have shape batch_size x (past_sequence_length + sequence_length)");
// Add operator supports broadcasting. Here we handle a case with only one element in the 2nd dimension.
if ((static_cast<int>(mask_dims[0]) == batch_size || static_cast<int>(mask_dims[0]) == 1) && static_cast<int>(mask_dims[1]) == 1) {
// Mask will have same value after propogation, which has same effect as no mask.
mask_index = nullptr;
}
else {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'mask_index' with raw attention mask shall have shape batch_size x (past_sequence_length + sequence_length)");
}
}
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'mask_index' is expected to have 1 or 2 dimensions, got ",

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

@ -12,10 +12,11 @@ namespace contrib {
class AttentionBase {
protected:
AttentionBase(const OpKernelInfo& info);
Status CheckInputs(const TensorShape& input_shape,
const TensorShape& weights_shape,
const TensorShape& bias_shape,
const Tensor* mask_index,
const Tensor*& mask_index, // For dummy mask with shape (1, 1) or (batch_size, 1), it will be updated to nullptr.
const Tensor* past) const;
Tensor* GetPresent(OpKernelContext* context,

2
onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc
View file

@ -107,7 +107,7 @@ Status QAttention<T>::Compute(OpKernelContext* context) const {
// Input 2 - bias : (3 * hidden_size)
// Input 3 - input_scale : scalar
// Input 4 - weight_scale : scalar
// Input 5 - mask_index : (batch_size)
// Input 5 - mask_index : nullptr, (batch_size), (2 * batch_size), (batch_size, 1), (1, 1) or (batch_size, past_sequence_length + sequence_length)
// Input 6 - input_zero_point : scalar
// Input 7 - weight_zero_point : scalar
// Input 8 - past : (2, batch_size, num_heads, past_sequence_length, head_size)

4
onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc
View file

@ -45,7 +45,7 @@ Status QAttention<T, int8_t>::CheckInputs(const Tensor* input,
const Tensor* bias,
const Tensor* input_scale_tensor,
const Tensor* weight_scale_tensor,
const Tensor* mask_index,
const Tensor*& mask_index,
const Tensor* i_zp_tensor,
const Tensor* w_zp_tensor,
const Tensor* past_tensor) const {
@ -55,7 +55,7 @@ Status QAttention<T, int8_t>::CheckInputs(const Tensor* input,
// Input 2 - bias : (3 * hidden_size)
// Input 3 - input_scale : scalar
// Input 4 - weight_scale : scalar
// Input 5 - mask_index : (batch_size)
// Input 5 - mask_index : nullptr, (batch_size), (2 * batch_size), (batch_size, 1), (1, 1) or (batch_size, past_sequence_length + sequence_length)
// Input 6 - input_zero_point : scalar
// Input 7 - weight_zero_point : scalar
// Output : (batch_size, sequence_length, hidden_size)

2
onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h
View file

@ -34,7 +34,7 @@ class QAttention<T, int8_t> final : public CudaKernel, public AttentionBase {
const Tensor* bias,
const Tensor* input_scale_tensor,
const Tensor* weight_scale_tensor,
const Tensor* mask_index,
const Tensor*& mask_index,
const Tensor* i_zp_tensor,
const Tensor* w_zp_tensor,
const Tensor* past_tensor) const;

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

@ -12,7 +12,8 @@ namespace test {
enum MaskIndexType {
kMaskIndexEnd = 0,
kMaskIndexEndAndStart,
kMaskRaw
kMaskRaw,
kMaskDummy // Dummy mask with shape [1, 1] or [batch_size, 1]
};
static void RunAttentionTest(
@ -49,7 +50,25 @@ static void RunAttentionTest(
std::vector<int64_t> mask_index_dims_1 = {batch_size};
std::vector<int64_t> mask_index_dims_2 = {2 * batch_size};
std::vector<int64_t> mask_index_dims_3 = {batch_size, past_sequence_length + sequence_length};
std::vector<int64_t> mask_index_dims = (mask_index_type == kMaskIndexEnd ? mask_index_dims_1 : (mask_index_type == kMaskIndexEndAndStart ? mask_index_dims_2 : mask_index_dims_3));
std::vector<int64_t> mask_index_dims_4 = {batch_size, 1};
std::vector<int64_t> mask_index_dims;
switch (mask_index_type) {
case kMaskIndexEnd:
mask_index_dims = mask_index_dims_1;
break;
case kMaskIndexEndAndStart:
mask_index_dims = mask_index_dims_2;
break;
case kMaskRaw:
mask_index_dims = mask_index_dims_3;
break;
case kMaskDummy:
mask_index_dims = mask_index_dims_4;
break;
default:
assert(0); // shall not reach here.
break;
}
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};
@ -1134,6 +1153,48 @@ TEST(AttentionTest, AttentionMask2DNoWord) {
use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskRaw);
}
TEST(AttentionTest, AttentionDummyMask2D) {
int batch_size = 2;
int sequence_length = 2;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
0.5f, 0.2f, 0.3f, -0.6f,
0.8f, -0.5f, 0.0f, 1.f,
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};
std::vector<int32_t> mask_index_data = {1, 1};
std::vector<float> output_data = {
3.9696791172027588f, 0.073143675923347473f, 4.25f, 5.65f,
3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.65f,
3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.65f,
3.9696791172027588f, 0.073143675923347473f, 4.25f, 5.65f};
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_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);
}
TEST(AttentionTest, AttentionMaskIndexOutOfRange) {
int batch_size = 2;
int sequence_length = 2;