r""" Stem layer of InternImage Args: in_channels (int): number of input channels out_channels (int): number of output channels act_layer (str): activation layer norm_layer (str): normalization layer
| 3 | |
| 4 | |
| 5 | class StemLayer(nn.Module): |
| 6 | r""" Stem layer of InternImage |
| 7 | Args: |
| 8 | in_channels (int): number of input channels |
| 9 | out_channels (int): number of output channels |
| 10 | act_layer (str): activation layer |
| 11 | norm_layer (str): normalization layer |
| 12 | """ |
| 13 | |
| 14 | def __init__(self, |
| 15 | in_channels=3+1, |
| 16 | inter_channels=48, |
| 17 | out_channels=96, |
| 18 | act_layer='GELU', |
| 19 | norm_layer='BN'): |
| 20 | super().__init__() |
| 21 | self.conv1 = nn.Conv2d(in_channels, |
| 22 | inter_channels, |
| 23 | kernel_size=3, |
| 24 | stride=1, |
| 25 | padding=1) |
| 26 | self.norm1 = build_norm_layer( |
| 27 | inter_channels, norm_layer, 'channels_first', 'channels_first' |
| 28 | ) |
| 29 | self.act = build_act_layer(act_layer) |
| 30 | self.conv2 = nn.Conv2d(inter_channels, |
| 31 | out_channels, |
| 32 | kernel_size=3, |
| 33 | stride=1, |
| 34 | padding=1) |
| 35 | self.norm2 = build_norm_layer( |
| 36 | out_channels, norm_layer, 'channels_first', 'channels_first' |
| 37 | ) |
| 38 | |
| 39 | def forward(self, x): |
| 40 | x = self.conv1(x) |
| 41 | x = self.norm1(x) |
| 42 | x = self.act(x) |
| 43 | x = self.conv2(x) |
| 44 | x = self.norm2(x) |
| 45 | return x |