Residual convolution module.
| 320 | |
| 321 | |
| 322 | class ResidualConvUnit_custom(nn.Module): |
| 323 | """Residual convolution module. |
| 324 | """ |
| 325 | |
| 326 | def __init__(self, features, activation, bn): |
| 327 | """Init. |
| 328 | |
| 329 | Args: |
| 330 | features (int): number of features |
| 331 | """ |
| 332 | super().__init__() |
| 333 | |
| 334 | self.bn = bn |
| 335 | |
| 336 | self.groups=1 |
| 337 | |
| 338 | self.conv1 = nn.Conv2d( |
| 339 | features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups |
| 340 | ) |
| 341 | |
| 342 | self.conv2 = nn.Conv2d( |
| 343 | features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups |
| 344 | ) |
| 345 | |
| 346 | if self.bn==True: |
| 347 | self.bn1 = nn.BatchNorm2d(features) |
| 348 | self.bn2 = nn.BatchNorm2d(features) |
| 349 | |
| 350 | self.activation = activation |
| 351 | |
| 352 | self.skip_add = nn.quantized.FloatFunctional() |
| 353 | |
| 354 | def forward(self, x): |
| 355 | """Forward pass. |
| 356 | |
| 357 | Args: |
| 358 | x (tensor): input |
| 359 | |
| 360 | Returns: |
| 361 | tensor: output |
| 362 | """ |
| 363 | |
| 364 | out = self.activation(x) |
| 365 | out = self.conv1(out) |
| 366 | if self.bn==True: |
| 367 | out = self.bn1(out) |
| 368 | |
| 369 | out = self.activation(out) |
| 370 | out = self.conv2(out) |
| 371 | if self.bn==True: |
| 372 | out = self.bn2(out) |
| 373 | |
| 374 | if self.groups > 1: |
| 375 | out = self.conv_merge(out) |
| 376 | |
| 377 | return self.skip_add.add(out, x) |
| 378 | |
| 379 | # return out + x |