Residual convolution module.
| 229 | |
| 230 | |
| 231 | class ResidualConvUnit_custom(nn.Module): |
| 232 | """Residual convolution module. |
| 233 | """ |
| 234 | |
| 235 | def __init__(self, features, activation, bn): |
| 236 | """Init. |
| 237 | |
| 238 | Args: |
| 239 | features (int): number of features |
| 240 | """ |
| 241 | super().__init__() |
| 242 | |
| 243 | self.bn = bn |
| 244 | |
| 245 | self.groups=1 |
| 246 | |
| 247 | self.conv1 = nn.Conv2d( |
| 248 | features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups |
| 249 | ) |
| 250 | |
| 251 | self.conv2 = nn.Conv2d( |
| 252 | features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups |
| 253 | ) |
| 254 | |
| 255 | if self.bn==True: |
| 256 | self.bn1 = nn.BatchNorm2d(features) |
| 257 | self.bn2 = nn.BatchNorm2d(features) |
| 258 | |
| 259 | self.activation = activation |
| 260 | |
| 261 | self.skip_add = nn.quantized.FloatFunctional() |
| 262 | |
| 263 | def forward(self, x): |
| 264 | """Forward pass. |
| 265 | |
| 266 | Args: |
| 267 | x (tensor): input |
| 268 | |
| 269 | Returns: |
| 270 | tensor: output |
| 271 | """ |
| 272 | |
| 273 | out = self.activation(x) |
| 274 | out = self.conv1(out) |
| 275 | if self.bn==True: |
| 276 | out = self.bn1(out) |
| 277 | |
| 278 | out = self.activation(out) |
| 279 | out = self.conv2(out) |
| 280 | if self.bn==True: |
| 281 | out = self.bn2(out) |
| 282 | |
| 283 | if self.groups > 1: |
| 284 | out = self.conv_merge(out) |
| 285 | |
| 286 | return self.skip_add.add(out, x) |
| 287 | |
| 288 | # return out + x |