| 247 | |
| 248 | class FeedForwardNetwork(nn.Module): |
| 249 | def __init__( |
| 250 | self, |
| 251 | embed_dim, |
| 252 | ffn_dim, |
| 253 | activation_fn=F.gelu, |
| 254 | dropout=0.0, |
| 255 | activation_dropout=0.0, |
| 256 | layernorm_eps=1e-6, |
| 257 | subln=False, |
| 258 | subconv=False |
| 259 | ): |
| 260 | super().__init__() |
| 261 | self.embed_dim = embed_dim |
| 262 | self.activation_fn = activation_fn |
| 263 | self.activation_dropout_module = torch.nn.Dropout(activation_dropout) |
| 264 | self.dropout_module = torch.nn.Dropout(dropout) |
| 265 | self.fc1 = nn.Linear(self.embed_dim, ffn_dim) |
| 266 | self.fc2 = nn.Linear(ffn_dim, self.embed_dim) |
| 267 | self.ffn_layernorm = nn.LayerNorm(ffn_dim, eps=layernorm_eps) if subln else None |
| 268 | self.dwconv = DWConv2d(ffn_dim, 3, 1, 1) if subconv else None |
| 269 | |
| 270 | def reset_parameters(self): |
| 271 | self.fc1.reset_parameters() |