Spaces:
Runtime error
Runtime error
| # MIT License | |
| # Copyright (c) 2022 Intelligent Systems Lab Org | |
| # Permission is hereby granted, free of charge, to any person obtaining a copy | |
| # of this software and associated documentation files (the "Software"), to deal | |
| # in the Software without restriction, including without limitation the rights | |
| # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell | |
| # copies of the Software, and to permit persons to whom the Software is | |
| # furnished to do so, subject to the following conditions: | |
| # The above copyright notice and this permission notice shall be included in all | |
| # copies or substantial portions of the Software. | |
| # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | |
| # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, | |
| # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE | |
| # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER | |
| # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, | |
| # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE | |
| # SOFTWARE. | |
| # File author: Zhenyu Li | |
| # This file mainly copy-and-paste from DETR (https://github.com/facebookresearch/detr/blob/main/models/transformer.py); author: Nicolas Carion | |
| import copy | |
| from typing import Optional, List | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn, Tensor | |
| from timm.models.layers import DropPath, to_2tuple, trunc_normal_ | |
| 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}.") | |
| class TransformerDecoderLayer(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.multihead_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.norm3 = nn.LayerNorm(d_model) | |
| self.dropout1 = nn.Dropout(dropout) | |
| self.dropout2 = nn.Dropout(dropout) | |
| self.dropout3 = 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, tgt, memory, | |
| tgt_mask: Optional[Tensor] = None, | |
| memory_mask: Optional[Tensor] = None, | |
| tgt_key_padding_mask: Optional[Tensor] = None, | |
| memory_key_padding_mask: Optional[Tensor] = None, | |
| pos: 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.dropout1(tgt2) | |
| tgt = self.norm1(tgt) | |
| tgt2 = 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)[0] | |
| tgt = tgt + self.dropout2(tgt2) | |
| tgt = self.norm2(tgt) | |
| tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) | |
| tgt = tgt + self.dropout3(tgt2) | |
| tgt = self.norm3(tgt) | |
| return tgt | |
| def forward_pre(self, tgt, memory, | |
| tgt_mask: Optional[Tensor] = None, | |
| memory_mask: Optional[Tensor] = None, | |
| tgt_key_padding_mask: Optional[Tensor] = None, | |
| memory_key_padding_mask: Optional[Tensor] = None, | |
| pos: Optional[Tensor] = None, | |
| query_pos: Optional[Tensor] = None): | |
| tgt2 = self.norm1(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.dropout1(tgt2) | |
| tgt2 = self.norm2(tgt) | |
| tgt2 = 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)[0] | |
| tgt = tgt + self.dropout2(tgt2) | |
| tgt2 = self.norm3(tgt) | |
| tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) | |
| tgt = tgt + self.dropout3(tgt2) | |
| return tgt | |
| def forward(self, tgt, memory, | |
| tgt_mask: Optional[Tensor] = None, | |
| memory_mask: Optional[Tensor] = None, | |
| tgt_key_padding_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, tgt_mask, memory_mask, | |
| tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) | |
| return self.forward_post(tgt, memory, tgt_mask, memory_mask, | |
| tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) | |
| def _get_clones(module, N): | |
| return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) | |
| class TransformerEncoder(nn.Module): | |
| def __init__(self, encoder_layer, num_layers, norm=None, num_patches=None, ape=True, embed_dim=None, input_dim=None): | |
| super().__init__() | |
| self.layers = _get_clones(encoder_layer, num_layers) | |
| self.num_layers = num_layers | |
| self.norm = norm | |
| self.embed_dim = embed_dim | |
| if input_dim != embed_dim: | |
| self.proj_x = nn.Conv2d(input_dim, embed_dim, 3, padding=1) | |
| else: | |
| self.proj_x = None | |
| self.embed_proj = nn.Conv2d(1, embed_dim, 1, 1, 0) # learnable | |
| self.ape = ape | |
| if self.ape: | |
| self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim), requires_grad=True) | |
| trunc_normal_(self.absolute_pos_embed, std=.02) | |
| def forward(self, x, | |
| mask: Optional[Tensor] = None, | |
| src_key_padding_mask: Optional[Tensor] = None, | |
| pos: Optional[Tensor] = None, | |
| area_prior = None): | |
| if self.proj_x is not None: | |
| x = self.proj_x(x) | |
| if area_prior is not None: | |
| prior_embed = self.embed_proj(area_prior) | |
| x = x + prior_embed | |
| Wh, Ww = x.size(2), x.size(3) | |
| x = x.flatten(2).transpose(1, 2) | |
| if self.ape: | |
| x = x + self.absolute_pos_embed # this line is later added | |
| for layer in self.layers: | |
| x = layer(x, src_mask=None, src_key_padding_mask=None, pos=None) | |
| if self.norm is not None: | |
| x = self.norm(x) | |
| x = x.view(-1, Wh, Ww, self.embed_dim).permute(0, 3, 1, 2).contiguous() | |
| return x | |
| class TransformerEncoderLayer(nn.Module): | |
| def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.0, | |
| activation="gelu", 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.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) |