Почему входной размер модуля MultiheadAttention в модуле Pytorch Transformer 1536? - PullRequest
Новая
       92

Почему входной размер модуля MultiheadAttention в модуле Pytorch Transformer 1536?

1 голос
/ 24 октября 2019

При использовании модуля / объекта torch.nn.modules.transformer.Transformer первый слой - это слой encoder.layers.0.self_attn, который является слоем MultiheadAttention, т.е. 

from torch.nn.modules.transformer import Transformer
bumblebee = Transformer()

bumblee.parameters

[out]: 

<bound method Module.parameters of Transformer(
  (encoder): TransformerEncoder(
    (layers): ModuleList(
      (0): TransformerEncoderLayer(
        (self_attn): MultiheadAttention(
          (out_proj): Linear(in_features=512, out_features=512, bias=True)
        )
        (linear1): Linear(in_features=512, out_features=2048, bias=True)
        (dropout): Dropout(p=0.1, inplace=False)
        (linear2): Linear(in_features=2048, out_features=512, bias=True)
        (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
        (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
        (dropout1): Dropout(p=0.1, inplace=False)
        (dropout2): Dropout(p=0.1, inplace=False)
      )

И если мы распечатываем размер слоя, мы видим: 

for name in bumblebee.encoder.state_dict():
    print(name, '\t', bumblebee.encoder.state_dict()[name].shape)

[out]: 

layers.0.self_attn.in_proj_weight    torch.Size([1536, 512])
layers.0.self_attn.in_proj_bias      torch.Size([1536])
layers.0.self_attn.out_proj.weight   torch.Size([512, 512])
layers.0.self_attn.out_proj.bias     torch.Size([512])
layers.0.linear1.weight      torch.Size([2048, 512])
layers.0.linear1.bias    torch.Size([2048])
layers.0.linear2.weight      torch.Size([512, 2048])
layers.0.linear2.bias    torch.Size([512])
layers.0.norm1.weight    torch.Size([512])
layers.0.norm1.bias      torch.Size([512])
layers.0.norm2.weight    torch.Size([512])
layers.0.norm2.bias      torch.Size([512])

Кажется, что 1536 - это 512 * 3 и как-тоПараметр layers.0.self_attn.in_proj_weight может хранить все три тензора QKV в архитектуре трансформатора в одной матрице.

С https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/activation.py#L649 

class MultiheadAttention(Module):
    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
        super(MultiheadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.kdim = kdim if kdim is not None else embed_dim
        self.vdim = vdim if vdim is not None else embed_dim
        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim

        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"

        if self._qkv_same_embed_dim is False:
            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
        else:
            self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))

И примечание в строке документа MultiheadAttention гласит:

Примечание: if kdim и vdimNone, они будут установлены в embed_dim так, чтобы запрос, ключ и значение имели одинаковое количество функций.

Это правильно?

1 Ответ

1 голос
/ 25 октября 2019

Из определения nn.Transformer со значениями по умолчанию, EncoderLayer создается с d_model=512, nhead=8.

MultiheadAttention создается с d_model, nhead, равным этим значениям, а k_dim, v_dim оставляются по умолчанию None.

Если они равны None, self._qkv_same_embed_dim в , эта строка оценивается как True. Когда это происходит, как вы правильно указали, self.in_proj_weight определяется как тензор формы (3 x embed_dim, embed_dim).

Короче говоря: да, это правильно.

...