比较与torch.nn.MultiheadAttention的功能差异

查看源文件

torch.nn.MultiheadAttention

torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, batch_first=False, device=None, dtype=None)

更多内容详见torch.nn.MultiheadAttention

mindspore.nn.transformer.MultiHeadAttention

class mindspore.nn.transformer.MultiHeadAttention(batch_size, src_seq_length, tgt_seq_length, hidden_size, num_heads, hidden_dropout_rate=0.1, attention_dropout_rate=0.1, compute_dtype=mstype.float16, softmax_compute_type=mstype.float32, param_init_type=mstype.float32, use_past=False, parallel_config=default_dpmp_config)(
    query_tensor, key_tensor, value_tensor, attention_mask, key_past=None,
    value_past=None, batch_valid_length=None
)

更多内容详见mindspore.nn.transformer.MultiHeadAttention

使用方式

mindspore.nn.transformer.MultiHeadAttention在初始化参数和torch.nn.MultiheadAttention并不完全相同,但是基本功能保持一致。具体的区别如下说明:

mindspore.nn.transformer.MultiHeadAttention

torch.nn.MultiheadAttention

说明

batch_size

MindSpore需要传入额外的batch size以作校验和增量推理使用。

src_seq_length

encoder输入序列长度。

tgt_seq_length

decoder输入序列长度。

hidden_size

embed_dim

参数名称不一致,含义相同。

num_heads

num_heads

Attention的head数目,含义相同。

hidden_dropout_rate

dropout

含义不同。hidden_dropout_rate表示在隐藏层处的dropout,而PyTorch的dropout参数额外控制了softmax处的dropout rate。

attention_dropout_rate

dropout

含义不同。attention_dropout_rate表示在softmax处的dropout,而PyTorch的dropout参数额外控制了隐藏层的dropout rate。

compute_type

控制内部matmul矩阵计算类型。

softmax_compute_type

控制attention中softmax的计算类型。

param_init_type

控制参数初始化的类型。

use_past

是否使用增量推理。

parallel_config

并行设置的配置参数。

bias

是否在projection层添加bias。MindSpore默认行为是添加的。

add_bias_kv

是否在第0维度的key和value序列上面添加bias。MindSpore未实现此功能。

add_zero_attn

是否在第1维度的key和value序列上面添加全零的数据。MindSpore未实现此功能。

kdims

key维度的feature数量。MindSpore未实现此功能。

vdims

value维度的feature数量。MindSpore未实现此功能。

batch_first

MindSpore默认配置为(batch,seq, feature),即Pytorch的batch_first=True

  • mindspore.nn.transformer.MultiHeadAttention缺少bias、add_bias_kv、add_zero_attn、kdim和vdim的输入,在前向计算中缺少key_padding_mask的输入。

  • mindspore.nn.transformer.MultiHeadAttention的dropout_rate分别为0.1,而pytorch对应的版本为0.0。

  • mindspore.nn.transformer.MultiHeadAttention的输入中attention_mask是必须的。

  • mindspore.nn.transformer.MultiHeadAttention提供了并行配置parallel_config入参,可以实现混合并行。

  • mindspore.nn.transformer.MultiHeadAttention会返回attention的key和value的历史值,而torch.nn.MultiheadAttention可以控制是否返回attention中计算的key和query之间的得分。

  • mindspore.nn.transformer.MultiHeadAttention提供了静态图的增量推理功能。

  • mindspore.nn.transformer.MultiHeadAttention默认采用fp16进行矩阵运算。

  • mindspore.nn.transformer.MultiHeadAttention的返回tensor的第0维是batch维度,而pytorch默认是第0维度为seq_length。

PyTorch:实例化Transformer时需要提供的参数较少。

MindSpore:在类初始化的时候,需要提供batch_size、源序列和目标序列长度等额外信息,并且在计算时需要输入attention_mask。

代码示例

import numpy as np
from mindspore.nn.transformer import MultiHeadAttention
import mindspore as ms
model = MultiHeadAttention(batch_size=32, hidden_size=512, src_seq_length=10, tgt_seq_length=20,
                           num_heads=8)
query = ms.Tensor(np.random.rand(32, 10, 512), ms.float32)
key = ms.Tensor(np.random.rand(32, 20, 512), ms.float32)
value = ms.Tensor(np.random.rand(32, 20, 512), ms.float32)
attention_mask = ms.Tensor(np.ones((32, 10, 20)), ms.float16)
output, past = model(query, key, value, attention_mask)
print(output.shape)
# output:
# (32, 20, 512)

import torch
query = torch.rand(20, 32, 512)
key = torch.rand(10, 32, 512)
value = torch.rand(10, 32, 512)
multihead_attn = torch.nn.MultiheadAttention(embed_dim=512, num_heads=8)
output, attn_output_weights = multihead_attn(query, key, value)
print(output.shape)
# output:
# torch.Size([20, 32, 512])