| from torch import nn, Tensor |
| from typing import Union |
|
|
| from .utils import _init_weights |
| from .blocks import BasicBlock, LightBasicBlock, conv1x1, conv3x3 |
|
|
|
|
| class ConvRefine(nn.Module): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| norm_layer: Union[nn.BatchNorm2d, nn.GroupNorm, None] = nn.BatchNorm2d, |
| activation: nn.Module = nn.ReLU(inplace=True), |
| groups: int = 1, |
| ) -> None: |
| super().__init__() |
| self.refine = BasicBlock( |
| in_channels=in_channels, |
| out_channels=out_channels, |
| norm_layer=norm_layer, |
| activation=activation, |
| groups=groups, |
| ) |
| self.apply(_init_weights) |
| |
| def forward(self, x: Tensor) -> Tensor: |
| return self.refine(x) |
|
|
|
|
| class LightConvRefine(nn.Module): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| norm_layer: Union[nn.BatchNorm2d, nn.GroupNorm, None] = nn.BatchNorm2d, |
| activation: nn.Module = nn.ReLU(inplace=True), |
| ) -> None: |
| super().__init__() |
| self.refine = LightBasicBlock( |
| in_channels=in_channels, |
| out_channels=out_channels, |
| norm_layer=norm_layer, |
| activation=activation, |
| ) |
| self.apply(_init_weights) |
| |
| def forward(self, x: Tensor) -> Tensor: |
| return self.refine(x) |
|
|
|
|
| class LighterConvRefine(nn.Module): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| norm_layer: Union[nn.BatchNorm2d, nn.GroupNorm, None] = nn.BatchNorm2d, |
| activation: nn.Module = nn.ReLU(inplace=True), |
| ) -> None: |
| super().__init__() |
| |
| self.conv1 = conv3x3( |
| in_channels=in_channels, |
| out_channels=in_channels, |
| stride=1, |
| groups=in_channels, |
| bias=not norm_layer, |
| ) |
| self.norm1 = norm_layer(in_channels) if norm_layer else nn.Identity() |
| self.act1 = activation |
|
|
| self.conv2 = conv1x1( |
| in_channels=in_channels, |
| out_channels=out_channels, |
| stride=1, |
| bias=not norm_layer, |
| ) |
| self.norm2 = norm_layer(out_channels) if norm_layer else nn.Identity() |
| self.act2 = activation |
|
|
| if in_channels != out_channels: |
| self.downsample = nn.Sequential( |
| conv1x1(in_channels, out_channels, stride=1, bias=not norm_layer), |
| norm_layer(out_channels) if norm_layer else nn.Identity(), |
| ) |
| else: |
| self.downsample = nn.Identity() |
|
|
| self.apply(_init_weights) |
| |
| def forward(self, x: Tensor) -> Tensor: |
| identity = x |
|
|
| out = self.conv1(x) |
| out = self.norm1(out) |
| out = self.act1(out) |
|
|
| out = self.conv2(out) |
| out = self.norm2(out) |
|
|
| out += self.downsample(identity) |
| out = self.act2(out) |
| return out |
|
|
