| 151 | return x |
| 152 | |
| 153 | class BlockMod(nn.Module): |
| 154 | def __init__(self, channels: int, dw_expand: float = 1., ffn_expand: int = 2, drop_out_rate: float = 0., |
| 155 | attention_type: Literal['CA', 'SCA'] = 'CA', |
| 156 | activation_type: Literal['GELU', 'SG', 'Identity', 'ReLU', 'LReLU', 'Tanh', 'Sigmoid'] = 'GELU', |
| 157 | inverted_conv: bool = True, kernel_size: int = 3, |
| 158 | use_pos_map: bool = False, depth: int = None): |
| 159 | """ |
| 160 | Modified block from NAFNet with optional CoordConv on th first convolution. |
| 161 | Args: |
| 162 | channels: Number of input channels |
| 163 | dw_expand: Expansion factor for sub block 1 |
| 164 | ffn_expand: Expansion factor for sub block 2 |
| 165 | drop_out_rate: Dropout rate |
| 166 | attention_type: Type of attention to use |
| 167 | activation_type: Type of activation to use |
| 168 | inverted_conv: Whether to use inverted convolution |
| 169 | kernel_size: size of all convolution kernels |
| 170 | use_pos_map: Whether coordconv is used |
| 171 | depth: Depth of the block in the network, used for positional map scaling |
| 172 | """ |
| 173 | super().__init__() |
| 174 | dw_channel = int(channels * dw_expand) if int(channels * dw_expand) % 2 == 0 else int(channels * dw_expand) + 1 |
| 175 | |
| 176 | self.cat = ConcatTensors() if use_pos_map else IdentityMod() |
| 177 | |
| 178 | if inverted_conv: |
| 179 | self.conv1 = InvertedConvolution(in_channels=channels + 2 if use_pos_map else channels, |
| 180 | out_channels=dw_channel, |
| 181 | kernel_size=kernel_size, padding='same', bias=True) |
| 182 | else: |
| 183 | self.conv1 = nn.Conv2d(in_channels=channels + 2 if use_pos_map else channels, out_channels=dw_channel, |
| 184 | kernel_size=kernel_size, padding='same', stride=1, bias=True) |
| 185 | |
| 186 | # Activation |
| 187 | self.activation = get_activation(activation_type) |
| 188 | |
| 189 | # Channel Attention |
| 190 | self.attention = get_cnn_attention(attention_type)( |
| 191 | dw_channel // 2 if activation_type == 'SG' else dw_channel) |
| 192 | |
| 193 | self.conv2 = nn.Conv2d(in_channels=dw_channel // 2 if activation_type == 'SG' else dw_channel, |
| 194 | out_channels=channels, kernel_size=1, padding=0, stride=1, groups=1, bias=True) |
| 195 | |
| 196 | |
| 197 | # sub block 1 done |
| 198 | ffn_channel = floor(ffn_expand * channels) |
| 199 | self.conv3 = nn.Conv2d(in_channels=channels, out_channels=ffn_channel, kernel_size=1, padding=0, stride=1, |
| 200 | groups=1, bias=True) |
| 201 | # second activation call |
| 202 | self.conv4 = nn.Conv2d(in_channels=ffn_channel // 2 if activation_type == 'SG' else ffn_channel, |
| 203 | out_channels=channels, kernel_size=1, padding=0, stride=1, groups=1, bias=True) |
| 204 | |
| 205 | # sub block 2 done |
| 206 | |
| 207 | self.norm1 = LayerNorm2d(channels) |
| 208 | self.norm2 = LayerNorm2d(channels) |
| 209 | |
| 210 | self.dropout1 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity() |