| | |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| |
|
| | from ..utils import xavier_init |
| | from .registry import UPSAMPLE_LAYERS |
| |
|
| | UPSAMPLE_LAYERS.register_module('nearest', module=nn.Upsample) |
| | UPSAMPLE_LAYERS.register_module('bilinear', module=nn.Upsample) |
| |
|
| |
|
| | @UPSAMPLE_LAYERS.register_module(name='pixel_shuffle') |
| | class PixelShufflePack(nn.Module): |
| | """Pixel Shuffle upsample layer. |
| | |
| | This module packs `F.pixel_shuffle()` and a nn.Conv2d module together to |
| | achieve a simple upsampling with pixel shuffle. |
| | |
| | Args: |
| | in_channels (int): Number of input channels. |
| | out_channels (int): Number of output channels. |
| | scale_factor (int): Upsample ratio. |
| | upsample_kernel (int): Kernel size of the conv layer to expand the |
| | channels. |
| | """ |
| |
|
| | def __init__(self, in_channels, out_channels, scale_factor, |
| | upsample_kernel): |
| | super(PixelShufflePack, self).__init__() |
| | self.in_channels = in_channels |
| | self.out_channels = out_channels |
| | self.scale_factor = scale_factor |
| | self.upsample_kernel = upsample_kernel |
| | self.upsample_conv = nn.Conv2d( |
| | self.in_channels, |
| | self.out_channels * scale_factor * scale_factor, |
| | self.upsample_kernel, |
| | padding=(self.upsample_kernel - 1) // 2) |
| | self.init_weights() |
| |
|
| | def init_weights(self): |
| | xavier_init(self.upsample_conv, distribution='uniform') |
| |
|
| | def forward(self, x): |
| | x = self.upsample_conv(x) |
| | x = F.pixel_shuffle(x, self.scale_factor) |
| | return x |
| |
|
| |
|
| | def build_upsample_layer(cfg, *args, **kwargs): |
| | """Build upsample layer. |
| | |
| | Args: |
| | cfg (dict): The upsample layer config, which should contain: |
| | |
| | - type (str): Layer type. |
| | - scale_factor (int): Upsample ratio, which is not applicable to |
| | deconv. |
| | - layer args: Args needed to instantiate a upsample layer. |
| | args (argument list): Arguments passed to the ``__init__`` |
| | method of the corresponding conv layer. |
| | kwargs (keyword arguments): Keyword arguments passed to the |
| | ``__init__`` method of the corresponding conv layer. |
| | |
| | Returns: |
| | nn.Module: Created upsample layer. |
| | """ |
| | if not isinstance(cfg, dict): |
| | raise TypeError(f'cfg must be a dict, but got {type(cfg)}') |
| | if 'type' not in cfg: |
| | raise KeyError( |
| | f'the cfg dict must contain the key "type", but got {cfg}') |
| | cfg_ = cfg.copy() |
| |
|
| | layer_type = cfg_.pop('type') |
| | if layer_type not in UPSAMPLE_LAYERS: |
| | raise KeyError(f'Unrecognized upsample type {layer_type}') |
| | else: |
| | upsample = UPSAMPLE_LAYERS.get(layer_type) |
| |
|
| | if upsample is nn.Upsample: |
| | cfg_['mode'] = layer_type |
| | layer = upsample(*args, **kwargs, **cfg_) |
| | return layer |
| |
|
