sneedium/captcha_pixelplanet
1
1# pytorch 1.5.02import copy3import math4import warnings5from typing import Optional6 7import torch8import torch.nn as nn9from torch import Tensor10from torch.nn import Dropout, LayerNorm, Linear, Module, ModuleList, Parameter11from torch.nn import functional as F12from torch.nn.init import constant_, xavier_uniform_13 14 15def multi_head_attention_forward(query, # type: Tensor16 key, # type: Tensor17 value, # type: Tensor18 embed_dim_to_check, # type: int19 num_heads, # type: int20 in_proj_weight, # type: Tensor21 in_proj_bias, # type: Tensor22 bias_k, # type: Optional[Tensor]23 bias_v, # type: Optional[Tensor]24 add_zero_attn, # type: bool25 dropout_p, # type: float26 out_proj_weight, # type: Tensor27 out_proj_bias, # type: Tensor28 training=True, # type: bool29 key_padding_mask=None, # type: Optional[Tensor]30 need_weights=True, # type: bool31 attn_mask=None, # type: Optional[Tensor]32 use_separate_proj_weight=False, # type: bool33 q_proj_weight=None, # type: Optional[Tensor]34 k_proj_weight=None, # type: Optional[Tensor]35 v_proj_weight=None, # type: Optional[Tensor]36 static_k=None, # type: Optional[Tensor]37 static_v=None # type: Optional[Tensor]38 ):39 # type: (...) -> Tuple[Tensor, Optional[Tensor]]40 r"""41 Args:42 query, key, value: map a query and a set of key-value pairs to an output.43 See "Attention Is All You Need" for more details.44 embed_dim_to_check: total dimension of the model.45 num_heads: parallel attention heads.46 in_proj_weight, in_proj_bias: input projection weight and bias.47 bias_k, bias_v: bias of the key and value sequences to be added at dim=0.48 add_zero_attn: add a new batch of zeros to the key and49 value sequences at dim=1.50 dropout_p: probability of an element to be zeroed.51 out_proj_weight, out_proj_bias: the output projection weight and bias.52 training: apply dropout if is ``True``.53 key_padding_mask: if provided, specified padding elements in the key will54 be ignored by the attention. This is an binary mask. When the value is True,55 the corresponding value on the attention layer will be filled with -inf.56 need_weights: output attn_output_weights.57 attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all58 the batches while a 3D mask allows to specify a different mask for the entries of each batch.59 use_separate_proj_weight: the function accept the proj. weights for query, key,60 and value in different forms. If false, in_proj_weight will be used, which is61 a combination of q_proj_weight, k_proj_weight, v_proj_weight.62 q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.63 static_k, static_v: static key and value used for attention operators.64 Shape:65 Inputs:66 - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is67 the embedding dimension.68 - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is69 the embedding dimension.70 - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is71 the embedding dimension.72 - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.73 If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions74 will be unchanged. If a BoolTensor is provided, the positions with the75 value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.76 - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.77 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,78 S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked79 positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend80 while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``81 are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor82 is provided, it will be added to the attention weight.83 - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,84 N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.85 - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,86 N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.87 Outputs:88 - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,89 E is the embedding dimension.90 - attn_output_weights: :math:`(N, L, S)` where N is the batch size,91 L is the target sequence length, S is the source sequence length.92 """93 # if not torch.jit.is_scripting():94 # tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v,95 # out_proj_weight, out_proj_bias)96 # if any([type(t) is not Tensor for t in tens_ops]) and has_torch_function(tens_ops):97 # return handle_torch_function(98 # multi_head_attention_forward, tens_ops, query, key, value,99 # embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias,100 # bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight,101 # out_proj_bias, training=training, key_padding_mask=key_padding_mask,102 # need_weights=need_weights, attn_mask=attn_mask,103 # use_separate_proj_weight=use_separate_proj_weight,104 # q_proj_weight=q_proj_weight, k_proj_weight=k_proj_weight,105 # v_proj_weight=v_proj_weight, static_k=static_k, static_v=static_v)106 tgt_len, bsz, embed_dim = query.size()107 assert embed_dim == embed_dim_to_check108 assert key.size() == value.size()109 110 head_dim = embed_dim // num_heads111 assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"112 scaling = float(head_dim) ** -0.5113 114 if not use_separate_proj_weight:115 if torch.equal(query, key) and torch.equal(key, value):116 # self-attention117 q, k, v = F.linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)118 119 elif torch.equal(key, value):120 # encoder-decoder attention121 # This is inline in_proj function with in_proj_weight and in_proj_bias122 _b = in_proj_bias123 _start = 0124 _end = embed_dim125 _w = in_proj_weight[_start:_end, :]126 if _b is not None:127 _b = _b[_start:_end]128 q = F.linear(query, _w, _b)129 130 if key is None:131 assert value is None132 k = None133 v = None134 else:135 136 # This is inline in_proj function with in_proj_weight and in_proj_bias137 _b = in_proj_bias138 _start = embed_dim139 _end = None140 _w = in_proj_weight[_start:, :]141 if _b is not None:142 _b = _b[_start:]143 k, v = F.linear(key, _w, _b).chunk(2, dim=-1)144 145 else:146 # This is inline in_proj function with in_proj_weight and in_proj_bias147 _b = in_proj_bias148 _start = 0149 _end = embed_dim150 _w = in_proj_weight[_start:_end, :]151 if _b is not None:152 _b = _b[_start:_end]153 q = F.linear(query, _w, _b)154 155 # This is inline in_proj function with in_proj_weight and in_proj_bias156 _b = in_proj_bias157 _start = embed_dim158 _end = embed_dim * 2159 _w = in_proj_weight[_start:_end, :]160 if _b is not None:161 _b = _b[_start:_end]162 k = F.linear(key, _w, _b)163 164 # This is inline in_proj function with in_proj_weight and in_proj_bias165 _b = in_proj_bias166 _start = embed_dim * 2167 _end = None168 _w = in_proj_weight[_start:, :]169 if _b is not None:170 _b = _b[_start:]171 v = F.linear(value, _w, _b)172 else:173 q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight)174 len1, len2 = q_proj_weight_non_opt.size()175 assert len1 == embed_dim and len2 == query.size(-1)176 177 k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight)178 len1, len2 = k_proj_weight_non_opt.size()179 assert len1 == embed_dim and len2 == key.size(-1)180 181 v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight)182 len1, len2 = v_proj_weight_non_opt.size()183 assert len1 == embed_dim and len2 == value.size(-1)184 185 if in_proj_bias is not None:186 q = F.linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim])187 k = F.linear(key, k_proj_weight_non_opt, in_proj_bias[embed_dim:(embed_dim * 2)])188 v = F.linear(value, v_proj_weight_non_opt, in_proj_bias[(embed_dim * 2):])189 else:190 q = F.linear(query, q_proj_weight_non_opt, in_proj_bias)191 k = F.linear(key, k_proj_weight_non_opt, in_proj_bias)192 v = F.linear(value, v_proj_weight_non_opt, in_proj_bias)193 q = q * scaling194 195 if attn_mask is not None:196 assert attn_mask.dtype == torch.float32 or attn_mask.dtype == torch.float64 or \197 attn_mask.dtype == torch.float16 or attn_mask.dtype == torch.uint8 or attn_mask.dtype == torch.bool, \198 'Only float, byte, and bool types are supported for attn_mask, not {}'.format(attn_mask.dtype)199 if attn_mask.dtype == torch.uint8:200 warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")201 attn_mask = attn_mask.to(torch.bool)202 203 if attn_mask.dim() == 2:204 attn_mask = attn_mask.unsqueeze(0)205 if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:206 raise RuntimeError('The size of the 2D attn_mask is not correct.')207 elif attn_mask.dim() == 3:208 if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:209 raise RuntimeError('The size of the 3D attn_mask is not correct.')210 else:211 raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))212 # attn_mask's dim is 3 now.213 214 # # convert ByteTensor key_padding_mask to bool215 # if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:216 # warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")217 # key_padding_mask = key_padding_mask.to(torch.bool)218 219 if bias_k is not None and bias_v is not None:220 if static_k is None and static_v is None:221 k = torch.cat([k, bias_k.repeat(1, bsz, 1)])222 v = torch.cat([v, bias_v.repeat(1, bsz, 1)])223 if attn_mask is not None:224 attn_mask = pad(attn_mask, (0, 1))225 if key_padding_mask is not None:226 key_padding_mask = pad(key_padding_mask, (0, 1))227 else:228 assert static_k is None, "bias cannot be added to static key."229 assert static_v is None, "bias cannot be added to static value."230 else:231 assert bias_k is None232 assert bias_v is None233 234 q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)235 if k is not None:236 k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)237 if v is not None:238 v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)239 240 if static_k is not None:241 assert static_k.size(0) == bsz * num_heads242 assert static_k.size(2) == head_dim243 k = static_k244 245 if static_v is not None:246 assert static_v.size(0) == bsz * num_heads247 assert static_v.size(2) == head_dim248 v = static_v249 250 src_len = k.size(1)251 252 if key_padding_mask is not None:253 assert key_padding_mask.size(0) == bsz254 assert key_padding_mask.size(1) == src_len255 256 if add_zero_attn:257 src_len += 1258 k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1)259 v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1)260 if attn_mask is not None:261 attn_mask = pad(attn_mask, (0, 1))262 if key_padding_mask is not None:263 key_padding_mask = pad(key_padding_mask, (0, 1))264 265 attn_output_weights = torch.bmm(q, k.transpose(1, 2))266 assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]267 268 if attn_mask is not None:269 if attn_mask.dtype == torch.bool:270 attn_output_weights.masked_fill_(attn_mask, float('-inf'))271 else:272 attn_output_weights += attn_mask273 274 275 if key_padding_mask is not None:276 attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)277 attn_output_weights = attn_output_weights.masked_fill(278 key_padding_mask.unsqueeze(1).unsqueeze(2),279 float('-inf'),280 )281 attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)282 283 attn_output_weights = F.softmax(284 attn_output_weights, dim=-1)285 attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)286 287 attn_output = torch.bmm(attn_output_weights, v)288 assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]289 attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)290 attn_output = F.linear(attn_output, out_proj_weight, out_proj_bias)291 292 if need_weights:293 # average attention weights over heads294 attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)295 return attn_output, attn_output_weights.sum(dim=1) / num_heads296 else:297 return attn_output, None298 299class MultiheadAttention(Module):300 r"""Allows the model to jointly attend to information301 from different representation subspaces.302 See reference: Attention Is All You Need303 .. math::304 \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O305 \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)306 Args:307 embed_dim: total dimension of the model.308 num_heads: parallel attention heads.309 dropout: a Dropout layer on attn_output_weights. Default: 0.0.310 bias: add bias as module parameter. Default: True.311 add_bias_kv: add bias to the key and value sequences at dim=0.312 add_zero_attn: add a new batch of zeros to the key and313 value sequences at dim=1.314 kdim: total number of features in key. Default: None.315 vdim: total number of features in value. Default: None.316 Note: if kdim and vdim are None, they will be set to embed_dim such that317 query, key, and value have the same number of features.318 Examples::319 >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)320 >>> attn_output, attn_output_weights = multihead_attn(query, key, value)321 """322 # __annotations__ = {323 # 'bias_k': torch._jit_internal.Optional[torch.Tensor],324 # 'bias_v': torch._jit_internal.Optional[torch.Tensor],325 # }326 __constants__ = ['q_proj_weight', 'k_proj_weight', 'v_proj_weight', 'in_proj_weight']327 328 def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):329 super(MultiheadAttention, self).__init__()330 self.embed_dim = embed_dim331 self.kdim = kdim if kdim is not None else embed_dim332 self.vdim = vdim if vdim is not None else embed_dim333 self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim334 335 self.num_heads = num_heads336 self.dropout = dropout337 self.head_dim = embed_dim // num_heads338 assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"339 340 if self._qkv_same_embed_dim is False:341 self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))342 self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))343 self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))344 self.register_parameter('in_proj_weight', None)345 else:346 self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))347 self.register_parameter('q_proj_weight', None)348 self.register_parameter('k_proj_weight', None)349 self.register_parameter('v_proj_weight', None)350 351 if bias:352 self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))353 else:354 self.register_parameter('in_proj_bias', None)355 self.out_proj = Linear(embed_dim, embed_dim, bias=bias)356 357 if add_bias_kv:358 self.bias_k = Parameter(torch.empty(1, 1, embed_dim))359 self.bias_v = Parameter(torch.empty(1, 1, embed_dim))360 else:361 self.bias_k = self.bias_v = None362 363 self.add_zero_attn = add_zero_attn364 365 self._reset_parameters()366 367 def _reset_parameters(self):368 if self._qkv_same_embed_dim:369 xavier_uniform_(self.in_proj_weight)370 else:371 xavier_uniform_(self.q_proj_weight)372 xavier_uniform_(self.k_proj_weight)373 xavier_uniform_(self.v_proj_weight)374 375 if self.in_proj_bias is not None:376 constant_(self.in_proj_bias, 0.)377 constant_(self.out_proj.bias, 0.)378 if self.bias_k is not None:379 xavier_normal_(self.bias_k)380 if self.bias_v is not None:381 xavier_normal_(self.bias_v)382 383 def __setstate__(self, state):384 # Support loading old MultiheadAttention checkpoints generated by v1.1.0385 if '_qkv_same_embed_dim' not in state:386 state['_qkv_same_embed_dim'] = True387 388 super(MultiheadAttention, self).__setstate__(state)389 390 def forward(self, query, key, value, key_padding_mask=None,391 need_weights=True, attn_mask=None):392 # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]393 r"""394 Args:395 query, key, value: map a query and a set of key-value pairs to an output.396 See "Attention Is All You Need" for more details.397 key_padding_mask: if provided, specified padding elements in the key will398 be ignored by the attention. This is an binary mask. When the value is True,399 the corresponding value on the attention layer will be filled with -inf.400 need_weights: output attn_output_weights.401 attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all402 the batches while a 3D mask allows to specify a different mask for the entries of each batch.403 Shape:404 - Inputs:405 - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is406 the embedding dimension.407 - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is408 the embedding dimension.409 - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is410 the embedding dimension.411 - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.412 If a ByteTensor is provided, the non-zero positions will be ignored while the position413 with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the414 value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.415 - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.416 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,417 S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked418 positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend419 while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``420 is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor421 is provided, it will be added to the attention weight.422 - Outputs:423 - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,424 E is the embedding dimension.425 - attn_output_weights: :math:`(N, L, S)` where N is the batch size,426 L is the target sequence length, S is the source sequence length.427 """428 if not self._qkv_same_embed_dim:429 return multi_head_attention_forward(430 query, key, value, self.embed_dim, self.num_heads,431 self.in_proj_weight, self.in_proj_bias,432 self.bias_k, self.bias_v, self.add_zero_attn,433 self.dropout, self.out_proj.weight, self.out_proj.bias,434 training=self.training,435 key_padding_mask=key_padding_mask, need_weights=need_weights,436 attn_mask=attn_mask, use_separate_proj_weight=True,437 q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,438 v_proj_weight=self.v_proj_weight)439 else:440 return multi_head_attention_forward(441 query, key, value, self.embed_dim, self.num_heads,442 self.in_proj_weight, self.in_proj_bias,443 self.bias_k, self.bias_v, self.add_zero_attn,444 self.dropout, self.out_proj.weight, self.out_proj.bias,445 training=self.training,446 key_padding_mask=key_padding_mask, need_weights=need_weights,447 attn_mask=attn_mask)448 449 450class Transformer(Module):451 r"""A transformer model. User is able to modify the attributes as needed. The architecture452 is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,453 Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and454 Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information455 Processing Systems, pages 6000-6010. Users can build the BERT(https://arxiv.org/abs/1810.04805)456 model with corresponding parameters.457 458 Args:459 d_model: the number of expected features in the encoder/decoder inputs (default=512).460 nhead: the number of heads in the multiheadattention models (default=8).461 num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).462 num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).463 dim_feedforward: the dimension of the feedforward network model (default=2048).464 dropout: the dropout value (default=0.1).465 activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).466 custom_encoder: custom encoder (default=None).467 custom_decoder: custom decoder (default=None).468 469 Examples::470 >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)471 >>> src = torch.rand((10, 32, 512))472 >>> tgt = torch.rand((20, 32, 512))473 >>> out = transformer_model(src, tgt)474 475 Note: A full example to apply nn.Transformer module for the word language model is available in476 https://github.com/pytorch/examples/tree/master/word_language_model477 """478 479 def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,480 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,481 activation="relu", custom_encoder=None, custom_decoder=None):482 super(Transformer, self).__init__()483 484 if custom_encoder is not None:485 self.encoder = custom_encoder486 else:487 encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)488 encoder_norm = LayerNorm(d_model)489 self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)490 491 if custom_decoder is not None:492 self.decoder = custom_decoder493 else:494 decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation)495 decoder_norm = LayerNorm(d_model)496 self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)497 498 self._reset_parameters()499 500 self.d_model = d_model501 self.nhead = nhead502 503 def forward(self, src, tgt, src_mask=None, tgt_mask=None,504 memory_mask=None, src_key_padding_mask=None,505 tgt_key_padding_mask=None, memory_key_padding_mask=None):506 # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor # noqa507 r"""Take in and process masked source/target sequences.508 509 Args:510 src: the sequence to the encoder (required).511 tgt: the sequence to the decoder (required).512 src_mask: the additive mask for the src sequence (optional).513 tgt_mask: the additive mask for the tgt sequence (optional).514 memory_mask: the additive mask for the encoder output (optional).515 src_key_padding_mask: the ByteTensor mask for src keys per batch (optional).516 tgt_key_padding_mask: the ByteTensor mask for tgt keys per batch (optional).517 memory_key_padding_mask: the ByteTensor mask for memory keys per batch (optional).518 519 Shape:520 - src: :math:`(S, N, E)`.521 - tgt: :math:`(T, N, E)`.522 - src_mask: :math:`(S, S)`.523 - tgt_mask: :math:`(T, T)`.524 - memory_mask: :math:`(T, S)`.525 - src_key_padding_mask: :math:`(N, S)`.526 - tgt_key_padding_mask: :math:`(N, T)`.527 - memory_key_padding_mask: :math:`(N, S)`.528 529 Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked530 positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend531 while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``532 are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor533 is provided, it will be added to the attention weight. 534 [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by535 the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero536 positions will be unchanged. If a BoolTensor is provided, the positions with the537 value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.538 539 - output: :math:`(T, N, E)`.540 541 Note: Due to the multi-head attention architecture in the transformer model,542 the output sequence length of a transformer is same as the input sequence543 (i.e. target) length of the decode.544 545 where S is the source sequence length, T is the target sequence length, N is the546 batch size, E is the feature number547 548 Examples:549 >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)550 """551 552 if src.size(1) != tgt.size(1):553 raise RuntimeError("the batch number of src and tgt must be equal")554 555 if src.size(2) != self.d_model or tgt.size(2) != self.d_model:556 raise RuntimeError("the feature number of src and tgt must be equal to d_model")557 558 memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)559 output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,560 tgt_key_padding_mask=tgt_key_padding_mask,561 memory_key_padding_mask=memory_key_padding_mask)562 return output563 564 def generate_square_subsequent_mask(self, sz):565 r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').566 Unmasked positions are filled with float(0.0).567 """568 mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)569 mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))570 return mask571 572 def _reset_parameters(self):573 r"""Initiate parameters in the transformer model."""574 575 for p in self.parameters():576 if p.dim() > 1:577 xavier_uniform_(p)578 579 580class TransformerEncoder(Module):581 r"""TransformerEncoder is a stack of N encoder layers582 583 Args:584 encoder_layer: an instance of the TransformerEncoderLayer() class (required).585 num_layers: the number of sub-encoder-layers in the encoder (required).586 norm: the layer normalization component (optional).587 588 Examples::589 >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)590 >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)591 >>> src = torch.rand(10, 32, 512)592 >>> out = transformer_encoder(src)593 """594 __constants__ = ['norm']595 596 def __init__(self, encoder_layer, num_layers, norm=None):597 super(TransformerEncoder, self).__init__()598 self.layers = _get_clones(encoder_layer, num_layers)599 self.num_layers = num_layers600 self.norm = norm601 602 def forward(self, src, mask=None, src_key_padding_mask=None):603 # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor604 r"""Pass the input through the encoder layers in turn.605 606 Args:607 src: the sequence to the encoder (required).608 mask: the mask for the src sequence (optional).609 src_key_padding_mask: the mask for the src keys per batch (optional).610 611 Shape:612 see the docs in Transformer class.613 """614 output = src615 616 for i, mod in enumerate(self.layers):617 output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)618 619 if self.norm is not None:620 output = self.norm(output)621 622 return output623 624 625class TransformerDecoder(Module):626 r"""TransformerDecoder is a stack of N decoder layers627 628 Args:629 decoder_layer: an instance of the TransformerDecoderLayer() class (required).630 num_layers: the number of sub-decoder-layers in the decoder (required).631 norm: the layer normalization component (optional).632 633 Examples::634 >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)635 >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)636 >>> memory = torch.rand(10, 32, 512)637 >>> tgt = torch.rand(20, 32, 512)638 >>> out = transformer_decoder(tgt, memory)639 """640 __constants__ = ['norm']641 642 def __init__(self, decoder_layer, num_layers, norm=None):643 super(TransformerDecoder, self).__init__()644 self.layers = _get_clones(decoder_layer, num_layers)645 self.num_layers = num_layers646 self.norm = norm647 648 def forward(self, tgt, memory, memory2=None, tgt_mask=None,649 memory_mask=None, memory_mask2=None, tgt_key_padding_mask=None,650 memory_key_padding_mask=None, memory_key_padding_mask2=None):651 # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor652 r"""Pass the inputs (and mask) through the decoder layer in turn.653 654 Args:655 tgt: the sequence to the decoder (required).656 memory: the sequence from the last layer of the encoder (required).657 tgt_mask: the mask for the tgt sequence (optional).658 memory_mask: the mask for the memory sequence (optional).659 tgt_key_padding_mask: the mask for the tgt keys per batch (optional).660 memory_key_padding_mask: the mask for the memory keys per batch (optional).661 662 Shape:663 see the docs in Transformer class.664 """665 output = tgt666 667 for mod in self.layers:668 output = mod(output, memory, memory2=memory2, tgt_mask=tgt_mask,669 memory_mask=memory_mask, memory_mask2=memory_mask2,670 tgt_key_padding_mask=tgt_key_padding_mask,671 memory_key_padding_mask=memory_key_padding_mask,672 memory_key_padding_mask2=memory_key_padding_mask2)673 674 if self.norm is not None:675 output = self.norm(output)676 677 return output678 679class TransformerEncoderLayer(Module):680 r"""TransformerEncoderLayer is made up of self-attn and feedforward network.681 This standard encoder layer is based on the paper "Attention Is All You Need".682 Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,683 Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in684 Neural Information Processing Systems, pages 6000-6010. Users may modify or implement685 in a different way during application.686 687 Args:688 d_model: the number of expected features in the input (required).689 nhead: the number of heads in the multiheadattention models (required).690 dim_feedforward: the dimension of the feedforward network model (default=2048).691 dropout: the dropout value (default=0.1).692 activation: the activation function of intermediate layer, relu or gelu (default=relu).693 694 Examples::695 >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)696 >>> src = torch.rand(10, 32, 512)697 >>> out = encoder_layer(src)698 """699 700 def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, 701 activation="relu", debug=False):702 super(TransformerEncoderLayer, self).__init__()703 self.debug = debug704 self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)705 # Implementation of Feedforward model706 self.linear1 = Linear(d_model, dim_feedforward)707 self.dropout = Dropout(dropout)708 self.linear2 = Linear(dim_feedforward, d_model)709 710 self.norm1 = LayerNorm(d_model)711 self.norm2 = LayerNorm(d_model)712 self.dropout1 = Dropout(dropout)713 self.dropout2 = Dropout(dropout)714 715 self.activation = _get_activation_fn(activation)716 717 def __setstate__(self, state):718 if 'activation' not in state:719 state['activation'] = F.relu720 super(TransformerEncoderLayer, self).__setstate__(state)721 722 def forward(self, src, src_mask=None, src_key_padding_mask=None):723 # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor724 r"""Pass the input through the encoder layer.725 726 Args:727 src: the sequence to the encoder layer (required).728 src_mask: the mask for the src sequence (optional).729 src_key_padding_mask: the mask for the src keys per batch (optional).730 731 Shape:732 see the docs in Transformer class.733 """734 src2, attn = self.self_attn(src, src, src, attn_mask=src_mask,735 key_padding_mask=src_key_padding_mask)736 if self.debug: self.attn = attn737 src = src + self.dropout1(src2)738 src = self.norm1(src)739 src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))740 src = src + self.dropout2(src2)741 src = self.norm2(src)742 743 return src744 745 746class TransformerDecoderLayer(Module):747 r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.748 This standard decoder layer is based on the paper "Attention Is All You Need".749 Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,750 Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in751 Neural Information Processing Systems, pages 6000-6010. Users may modify or implement752 in a different way during application.753 754 Args:755 d_model: the number of expected features in the input (required).756 nhead: the number of heads in the multiheadattention models (required).757 dim_feedforward: the dimension of the feedforward network model (default=2048).758 dropout: the dropout value (default=0.1).759 activation: the activation function of intermediate layer, relu or gelu (default=relu).760 761 Examples::762 >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)763 >>> memory = torch.rand(10, 32, 512)764 >>> tgt = torch.rand(20, 32, 512)765 >>> out = decoder_layer(tgt, memory)766 """767 768 def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, 769 activation="relu", self_attn=True, siamese=False, debug=False):770 super(TransformerDecoderLayer, self).__init__()771 self.has_self_attn, self.siamese = self_attn, siamese772 self.debug = debug773 if self.has_self_attn:774 self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)775 self.norm1 = LayerNorm(d_model)776 self.dropout1 = Dropout(dropout)777 self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)778 # Implementation of Feedforward model779 self.linear1 = Linear(d_model, dim_feedforward)780 self.dropout = Dropout(dropout)781 self.linear2 = Linear(dim_feedforward, d_model)782 783 self.norm2 = LayerNorm(d_model)784 self.norm3 = LayerNorm(d_model)785 self.dropout2 = Dropout(dropout)786 self.dropout3 = Dropout(dropout)787 if self.siamese:788 self.multihead_attn2 = MultiheadAttention(d_model, nhead, dropout=dropout)789 790 self.activation = _get_activation_fn(activation)791 792 def __setstate__(self, state):793 if 'activation' not in state:794 state['activation'] = F.relu795 super(TransformerDecoderLayer, self).__setstate__(state)796 797 def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,798 tgt_key_padding_mask=None, memory_key_padding_mask=None,799 memory2=None, memory_mask2=None, memory_key_padding_mask2=None):800 # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor801 r"""Pass the inputs (and mask) through the decoder layer.802 803 Args:804 tgt: the sequence to the decoder layer (required).805 memory: the sequence from the last layer of the encoder (required).806 tgt_mask: the mask for the tgt sequence (optional).807 memory_mask: the mask for the memory sequence (optional).808 tgt_key_padding_mask: the mask for the tgt keys per batch (optional).809 memory_key_padding_mask: the mask for the memory keys per batch (optional).810 811 Shape:812 see the docs in Transformer class.813 """814 if self.has_self_attn:815 tgt2, attn = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,816 key_padding_mask=tgt_key_padding_mask)817 tgt = tgt + self.dropout1(tgt2)818 tgt = self.norm1(tgt)819 if self.debug: self.attn = attn820 tgt2, attn2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,821 key_padding_mask=memory_key_padding_mask)822 if self.debug: self.attn2 = attn2823 824 if self.siamese:825 tgt3, attn3 = self.multihead_attn2(tgt, memory2, memory2, attn_mask=memory_mask2,826 key_padding_mask=memory_key_padding_mask2)827 tgt = tgt + self.dropout2(tgt3)828 if self.debug: self.attn3 = attn3829 830 tgt = tgt + self.dropout2(tgt2)831 tgt = self.norm2(tgt)832 tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))833 tgt = tgt + self.dropout3(tgt2)834 tgt = self.norm3(tgt)835 836 return tgt837 838 839def _get_clones(module, N):840 return ModuleList([copy.deepcopy(module) for i in range(N)])841 842 843def _get_activation_fn(activation):844 if activation == "relu":845 return F.relu846 elif activation == "gelu":847 return F.gelu848 849 raise RuntimeError("activation should be relu/gelu, not {}".format(activation))850 851 852class PositionalEncoding(nn.Module):853 r"""Inject some information about the relative or absolute position of the tokens854 in the sequence. The positional encodings have the same dimension as855 the embeddings, so that the two can be summed. Here, we use sine and cosine856 functions of different frequencies.857 .. math::858 \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))859 \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))860 \text{where pos is the word position and i is the embed idx)861 Args:862 d_model: the embed dim (required).863 dropout: the dropout value (default=0.1).864 max_len: the max. length of the incoming sequence (default=5000).865 Examples:866 >>> pos_encoder = PositionalEncoding(d_model)867 """868 869 def __init__(self, d_model, dropout=0.1, max_len=5000):870 super(PositionalEncoding, self).__init__()871 self.dropout = nn.Dropout(p=dropout)872 873 pe = torch.zeros(max_len, d_model)874 position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)875 div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))876 pe[:, 0::2] = torch.sin(position * div_term)877 pe[:, 1::2] = torch.cos(position * div_term)878 pe = pe.unsqueeze(0).transpose(0, 1)879 self.register_buffer('pe', pe)880 881 def forward(self, x):882 r"""Inputs of forward function883 Args:884 x: the sequence fed to the positional encoder model (required).885 Shape:886 x: [sequence length, batch size, embed dim]887 output: [sequence length, batch size, embed dim]888 Examples:889 >>> output = pos_encoder(x)890 """891 892 x = x + self.pe[:x.size(0), :]893 return self.dropout(x)894 895 896if __name__ == '__main__':897 transformer_model = Transformer(nhead=16, num_encoder_layers=12)898 src = torch.rand((10, 32, 512))899 tgt = torch.rand((20, 32, 512))900 out = transformer_model(src, tgt)901 print(out)902