| 260 | return output |
| 261 | |
| 262 | class FeedForwardNetwork(nn.Module): |
| 263 | def __init__( |
| 264 | self, |
| 265 | embed_dim, |
| 266 | ffn_dim, |
| 267 | activation_fn=F.gelu, |
| 268 | dropout=0.0, |
| 269 | activation_dropout=0.0, |
| 270 | layernorm_eps=1e-6, |
| 271 | subln=False, |
| 272 | subconv=False |
| 273 | ): |
| 274 | super().__init__() |
| 275 | self.embed_dim = embed_dim |
| 276 | # self.out_dim = out_dim |
| 277 | self.activation_fn = activation_fn |
| 278 | self.activation_dropout_module = torch.nn.Dropout(activation_dropout) |
| 279 | self.dropout_module = torch.nn.Dropout(dropout) |
| 280 | self.fc1 = nn.Linear(self.embed_dim, ffn_dim) |
| 281 | self.fc2 = nn.Linear(ffn_dim, self.embed_dim) |
| 282 | self.ffn_layernorm = nn.LayerNorm(ffn_dim, eps=layernorm_eps) if subln else None |
| 283 | self.dwconv = DWConv2d(ffn_dim, 3, 1, 1) if subconv else None |
| 284 | |
| 285 | def reset_parameters(self): |
| 286 | self.fc1.reset_parameters() |
| 287 | self.fc2.reset_parameters() |
| 288 | if self.ffn_layernorm is not None: |
| 289 | self.ffn_layernorm.reset_parameters() |
| 290 | |
| 291 | def forward(self, x: torch.Tensor): |
| 292 | ''' |
| 293 | x: (b h w c) |
| 294 | ''' |
| 295 | x = self.fc1(x) |
| 296 | x = self.activation_fn(x) |
| 297 | x = self.activation_dropout_module(x) |
| 298 | if self.dwconv is not None: |
| 299 | residual = x |
| 300 | x = self.dwconv(x) |
| 301 | x = x + residual |
| 302 | if self.ffn_layernorm is not None: |
| 303 | x = self.ffn_layernorm(x) |
| 304 | x = self.fc2(x) |
| 305 | x = self.dropout_module(x) |
| 306 | return x |
| 307 | |
| 308 | class FeedForward(nn.Module): |
| 309 | def __init__(self, in_dim, hidden_dim, out_chans=None, act_layer=nn.GELU, dropout=0.): |