Spaces:
Runtime error
Runtime error
| # Code copy from PyTorch, modified by Xueyan Zou | |
| import warnings | |
| from typing import Optional, Tuple | |
| import torch | |
| import torch.nn as nn | |
| from torch import Tensor | |
| from torch.nn.init import constant_, xavier_normal_, xavier_uniform_ | |
| from torch.nn.parameter import Parameter | |
| from torch.overrides import has_torch_function, handle_torch_function | |
| from torch.nn.functional import pad, linear, softmax, dropout | |
| def multi_head_attention_forward( | |
| query: Tensor, | |
| key: Tensor, | |
| value: Tensor, | |
| embed_dim_to_check: int, | |
| num_heads: int, | |
| in_proj_weight: Tensor, | |
| in_proj_bias: Tensor, | |
| bias_k: Optional[Tensor], | |
| bias_v: Optional[Tensor], | |
| add_zero_attn: bool, | |
| dropout_p: float, | |
| out_proj_weight: Tensor, | |
| out_proj_bias: Tensor, | |
| training: bool = True, | |
| key_padding_mask: Optional[Tensor] = None, | |
| need_weights: bool = True, | |
| attn_mask: Optional[Tensor] = None, | |
| use_separate_proj_weight: bool = False, | |
| q_proj_weight: Optional[Tensor] = None, | |
| k_proj_weight: Optional[Tensor] = None, | |
| v_proj_weight: Optional[Tensor] = None, | |
| static_k: Optional[Tensor] = None, | |
| static_v: Optional[Tensor] = None, | |
| ) -> Tuple[Tensor, Optional[Tensor]]: | |
| r""" | |
| Args: | |
| query, key, value: map a query and a set of key-value pairs to an output. | |
| See "Attention Is All You Need" for more details. | |
| embed_dim_to_check: total dimension of the model. | |
| num_heads: parallel attention heads. | |
| in_proj_weight, in_proj_bias: input projection weight and bias. | |
| bias_k, bias_v: bias of the key and value sequences to be added at dim=0. | |
| add_zero_attn: add a new batch of zeros to the key and | |
| value sequences at dim=1. | |
| dropout_p: probability of an element to be zeroed. | |
| out_proj_weight, out_proj_bias: the output projection weight and bias. | |
| training: apply dropout if is ``True``. | |
| key_padding_mask: if provided, specified padding elements in the key will | |
| be ignored by the attention. This is an binary mask. When the value is True, | |
| the corresponding value on the attention layer will be filled with -inf. | |
| need_weights: output attn_output_weights. | |
| attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all | |
| the batches while a 3D mask allows to specify a different mask for the entries of each batch. | |
| use_separate_proj_weight: the function accept the proj. weights for query, key, | |
| and value in different forms. If false, in_proj_weight will be used, which is | |
| a combination of q_proj_weight, k_proj_weight, v_proj_weight. | |
| q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias. | |
| static_k, static_v: static key and value used for attention operators. | |
| Shape: | |
| Inputs: | |
| - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is | |
| the embedding dimension. | |
| - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is | |
| the embedding dimension. | |
| - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is | |
| the embedding dimension. | |
| - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length. | |
| If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions | |
| will be unchanged. If a BoolTensor is provided, the positions with the | |
| value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged. | |
| - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length. | |
| 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length, | |
| S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked | |
| positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend | |
| while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True`` | |
| are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor | |
| is provided, it will be added to the attention weight. | |
| - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length, | |
| N is the batch size, E is the embedding dimension. E/num_heads is the head dimension. | |
| - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length, | |
| N is the batch size, E is the embedding dimension. E/num_heads is the head dimension. | |
| Outputs: | |
| - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, | |
| E is the embedding dimension. | |
| - attn_output_weights: :math:`(N, L, S)` where N is the batch size, | |
| L is the target sequence length, S is the source sequence length. | |
| """ | |
| tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v, out_proj_weight, out_proj_bias) | |
| if has_torch_function(tens_ops): | |
| return handle_torch_function( | |
| multi_head_attention_forward, | |
| tens_ops, | |
| query, | |
| key, | |
| value, | |
| embed_dim_to_check, | |
| num_heads, | |
| in_proj_weight, | |
| in_proj_bias, | |
| bias_k, | |
| bias_v, | |
| add_zero_attn, | |
| dropout_p, | |
| out_proj_weight, | |
| out_proj_bias, | |
| training=training, | |
| key_padding_mask=key_padding_mask, | |
| need_weights=need_weights, | |
| attn_mask=attn_mask, | |
| use_separate_proj_weight=use_separate_proj_weight, | |
| q_proj_weight=q_proj_weight, | |
| k_proj_weight=k_proj_weight, | |
| v_proj_weight=v_proj_weight, | |
| static_k=static_k, | |
| static_v=static_v, | |
| ) | |
| tgt_len, bsz, embed_dim = query.size() | |
| assert embed_dim == embed_dim_to_check | |
| # allow MHA to have different sizes for the feature dimension | |
| assert key.size(0) == value.size(0) and key.size(1) == value.size(1) | |
| head_dim = embed_dim // num_heads | |
| assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads" | |
| scaling = float(head_dim) ** -0.5 | |
| if not use_separate_proj_weight: | |
| if (query is key or torch.equal(query, key)) and (key is value or torch.equal(key, value)): | |
| # self-attention | |
| q, k, v = linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1) | |
| elif key is value or torch.equal(key, value): | |
| # encoder-decoder attention | |
| # This is inline in_proj function with in_proj_weight and in_proj_bias | |
| _b = in_proj_bias | |
| _start = 0 | |
| _end = embed_dim | |
| _w = in_proj_weight[_start:_end, :] | |
| if _b is not None: | |
| _b = _b[_start:_end] | |
| q = linear(query, _w, _b) | |
| if key is None: | |
| assert value is None | |
| k = None | |
| v = None | |
| else: | |
| # This is inline in_proj function with in_proj_weight and in_proj_bias | |
| _b = in_proj_bias | |
| _start = embed_dim | |
| _end = None | |
| _w = in_proj_weight[_start:, :] | |
| if _b is not None: | |
| _b = _b[_start:] | |
| k, v = linear(key, _w, _b).chunk(2, dim=-1) | |
| else: | |
| # This is inline in_proj function with in_proj_weight and in_proj_bias | |
| _b = in_proj_bias | |
| _start = 0 | |
| _end = embed_dim | |
| _w = in_proj_weight[_start:_end, :] | |
| if _b is not None: | |
| _b = _b[_start:_end] | |
| q = linear(query, _w, _b) | |
| # This is inline in_proj function with in_proj_weight and in_proj_bias | |
| _b = in_proj_bias | |
| _start = embed_dim | |
| _end = embed_dim * 2 | |
| _w = in_proj_weight[_start:_end, :] | |
| if _b is not None: | |
| _b = _b[_start:_end] | |
| k = linear(key, _w, _b) | |
| # This is inline in_proj function with in_proj_weight and in_proj_bias | |
| _b = in_proj_bias | |
| _start = embed_dim * 2 | |
| _end = None | |
| _w = in_proj_weight[_start:, :] | |
| if _b is not None: | |
| _b = _b[_start:] | |
| v = linear(value, _w, _b) | |
| else: | |
| q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight) | |
| len1, len2 = q_proj_weight_non_opt.size() | |
| assert len1 == embed_dim and len2 == query.size(-1) | |
| k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight) | |
| len1, len2 = k_proj_weight_non_opt.size() | |
| assert len1 == embed_dim and len2 == key.size(-1) | |
| v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight) | |
| len1, len2 = v_proj_weight_non_opt.size() | |
| assert len1 == embed_dim and len2 == value.size(-1) | |
| if in_proj_bias is not None: | |
| q = linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim]) | |
| k = linear(key, k_proj_weight_non_opt, in_proj_bias[embed_dim : (embed_dim * 2)]) | |
| v = linear(value, v_proj_weight_non_opt, in_proj_bias[(embed_dim * 2) :]) | |
| else: | |
| q = linear(query, q_proj_weight_non_opt, in_proj_bias) | |
| k = linear(key, k_proj_weight_non_opt, in_proj_bias) | |
| v = linear(value, v_proj_weight_non_opt, in_proj_bias) | |
| q = q * scaling | |
| if attn_mask is not None: | |
| assert ( | |
| attn_mask.dtype == torch.float32 | |
| or attn_mask.dtype == torch.float64 | |
| or attn_mask.dtype == torch.float16 | |
| or attn_mask.dtype == torch.uint8 | |
| or attn_mask.dtype == torch.bool | |
| ), "Only float, byte, and bool types are supported for attn_mask, not {}".format(attn_mask.dtype) | |
| if attn_mask.dtype == torch.uint8: | |
| warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.") | |
| attn_mask = attn_mask.to(torch.bool) | |
| if attn_mask.dim() == 2: | |
| attn_mask = attn_mask.unsqueeze(0) | |
| if list(attn_mask.size()) != [1, query.size(0), key.size(0)]: | |
| raise RuntimeError("The size of the 2D attn_mask is not correct.") | |
| elif attn_mask.dim() == 3: | |
| if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]: | |
| raise RuntimeError("The size of the 3D attn_mask is not correct.") | |
| else: | |
| raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim())) | |
| # attn_mask's dim is 3 now. | |
| # convert ByteTensor key_padding_mask to bool | |
| if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8: | |
| warnings.warn( | |
| "Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead." | |
| ) | |
| key_padding_mask = key_padding_mask.to(torch.bool) | |
| if bias_k is not None and bias_v is not None: | |
| if static_k is None and static_v is None: | |
| k = torch.cat([k, bias_k.repeat(1, bsz, 1)]) | |
| v = torch.cat([v, bias_v.repeat(1, bsz, 1)]) | |
| if attn_mask is not None: | |
| attn_mask = pad(attn_mask, (0, 1)) | |
| if key_padding_mask is not None: | |
| key_padding_mask = pad(key_padding_mask, (0, 1)) | |
| else: | |
| assert static_k is None, "bias cannot be added to static key." | |
| assert static_v is None, "bias cannot be added to static value." | |
| else: | |
| assert bias_k is None | |
| assert bias_v is None | |
| q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1) | |
| if k is not None: | |
| k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1) | |
| if v is not None: | |
| v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1) | |
| if static_k is not None: | |
| assert static_k.size(0) == bsz * num_heads | |
| assert static_k.size(2) == head_dim | |
| k = static_k | |
| if static_v is not None: | |
| assert static_v.size(0) == bsz * num_heads | |
| assert static_v.size(2) == head_dim | |
| v = static_v | |
| src_len = k.size(1) | |
| if key_padding_mask is not None: | |
| # assert key_padding_mask.size(0) == bsz | |
| assert key_padding_mask.size(1) == src_len | |
| if add_zero_attn: | |
| src_len += 1 | |
| k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1) | |
| v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1) | |
| if attn_mask is not None: | |
| attn_mask = pad(attn_mask, (0, 1)) | |
| if key_padding_mask is not None: | |
| key_padding_mask = pad(key_padding_mask, (0, 1)) | |
| attn_output_weights = torch.bmm(q, k.transpose(1, 2)) | |
| assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len] | |
| if attn_mask is not None: | |
| if attn_mask.dtype == torch.bool: | |
| attn_output_weights.masked_fill_(attn_mask, float("-inf")) | |
| else: | |
| attn_output_weights += attn_mask | |
| if key_padding_mask is not None: | |
| attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len) | |
| attn_output_weights = attn_output_weights.masked_fill( | |
| key_padding_mask.unsqueeze(1), | |
| float("-inf"), | |
| ) | |
| attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len) | |
| attn_output_weights = softmax(attn_output_weights, dim=-1) | |
| attn_output_weights = dropout(attn_output_weights, p=dropout_p, training=training) | |
| attn_output = torch.bmm(attn_output_weights, v) | |
| assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim] | |
| attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim) | |
| attn_output = linear(attn_output, out_proj_weight, out_proj_bias) | |
| if need_weights: | |
| # average attention weights over heads | |
| attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len) | |
| return attn_output, attn_output_weights.sum(dim=1) / num_heads | |
| else: | |
| return attn_output, None | |
| # This class exists solely for Transformer; it has an annotation stating | |
| # that bias is never None, which appeases TorchScript | |
| class _LinearWithBias(nn.Linear): | |
| bias: Tensor # type: ignore | |
| def __init__(self, in_features: int, out_features: int) -> None: | |
| super().__init__(in_features, out_features, bias=True) # type: ignore | |
| class MultiheadAttention(nn.Module): | |
| r"""Allows the model to jointly attend to information | |
| from different representation subspaces. | |
| See `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_ | |
| .. math:: | |
| \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O | |
| where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`. | |
| Args: | |
| embed_dim: total dimension of the model. | |
| num_heads: parallel attention heads. | |
| dropout: a Dropout layer on attn_output_weights. Default: 0.0. | |
| bias: add bias as module parameter. Default: True. | |
| add_bias_kv: add bias to the key and value sequences at dim=0. | |
| add_zero_attn: add a new batch of zeros to the key and | |
| value sequences at dim=1. | |
| kdim: total number of features in key. Default: None. | |
| vdim: total number of features in value. Default: None. | |
| Note that if :attr:`kdim` and :attr:`vdim` are None, they will be set | |
| to :attr:`embed_dim` such that query, key, and value have the same | |
| number of features. | |
| Examples:: | |
| >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads) | |
| >>> attn_output, attn_output_weights = multihead_attn(query, key, value) | |
| """ | |
| bias_k: Optional[torch.Tensor] | |
| bias_v: Optional[torch.Tensor] | |
| 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)) | |
| self.register_parameter('in_proj_weight', None) | |
| else: | |
| self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim)) | |
| self.register_parameter('q_proj_weight', None) | |
| self.register_parameter('k_proj_weight', None) | |
| self.register_parameter('v_proj_weight', None) | |
| if bias: | |
| self.in_proj_bias = Parameter(torch.empty(3 * embed_dim)) | |
| else: | |
| self.register_parameter('in_proj_bias', None) | |
| self.out_proj = _LinearWithBias(embed_dim, embed_dim) | |
| if add_bias_kv: | |
| self.bias_k = Parameter(torch.empty(1, 1, embed_dim)) | |
| self.bias_v = Parameter(torch.empty(1, 1, embed_dim)) | |
| else: | |
| self.bias_k = self.bias_v = None | |
| self.add_zero_attn = add_zero_attn | |
| self._reset_parameters() | |
| def _reset_parameters(self): | |
| if self._qkv_same_embed_dim: | |
| xavier_uniform_(self.in_proj_weight) | |
| else: | |
| xavier_uniform_(self.q_proj_weight) | |
| xavier_uniform_(self.k_proj_weight) | |
| xavier_uniform_(self.v_proj_weight) | |
| if self.in_proj_bias is not None: | |
| constant_(self.in_proj_bias, 0.) | |
| constant_(self.out_proj.bias, 0.) | |
| if self.bias_k is not None: | |
| xavier_normal_(self.bias_k) | |
| if self.bias_v is not None: | |
| xavier_normal_(self.bias_v) | |
| def __setstate__(self, state): | |
| # Support loading old MultiheadAttention checkpoints generated by v1.1.0 | |
| if '_qkv_same_embed_dim' not in state: | |
| state['_qkv_same_embed_dim'] = True | |
| super(MultiheadAttention, self).__setstate__(state) | |
| def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None, | |
| need_weights: bool = True, attn_mask: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]: | |
| r""" | |
| Args: | |
| query, key, value: map a query and a set of key-value pairs to an output. | |
| See "Attention Is All You Need" for more details. | |
| key_padding_mask: if provided, specified padding elements in the key will | |
| be ignored by the attention. When given a binary mask and a value is True, | |
| the corresponding value on the attention layer will be ignored. When given | |
| a byte mask and a value is non-zero, the corresponding value on the attention | |
| layer will be ignored | |
| need_weights: output attn_output_weights. | |
| attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all | |
| the batches while a 3D mask allows to specify a different mask for the entries of each batch. | |
| Shapes for inputs: | |
| - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is | |
| the embedding dimension. | |
| - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is | |
| the embedding dimension. | |
| - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is | |
| the embedding dimension. | |
| - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length. | |
| If a ByteTensor is provided, the non-zero positions will be ignored while the position | |
| with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the | |
| value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged. | |
| - attn_mask: if a 2D mask: :math:`(L, S)` where L is the target sequence length, S is the | |
| source sequence length. | |
| If a 3D mask: :math:`(N\cdot\text{num\_heads}, L, S)` where N is the batch size, L is the target sequence | |
| length, S is the source sequence length. ``attn_mask`` ensure that position i is allowed to attend | |
| the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend | |
| while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True`` | |
| is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor | |
| is provided, it will be added to the attention weight. | |
| Shapes for outputs: | |
| - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, | |
| E is the embedding dimension. | |
| - attn_output_weights: :math:`(N, L, S)` where N is the batch size, | |
| L is the target sequence length, S is the source sequence length. | |
| """ | |
| if not self._qkv_same_embed_dim: | |
| return multi_head_attention_forward( | |
| query, key, value, self.embed_dim, self.num_heads, | |
| self.in_proj_weight, self.in_proj_bias, | |
| self.bias_k, self.bias_v, self.add_zero_attn, | |
| self.dropout, self.out_proj.weight, self.out_proj.bias, | |
| training=self.training, | |
| key_padding_mask=key_padding_mask, need_weights=need_weights, | |
| attn_mask=attn_mask, use_separate_proj_weight=True, | |
| q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight, | |
| v_proj_weight=self.v_proj_weight) | |
| else: | |
| return multi_head_attention_forward( | |
| query, key, value, self.embed_dim, self.num_heads, | |
| self.in_proj_weight, self.in_proj_bias, | |
| self.bias_k, self.bias_v, self.add_zero_attn, | |
| self.dropout, self.out_proj.weight, self.out_proj.bias, | |
| training=self.training, | |
| key_padding_mask=key_padding_mask, need_weights=need_weights, | |
| attn_mask=attn_mask) |