| from torch import nn, Tensor |
| from torch.nn import functional as F |
|
|
| from typing import Union |
| from functools import partial |
|
|
| from .utils import _init_weights |
| from .refine import ConvRefine, LightConvRefine, LighterConvRefine |
|
|
|
|
| class ConvUpsample(nn.Module): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| scale_factor: int = 2, |
| norm_layer: Union[nn.BatchNorm2d, nn.GroupNorm, None] = nn.BatchNorm2d, |
| activation: nn.Module = nn.ReLU(inplace=True), |
| groups: int = 1, |
| ) -> None: |
| super().__init__() |
| assert scale_factor >= 1, f"Scale factor should be greater than or equal to 1, but got {scale_factor}" |
| self.scale_factor = scale_factor |
| self.upsample = partial( |
| F.interpolate, |
| scale_factor=scale_factor, |
| mode="bilinear", |
| align_corners=False, |
| recompute_scale_factor=False, |
| antialias=False, |
| ) if scale_factor > 1 else nn.Identity() |
|
|
| self.refine = ConvRefine( |
| 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: |
| x = self.upsample(x) |
| x = self.refine(x) |
| return x |
|
|
|
|
| class LightConvUpsample(nn.Module): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| scale_factor: int = 2, |
| norm_layer: Union[nn.BatchNorm2d, nn.GroupNorm, None] = nn.BatchNorm2d, |
| activation: nn.Module = nn.ReLU(inplace=True), |
| ) -> None: |
| super().__init__() |
| assert scale_factor >= 1, f"Scale factor should be greater than or equal to 1, but got {scale_factor}" |
| self.scale_factor = scale_factor |
| self.upsample = partial( |
| F.interpolate, |
| scale_factor=scale_factor, |
| mode="bilinear", |
| align_corners=False, |
| recompute_scale_factor=False, |
| antialias=False, |
| ) if scale_factor > 1 else nn.Identity() |
|
|
| self.refine = LightConvRefine( |
| in_channels=in_channels, |
| out_channels=out_channels, |
| norm_layer=norm_layer, |
| activation=activation, |
| ) |
|
|
| self.apply(_init_weights) |
| |
| def forward(self, x: Tensor) -> Tensor: |
| x = self.upsample(x) |
| x = self.refine(x) |
| return x |
|
|
|
|
| class LighterConvUpsample(nn.Module): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| scale_factor: int = 2, |
| norm_layer: Union[nn.BatchNorm2d, nn.GroupNorm, None] = nn.BatchNorm2d, |
| activation: nn.Module = nn.ReLU(inplace=True), |
| ) -> None: |
| super().__init__() |
| assert scale_factor >= 1, f"Scale factor should be greater than or equal to 1, but got {scale_factor}" |
| self.scale_factor = scale_factor |
| self.upsample = partial( |
| F.interpolate, |
| scale_factor=scale_factor, |
| mode="bilinear", |
| align_corners=False, |
| recompute_scale_factor=False, |
| antialias=False, |
| ) if scale_factor > 1 else nn.Identity() |
|
|
| self.refine = LighterConvRefine( |
| in_channels=in_channels, |
| out_channels=out_channels, |
| norm_layer=norm_layer, |
| activation=activation, |
| ) |
|
|
| self.apply(_init_weights) |
| |
| def forward(self, x: Tensor) -> Tensor: |
| x = self.upsample(x) |
| x = self.refine(x) |
| return x |
|
|
