An implementation of MixFFN of Segformer. The differences between MixFFN & FFN: 1. Use 1X1 Conv to replace Linear layer. 2. Introduce 3X3 Conv to encode positional information. Args: embed_dims (int): The feature dimension. Same as `MultiheadAttention`.
| 16 | |
| 17 | |
| 18 | class MixFFN(BaseModule): |
| 19 | """An implementation of MixFFN of Segformer. |
| 20 | |
| 21 | The differences between MixFFN & FFN: |
| 22 | 1. Use 1X1 Conv to replace Linear layer. |
| 23 | 2. Introduce 3X3 Conv to encode positional information. |
| 24 | |
| 25 | Args: |
| 26 | embed_dims (int): The feature dimension. Same as |
| 27 | `MultiheadAttention`. Defaults: 256. |
| 28 | feedforward_channels (int): The hidden dimension of FFNs. |
| 29 | Defaults: 1024. |
| 30 | act_cfg (dict, optional): The activation config for FFNs. |
| 31 | Default: dict(type='ReLU') |
| 32 | ffn_drop (float, optional): Probability of an element to be |
| 33 | zeroed in FFN. Default 0.0. |
| 34 | dropout_layer (obj:`ConfigDict`): The dropout_layer used |
| 35 | when adding the shortcut. |
| 36 | init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization. |
| 37 | Default: None. |
| 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) |