Spaces:
Runtime error
Runtime error
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| from typing import Optional | |
| import torch | |
| from torch import nn, Tensor | |
| from torch.nn import functional as F | |
| from timm.models.layers import trunc_normal_ | |
| from detectron2.layers import Conv2d | |
| import fvcore.nn.weight_init as weight_init | |
| from ..utils import MultiheadAttention | |
| class SelfAttentionLayer(nn.Module): | |
| def __init__(self, d_model, nhead, dropout=0.0, | |
| activation="relu", normalize_before=False): | |
| super().__init__() | |
| self.self_attn = 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__() | |
| # 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.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) | |
| 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 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 | |