| |
| |
|
|
| import math |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| try: |
| from apex.normalization import FusedLayerNorm as LayerNorm |
| except ModuleNotFoundError: |
| from torch.nn import LayerNorm |
|
|
| from .multiway_network import MultiwayWrapper |
| from .xpos_relative_position import XPOS |
|
|
|
|
| class MultiheadAttention(nn.Module): |
| def __init__( |
| self, |
| args, |
| embed_dim, |
| num_heads, |
| dropout=0.0, |
| self_attention=False, |
| encoder_decoder_attention=False, |
| subln=False, |
| one_attn=False, |
| ): |
| super().__init__() |
| self.args = args |
| self.embed_dim = embed_dim |
| self.num_heads = num_heads |
| self.head_dim = embed_dim // num_heads |
| self.scaling = self.head_dim ** (-0.5) |
| self.self_attention = self_attention |
| self.encoder_decoder_attention = encoder_decoder_attention |
| assert self.self_attention ^ self.encoder_decoder_attention |
| if one_attn: |
| self.k_proj = nn.Linear(embed_dim, embed_dim, bias=True) |
| self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True) |
| self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True) |
| self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True) |
| else: |
| self.k_proj = MultiwayWrapper(args, nn.Linear(embed_dim, embed_dim, bias=True)) |
| self.v_proj = MultiwayWrapper(args, nn.Linear(embed_dim, embed_dim, bias=True)) |
| self.q_proj = MultiwayWrapper(args, nn.Linear(embed_dim, embed_dim, bias=True)) |
| |
| self.out_proj = MultiwayWrapper(args, nn.Linear(embed_dim, embed_dim, bias=True)) |
| self.inner_attn_ln = ( |
| MultiwayWrapper(args, LayerNorm(self.embed_dim, eps=args.layernorm_eps)) |
| if subln and self.self_attention |
| else None |
| ) |
| self.dropout_module = torch.nn.Dropout(dropout) |
| self.xpos = XPOS(self.head_dim, args.xpos_scale_base) if args.xpos_rel_pos and self.self_attention else None |
|
|
| def reset_parameters(self): |
| nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2)) |
| nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2)) |
| nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2)) |
| nn.init.xavier_uniform_(self.out_proj.weight) |
| nn.init.constant_(self.out_proj.bias, 0.0) |
|
|
| def forward( |
| self, |
| query, |
| key, |
| value, |
| incremental_state=None, |
| key_padding_mask=None, |
| attn_mask=None, |
| rel_pos=None, |
| ): |
| bsz, tgt_len, embed_dim = query.size() |
| src_len = tgt_len |
| assert embed_dim == self.embed_dim, f"query dim {embed_dim} != {self.embed_dim}" |
|
|
| key_bsz, src_len, _ = key.size() |
| assert key_bsz == bsz, f"{query.size(), key.size()}" |
| assert value is not None |
| assert bsz, src_len == value.shape[:2] |
| |
| |
| |
| |
| |
| |
|
|
| q = self.q_proj(query) |
| k = self.k_proj(key) |
| v = self.v_proj(value) |
|
|
| q = (q * self.scaling).view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2) |
| k = k.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2) |
| v = v.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2) |
| q = q.reshape(bsz * self.num_heads, tgt_len, self.head_dim) |
| k = k.reshape(bsz * self.num_heads, src_len, self.head_dim) |
| v = v.reshape(bsz * self.num_heads, src_len, self.head_dim) |
|
|
| if incremental_state is not None: |
| if "prev_key" in incremental_state: |
| prev_key = incremental_state["prev_key"].view(bsz * self.num_heads, -1, self.head_dim) |
| prev_value = incremental_state["prev_value"].view(bsz * self.num_heads, -1, self.head_dim) |
| k = torch.cat([prev_key, k], dim=1) |
| v = torch.cat([prev_value, v], dim=1) |
| incremental_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim) |
| incremental_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim) |
| src_len = k.size(1) |
|
|
| if self.xpos is not None: |
| if incremental_state is not None: |
| offset = src_len - 1 |
| else: |
| offset = 0 |
| k = self.xpos(k, offset=0, downscale=True) |
| q = self.xpos(q, offset=offset, downscale=False) |
|
|
| attn_weights = torch.bmm(q, k.transpose(1, 2)) |
|
|
| if attn_mask is not None: |
| attn_weights = torch.nan_to_num(attn_weights) |
| if len(attn_mask.shape) != len(attn_weights.shape): |
| attn_mask = attn_mask.unsqueeze(0) |
| else: |
| attn_mask = attn_mask.repeat_interleave(self.num_heads, dim=0) |
| attn_weights += attn_mask |
|
|
| if key_padding_mask is not None: |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) |
| attn_weights = attn_weights.masked_fill( |
| key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), |
| float("-inf"), |
| ) |
| attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) |
|
|
| if rel_pos is not None: |
| rel_pos = rel_pos.view(attn_weights.size()) |
| attn_weights = attn_weights + rel_pos |
|
|
| attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as(attn_weights) |
| attn_probs = self.dropout_module(attn_weights) |
|
|
| attn = torch.bmm(attn_probs, v) |
| attn = attn.transpose(0, 1).reshape(tgt_len, bsz, embed_dim).transpose(0, 1) |
|
|
| if self.inner_attn_ln is not None: |
| attn = self.inner_attn_ln(attn) |
|
|
| attn = self.out_proj(attn) |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0) |
|
|
| return attn, attn_weights |
|
|
