| 152 | |
| 153 | |
| 154 | class TransformerEncoderLayer(nn.Module): |
| 155 | def __init__( |
| 156 | self, |
| 157 | d_model, |
| 158 | nhead, |
| 159 | dim_feedforward=2048, |
| 160 | dropout=0.1, |
| 161 | activation="relu", |
| 162 | normalize_before=False, |
| 163 | ): |
| 164 | super().__init__() |
| 165 | self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) |
| 166 | # Implementation of Feedforward model |
| 167 | self.linear1 = nn.Linear(d_model, dim_feedforward) |
| 168 | self.dropout = nn.Dropout(dropout) |
| 169 | self.linear2 = nn.Linear(dim_feedforward, d_model) |
| 170 | |
| 171 | self.norm1 = nn.LayerNorm(d_model) |
| 172 | self.norm2 = nn.LayerNorm(d_model) |
| 173 | self.dropout1 = nn.Dropout(dropout) |
| 174 | self.dropout2 = nn.Dropout(dropout) |
| 175 | |
| 176 | self.activation = _get_activation_fn(activation) |
| 177 | self.normalize_before = normalize_before |
| 178 | |
| 179 | def with_pos_embed(self, tensor, pos: Optional[Tensor]): |
| 180 | return tensor if pos is None else tensor + pos |
| 181 | |
| 182 | def forward_post( |
| 183 | self, |
| 184 | src, |
| 185 | src_mask: Optional[Tensor] = None, |
| 186 | src_key_padding_mask: Optional[Tensor] = None, |
| 187 | pos: Optional[Tensor] = None, |
| 188 | ): |
| 189 | q = k = self.with_pos_embed(src, pos) |
| 190 | |
| 191 | src2 = self.self_attn( |
| 192 | q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask |
| 193 | )[0] |
| 194 | src = src + self.dropout1(src2) |
| 195 | src = self.norm1(src) |
| 196 | src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) |
| 197 | src = src + self.dropout2(src2) |
| 198 | src = self.norm2(src) |
| 199 | return src |
| 200 | |
| 201 | def forward_pre( |
| 202 | self, |
| 203 | src, |
| 204 | src_mask: Optional[Tensor] = None, |
| 205 | src_key_padding_mask: Optional[Tensor] = None, |
| 206 | pos: Optional[Tensor] = None, |
| 207 | ): |
| 208 | src2 = self.norm1(src) |
| 209 | q = k = self.with_pos_embed(src2, pos) |
| 210 | src2 = self.self_attn( |
| 211 | q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask |