| 38 | """ |
| 39 | |
| 40 | def __init__(self, |
| 41 | embed_dims, |
| 42 | feedforward_channels, |
| 43 | act_cfg=dict(type='GELU'), |
| 44 | ffn_drop=0., |
| 45 | dropout_layer=None, |
| 46 | init_cfg=None): |
| 47 | super(MixFFN, self).__init__(init_cfg) |
| 48 | |
| 49 | self.embed_dims = embed_dims |
| 50 | self.feedforward_channels = feedforward_channels |
| 51 | self.act_cfg = act_cfg |
| 52 | self.activate = build_activation_layer(act_cfg) |
| 53 | |
| 54 | in_channels = embed_dims |
| 55 | fc1 = Conv2d( |
| 56 | in_channels=in_channels, |
| 57 | out_channels=feedforward_channels, |
| 58 | kernel_size=1, |
| 59 | stride=1, |
| 60 | bias=True) |
| 61 | # 3x3 depth wise conv to provide positional encode information |
| 62 | pe_conv = Conv2d( |
| 63 | in_channels=feedforward_channels, |
| 64 | out_channels=feedforward_channels, |
| 65 | kernel_size=3, |
| 66 | stride=1, |
| 67 | padding=(3 - 1) // 2, |
| 68 | bias=True, |
| 69 | groups=feedforward_channels) |
| 70 | fc2 = Conv2d( |
| 71 | in_channels=feedforward_channels, |
| 72 | out_channels=in_channels, |
| 73 | kernel_size=1, |
| 74 | stride=1, |
| 75 | bias=True) |
| 76 | drop = nn.Dropout(ffn_drop) |
| 77 | layers = [fc1, pe_conv, self.activate, drop, fc2, drop] |
| 78 | self.layers = Sequential(*layers) |
| 79 | self.dropout_layer = build_dropout( |
| 80 | dropout_layer) if dropout_layer else torch.nn.Identity() |
| 81 | |
| 82 | def forward(self, x, hw_shape, identity=None): |
| 83 | out = nlc_to_nchw(x, hw_shape) |