|
|
import copy |
|
|
import math |
|
|
from typing import Optional |
|
|
|
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
from torch import Tensor, nn |
|
|
|
|
|
|
|
|
"""Transformer class. |
|
|
|
|
|
Copy-paste from torch.nn.Transformer with modifications: |
|
|
* positional encodings are passed in MHattention |
|
|
* extra LN at the end of encoder is removed |
|
|
* decoder returns a stack of activations from all decoding layers |
|
|
""" |
|
|
|
|
|
|
|
|
class Conv2d(torch.nn.Conv2d): |
|
|
"""A wrapper around :class:`torch.nn.Conv2d` to support empty inputs and |
|
|
more features.""" |
|
|
|
|
|
def __init__(self, *args, **kwargs): |
|
|
"""Extra keyword arguments supported in addition to those in |
|
|
`torch.nn.Conv2d`: |
|
|
|
|
|
Args: |
|
|
norm (nn.Module, optional): a normalization layer |
|
|
activation (callable(Tensor) -> Tensor): a callable |
|
|
activation function |
|
|
|
|
|
It assumes that norm layer is used before activation. |
|
|
""" |
|
|
norm = kwargs.pop('norm', None) |
|
|
activation = kwargs.pop('activation', None) |
|
|
super().__init__(*args, **kwargs) |
|
|
|
|
|
self.norm = norm |
|
|
self.activation = activation |
|
|
|
|
|
def forward(self, x): |
|
|
x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding, |
|
|
self.dilation, self.groups) |
|
|
if self.norm is not None: |
|
|
x = self.norm(x) |
|
|
if self.activation is not None: |
|
|
x = self.activation(x) |
|
|
return x |
|
|
|
|
|
|
|
|
class PositionEmbeddingSine(nn.Module): |
|
|
"""This is a more standard version of the position embedding, very similar |
|
|
to the one used by the Attention is all you need paper, generalized to work |
|
|
on images.""" |
|
|
|
|
|
def __init__(self, |
|
|
num_pos_feats=64, |
|
|
temperature=10000, |
|
|
normalize=False, |
|
|
scale=None): |
|
|
super().__init__() |
|
|
self.num_pos_feats = num_pos_feats |
|
|
self.temperature = temperature |
|
|
self.normalize = normalize |
|
|
if scale is not None and normalize is False: |
|
|
raise ValueError('normalize should be True if scale is passed') |
|
|
if scale is None: |
|
|
scale = 2 * math.pi |
|
|
self.scale = scale |
|
|
|
|
|
def forward(self, x, mask=None): |
|
|
if mask is None: |
|
|
mask = torch.zeros((x.size(0), x.size(2), x.size(3)), |
|
|
device=x.device, |
|
|
dtype=torch.bool) |
|
|
not_mask = ~mask |
|
|
y_embed = not_mask.cumsum(1, dtype=x.dtype) |
|
|
x_embed = not_mask.cumsum(2, dtype=x.dtype) |
|
|
if self.normalize: |
|
|
eps = 1e-6 |
|
|
y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale |
|
|
x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale |
|
|
|
|
|
dim_t = torch.arange( |
|
|
self.num_pos_feats, dtype=x.dtype, device=x.device) |
|
|
dim_t = self.temperature**(2 * (dim_t // 2) / self.num_pos_feats) |
|
|
|
|
|
pos_x = x_embed[:, :, :, None] / dim_t |
|
|
pos_y = y_embed[:, :, :, None] / dim_t |
|
|
pos_x = torch.stack( |
|
|
(pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), |
|
|
dim=4).flatten(3) |
|
|
pos_y = torch.stack( |
|
|
(pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), |
|
|
dim=4).flatten(3) |
|
|
pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) |
|
|
return pos |
|
|
|
|
|
def __repr__(self, _repr_indent=4): |
|
|
head = 'Positional encoding ' + self.__class__.__name__ |
|
|
body = [ |
|
|
'num_pos_feats: {}'.format(self.num_pos_feats), |
|
|
'temperature: {}'.format(self.temperature), |
|
|
'normalize: {}'.format(self.normalize), |
|
|
'scale: {}'.format(self.scale), |
|
|
] |
|
|
|
|
|
lines = [head] + [' ' * _repr_indent + line for line in body] |
|
|
return '\n'.join(lines) |
|
|
|
|
|
|
|
|
class TransformerEncoder(nn.Module): |
|
|
|
|
|
def __init__(self, encoder_layer, num_layers, norm=None): |
|
|
super().__init__() |
|
|
self.layers = _get_clones(encoder_layer, num_layers) |
|
|
self.num_layers = num_layers |
|
|
self.norm = norm |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
src, |
|
|
mask: Optional[Tensor] = None, |
|
|
src_key_padding_mask: Optional[Tensor] = None, |
|
|
pos: Optional[Tensor] = None, |
|
|
): |
|
|
output = src |
|
|
|
|
|
for layer in self.layers: |
|
|
output = layer( |
|
|
output, |
|
|
src_mask=mask, |
|
|
src_key_padding_mask=src_key_padding_mask, |
|
|
pos=pos) |
|
|
|
|
|
if self.norm is not None: |
|
|
output = self.norm(output) |
|
|
|
|
|
return output |
|
|
|
|
|
|
|
|
class TransformerEncoderLayer(nn.Module): |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
d_model, |
|
|
nhead, |
|
|
dim_feedforward=2048, |
|
|
dropout=0.1, |
|
|
activation='relu', |
|
|
normalize_before=False, |
|
|
): |
|
|
super().__init__() |
|
|
self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) |
|
|
|
|
|
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.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] |
|
|
src = src + self.dropout1(src2) |
|
|
src = self.norm1(src) |
|
|
src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) |
|
|
src = src + self.dropout2(src2) |
|
|
src = self.norm2(src) |
|
|
return src |
|
|
|
|
|
def forward_pre( |
|
|
self, |
|
|
src, |
|
|
src_mask: Optional[Tensor] = None, |
|
|
src_key_padding_mask: Optional[Tensor] = None, |
|
|
pos: Optional[Tensor] = None, |
|
|
): |
|
|
src2 = self.norm1(src) |
|
|
q = k = self.with_pos_embed(src2, pos) |
|
|
src2 = self.self_attn( |
|
|
q, |
|
|
k, |
|
|
value=src2, |
|
|
attn_mask=src_mask, |
|
|
key_padding_mask=src_key_padding_mask)[0] |
|
|
src = src + self.dropout1(src2) |
|
|
src2 = self.norm2(src) |
|
|
src2 = self.linear2(self.dropout(self.activation(self.linear1(src2)))) |
|
|
src = src + self.dropout2(src2) |
|
|
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) |
|
|
|
|
|
|
|
|
class SelfAttentionLayer(nn.Module): |
|
|
|
|
|
def __init__(self, |
|
|
d_model, |
|
|
nhead, |
|
|
dropout=0.0, |
|
|
activation='relu', |
|
|
normalize_before=False): |
|
|
super().__init__() |
|
|
self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) |
|
|
|
|
|
self.norm = nn.LayerNorm(d_model) |
|
|
self.dropout = nn.Dropout(dropout) |
|
|
|
|
|
self.activation = _get_activation_fn(activation) |
|
|
self.normalize_before = normalize_before |
|
|
|
|
|
self._reset_parameters() |
|
|
|
|
|
def _reset_parameters(self): |
|
|
for p in self.parameters(): |
|
|
if p.dim() > 1: |
|
|
nn.init.xavier_uniform_(p) |
|
|
|
|
|
def with_pos_embed(self, tensor, pos: Optional[Tensor]): |
|
|
return tensor if pos is None else tensor + pos |
|
|
|
|
|
def forward_post(self, |
|
|
tgt, |
|
|
tgt_mask: Optional[Tensor] = None, |
|
|
tgt_key_padding_mask: Optional[Tensor] = None, |
|
|
query_pos: Optional[Tensor] = None): |
|
|
q = k = self.with_pos_embed(tgt, query_pos) |
|
|
tgt2 = self.self_attn( |
|
|
q, |
|
|
k, |
|
|
value=tgt, |
|
|
attn_mask=tgt_mask, |
|
|
key_padding_mask=tgt_key_padding_mask)[0] |
|
|
tgt = tgt + self.dropout(tgt2) |
|
|
tgt = self.norm(tgt) |
|
|
|
|
|
return tgt |
|
|
|
|
|
def forward_pre(self, |
|
|
tgt, |
|
|
tgt_mask: Optional[Tensor] = None, |
|
|
tgt_key_padding_mask: Optional[Tensor] = None, |
|
|
query_pos: Optional[Tensor] = None): |
|
|
tgt2 = self.norm(tgt) |
|
|
q = k = self.with_pos_embed(tgt2, query_pos) |
|
|
tgt2 = self.self_attn( |
|
|
q, |
|
|
k, |
|
|
value=tgt2, |
|
|
attn_mask=tgt_mask, |
|
|
key_padding_mask=tgt_key_padding_mask)[0] |
|
|
tgt = tgt + self.dropout(tgt2) |
|
|
|
|
|
return tgt |
|
|
|
|
|
def forward(self, |
|
|
tgt, |
|
|
tgt_mask: Optional[Tensor] = None, |
|
|
tgt_key_padding_mask: Optional[Tensor] = None, |
|
|
query_pos: Optional[Tensor] = None): |
|
|
if self.normalize_before: |
|
|
return self.forward_pre(tgt, tgt_mask, tgt_key_padding_mask, |
|
|
query_pos) |
|
|
return self.forward_post(tgt, tgt_mask, tgt_key_padding_mask, |
|
|
query_pos) |
|
|
|
|
|
|
|
|
class CrossAttentionLayer(nn.Module): |
|
|
|
|
|
def __init__(self, |
|
|
d_model, |
|
|
nhead, |
|
|
dropout=0.0, |
|
|
activation='relu', |
|
|
normalize_before=False): |
|
|
super().__init__() |
|
|
self.multihead_attn = nn.MultiheadAttention( |
|
|
d_model, nhead, dropout=dropout) |
|
|
|
|
|
self.norm = nn.LayerNorm(d_model) |
|
|
self.dropout = nn.Dropout(dropout) |
|
|
|
|
|
self.activation = _get_activation_fn(activation) |
|
|
self.normalize_before = normalize_before |
|
|
|
|
|
self._reset_parameters() |
|
|
|
|
|
def _reset_parameters(self): |
|
|
for p in self.parameters(): |
|
|
if p.dim() > 1: |
|
|
nn.init.xavier_uniform_(p) |
|
|
|
|
|
def with_pos_embed(self, tensor, pos: Optional[Tensor]): |
|
|
return tensor if pos is None else tensor + pos |
|
|
|
|
|
def forward_post(self, |
|
|
tgt, |
|
|
memory, |
|
|
memory_mask: Optional[Tensor] = None, |
|
|
memory_key_padding_mask: Optional[Tensor] = None, |
|
|
pos: Optional[Tensor] = None, |
|
|
query_pos: Optional[Tensor] = None): |
|
|
tgt2, avg_attn = self.multihead_attn( |
|
|
query=self.with_pos_embed(tgt, query_pos), |
|
|
key=self.with_pos_embed(memory, pos), |
|
|
value=memory, |
|
|
attn_mask=memory_mask, |
|
|
key_padding_mask=memory_key_padding_mask) |
|
|
tgt = tgt + self.dropout(tgt2) |
|
|
tgt = self.norm(tgt) |
|
|
return tgt, avg_attn |
|
|
|
|
|
def forward_pre(self, |
|
|
tgt, |
|
|
memory, |
|
|
memory_mask: Optional[Tensor] = None, |
|
|
memory_key_padding_mask: Optional[Tensor] = None, |
|
|
pos: Optional[Tensor] = None, |
|
|
query_pos: Optional[Tensor] = None): |
|
|
tgt2 = self.norm(tgt) |
|
|
tgt2, avg_attn = self.multihead_attn( |
|
|
query=self.with_pos_embed(tgt2, query_pos), |
|
|
key=self.with_pos_embed(memory, pos), |
|
|
value=memory, |
|
|
attn_mask=memory_mask, |
|
|
key_padding_mask=memory_key_padding_mask) |
|
|
tgt = tgt + self.dropout(tgt2) |
|
|
|
|
|
return tgt, avg_attn |
|
|
|
|
|
def forward(self, |
|
|
tgt, |
|
|
memory, |
|
|
memory_mask: Optional[Tensor] = None, |
|
|
memory_key_padding_mask: Optional[Tensor] = None, |
|
|
pos: Optional[Tensor] = None, |
|
|
query_pos: Optional[Tensor] = None): |
|
|
if self.normalize_before: |
|
|
return self.forward_pre(tgt, memory, memory_mask, |
|
|
memory_key_padding_mask, pos, query_pos) |
|
|
return self.forward_post(tgt, memory, memory_mask, |
|
|
memory_key_padding_mask, pos, query_pos) |
|
|
|
|
|
|
|
|
class FFNLayer(nn.Module): |
|
|
|
|
|
def __init__(self, |
|
|
d_model, |
|
|
dim_feedforward=2048, |
|
|
dropout=0.0, |
|
|
activation='relu', |
|
|
normalize_before=False): |
|
|
super().__init__() |
|
|
|
|
|
self.linear1 = nn.Linear(d_model, dim_feedforward) |
|
|
self.dropout = nn.Dropout(dropout) |
|
|
self.linear2 = nn.Linear(dim_feedforward, d_model) |
|
|
|
|
|
self.norm = nn.LayerNorm(d_model) |
|
|
|
|
|
self.activation = _get_activation_fn(activation) |
|
|
self.normalize_before = normalize_before |
|
|
|
|
|
self._reset_parameters() |
|
|
|
|
|
def _reset_parameters(self): |
|
|
for p in self.parameters(): |
|
|
if p.dim() > 1: |
|
|
nn.init.xavier_uniform_(p) |
|
|
|
|
|
def with_pos_embed(self, tensor, pos: Optional[Tensor]): |
|
|
return tensor if pos is None else tensor + pos |
|
|
|
|
|
def forward_post(self, tgt): |
|
|
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) |
|
|
tgt = tgt + self.dropout(tgt2) |
|
|
tgt = self.norm(tgt) |
|
|
return tgt |
|
|
|
|
|
def forward_pre(self, tgt): |
|
|
tgt2 = self.norm(tgt) |
|
|
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) |
|
|
tgt = tgt + self.dropout(tgt2) |
|
|
return tgt |
|
|
|
|
|
def forward(self, tgt): |
|
|
if self.normalize_before: |
|
|
return self.forward_pre(tgt) |
|
|
return self.forward_post(tgt) |
|
|
|
|
|
|
|
|
class MLP(nn.Module): |
|
|
"""Very simple multi-layer perceptron (also called FFN)""" |
|
|
|
|
|
def __init__(self, input_dim, hidden_dim, output_dim, num_layers): |
|
|
super().__init__() |
|
|
self.num_layers = num_layers |
|
|
h = [hidden_dim] * (num_layers - 1) |
|
|
self.layers = nn.ModuleList( |
|
|
nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) |
|
|
|
|
|
def forward(self, x): |
|
|
for i, layer in enumerate(self.layers): |
|
|
x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) |
|
|
return x |
|
|
|
|
|
|
|
|
def get_norm(norm, out_channels): |
|
|
""" |
|
|
Args: |
|
|
norm (str or callable): either one of BN, SyncBN, FrozenBN, GN; |
|
|
or a callable that takes a channel number and returns |
|
|
the normalization layer as a nn.Module. |
|
|
|
|
|
Returns: |
|
|
nn.Module or None: the normalization layer |
|
|
""" |
|
|
if norm is None: |
|
|
return None |
|
|
if isinstance(norm, str): |
|
|
if len(norm) == 0: |
|
|
return None |
|
|
norm = { |
|
|
'BN': nn.BatchNorm2d, |
|
|
'GN': lambda channels: nn.GroupNorm(32, channels), |
|
|
}[norm] |
|
|
return norm(out_channels) |
|
|
|
|
|
|
|
|
def _get_clones(module, N): |
|
|
return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) |
|
|
|
|
|
|
|
|
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}.') |
|
|
|
