| 15 | |
| 16 | |
| 17 | class DeformableTransformerEncoderLayer(nn.Module): |
| 18 | def __init__( |
| 19 | self, |
| 20 | d_model=256, |
| 21 | d_ffn=1024, |
| 22 | dropout=0.1, |
| 23 | activation='relu', |
| 24 | n_levels=4, |
| 25 | n_heads=8, |
| 26 | n_points=4, |
| 27 | ): |
| 28 | super().__init__() |
| 29 | # pdb.set_trace() |
| 30 | # self attention |
| 31 | self.self_attn = MSDeformAttn(d_model, n_levels, n_heads, |
| 32 | n_points) # 256 4 8 4 |
| 33 | self.dropout1 = nn.Dropout(dropout) |
| 34 | self.norm1 = nn.LayerNorm(d_model) |
| 35 | |
| 36 | # ffn |
| 37 | self.linear1 = nn.Linear(d_model, d_ffn) |
| 38 | self.activation = _get_activation_fn(activation, d_model=d_ffn) |
| 39 | self.dropout2 = nn.Dropout(dropout) |
| 40 | self.linear2 = nn.Linear(d_ffn, d_model) |
| 41 | self.dropout3 = nn.Dropout(dropout) |
| 42 | self.norm2 = nn.LayerNorm(d_model) |
| 43 | |
| 44 | @staticmethod |
| 45 | def with_pos_embed(tensor, pos): |
| 46 | return tensor if pos is None else tensor + pos |
| 47 | |
| 48 | def forward_ffn(self, src): |
| 49 | src2 = self.linear2(self.dropout2(self.activation(self.linear1(src)))) |
| 50 | src = src + self.dropout3(src2) |
| 51 | src = self.norm2(src) |
| 52 | return src |
| 53 | |
| 54 | def forward(self, |
| 55 | src, |
| 56 | pos, |
| 57 | reference_points, |
| 58 | spatial_shapes, |
| 59 | level_start_index, |
| 60 | key_padding_mask=None): |
| 61 | # pdb.set_trace() |
| 62 | src2 = self.self_attn(self.with_pos_embed(src, pos), reference_points, |
| 63 | src, spatial_shapes, level_start_index, |
| 64 | key_padding_mask) |
| 65 | src = src + self.dropout1(src2) |
| 66 | src = self.norm1(src) |
| 67 | src = self.forward_ffn(src) |
| 68 | |
| 69 | return src |
| 70 | |
| 71 | |
| 72 | class DeformableTransformerDecoderLayer(nn.Module): |