import copy import pdb from typing import Optional import torch import torch.nn.functional as F from torch import nn, Tensor def mask_logits(inputs, mask, mask_value=-1e30): mask = mask.type(torch.float32) return inputs + (1.0 - mask) * mask_value class Transformer(nn.Module): def __init__(self, d_model=512, nhead=8, num_encoder_layers=4, num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, droppath=0.1, activation="gelu", normalize_before=False, # False as default return_intermediate_dec=False): super().__init__() encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, droppath, activation, normalize_before) encoder_norm = nn.LayerNorm(d_model) if normalize_before else None self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) self._reset_parameters() self.d_model = d_model self.nhead = nhead def _reset_parameters(self): for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) def forward(self, src, mask, pos_embed): """ Args: src: (batch_size, L, d) mask: (batch_size, L) query_embed: (#queries, d) -> my imple (batch_size, d) and #queries=1 pos_embed: (batch_size, L, d) the same as src Returns: """ # flatten NxCxHxW to HWxNxC src = src.permute(1, 0, 2) # (L, batch_size, d) pos_embed = pos_embed.permute(1, 0, 2) # (L, batch_size, d) memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed) memory = memory.transpose(0, 1) return memory class TransformerEncoder(nn.Module): def __init__(self, encoder_layer, num_layers, norm=None, return_intermediate=False): super().__init__() self.layers = _get_clones(encoder_layer, num_layers) self.num_layers = num_layers self.norm = norm self.return_intermediate = return_intermediate def forward(self, src, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): output = src intermediate = [] for layer in self.layers: output = layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos) if self.return_intermediate: intermediate.append(output) if self.norm is not None: output = self.norm(output) if self.return_intermediate: return torch.stack(intermediate) return output class TransformerEncoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, droppath=0.1, activation="relu", normalize_before=False): super().__init__() self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) # Implementation of Feedforward model self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) # self.dropout1 = nn.Dropout(dropout) # self.dropout2 = nn.Dropout(dropout) self.droppath1 = DropPath(droppath) self.droppath2 = DropPath(droppath) self.activation = _get_activation_fn(activation) self.normalize_before = normalize_before def with_pos_embed(self, tensor, pos: Optional[Tensor]): return tensor if pos is None else tensor + pos def forward_post(self, src, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): q = k = self.with_pos_embed(src, pos) src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] # src2 = self.self_attn_eff(q=q, k=k, v=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] src = src + self.droppath1(src2) src = self.norm1(src) src2 = self.linear2(self.activation(self.linear1(src))) # src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) src = src + self.droppath2(src2) src = self.norm2(src) return src def forward(self, src, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): if self.normalize_before: return self.forward_pre(src, src_mask, src_key_padding_mask, pos) return self.forward_post(src, src_mask, src_key_padding_mask, pos) def _get_clones(module, N): return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) def build_transformer(args): return Transformer( d_model=args.hidden_dim, dropout=args.dropout, droppath=args.droppath, nhead=args.nheads, dim_feedforward=args.dim_feedforward, num_encoder_layers=args.enc_layers, num_decoder_layers=args.dec_layers, normalize_before=args.pre_norm, return_intermediate_dec=True, ) def drop_path(x, drop_prob=0.0, training=False): """ Stochastic Depth per sample. """ if drop_prob == 0.0 or not training: return x keep_prob = 1 - drop_prob shape = (x.shape[0],) + (1,) * (x.ndim - 1) mask = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) mask.floor_() x = x.div(keep_prob) * mask return x class DropPath(nn.Module): """ Drop paths per sample (when applied in main path of residual blocks). """ def __init__(self, drop_prob=None): super(DropPath, self).__init__() self.drop_prob = drop_prob def forward(self, x): x = x.permute(1, 0, 2) res = drop_path(x, self.drop_prob, self.training) return res.permute(1, 0, 2) # return drop_path(x, self.drop_prob, self.training) def _get_activation_fn(activation): """Return an activation function given a string""" if activation == "relu": return F.relu if activation == "gelu": return F.gelu if activation == "glu": return F.glu raise RuntimeError(F"activation should be relu/gelu, not {activation}.")