def __init__(self, num_blocks=None, num_channel=64, block=ResBlock): super().__init__() if num_blocks is None: num_blocks = [5, 10, 15] assert len(num_blocks) == 3 self.init = nn.Sequential( MeanShift(sign=-1), ConvolutionBlock(in_channels=3, out_channels=num_channel)) self.model_1 = nn.Sequential( *tuple([ block(in_channels=num_channel) for _ in range(num_blocks[0]) ]), ConvolutionDownscale(channels=num_channel), ConvolutionDownscale(channels=num_channel)) self.model_2 = nn.Sequential( *tuple([ block(in_channels=num_channel) for _ in range(num_blocks[1]) ]), ConvolutionDownscale(channels=num_channel)) self.model_3 = nn.Sequential(*tuple( [block(in_channels=num_channel) for _ in range(num_blocks[2])])) self.downscale_1 = ConvolutionDownscale(channels=num_channel) self.downscale_2 = ConvolutionDownscale(channels=num_channel) self.end = nn.Sequential( ConvolutionBlock(in_channels=num_channel * 3, out_channels=3), MeanShift(sign=1))
def __init__(self, block_sequence=None, num_channel=256, rgb_mean=(0.4488, 0.4371, 0.4040), rgb_std=(1.0, 1.0, 1.0), rgb_range=255): super().__init__() if block_sequence is None: block_sequence = [15, 10, 5] self.init = nn.Sequential( MeanShift(rgb_range, rgb_mean, rgb_std, sign=-1), ConvolutionBlock(in_channels=3, out_channels=num_channel) ) self.skip_upscale = nn.Sequential( TransposeUpscale(channels=num_channel), ConvolutionBlock(in_channels=num_channel, out_channels=3) ) model_body = [] for seq in block_sequence: model_body.extend([ResBlock(in_channels=num_channel) for _ in range(seq)]) model_body.append(PixelShuffleUpscale(channels=num_channel)) model_body.append(ConvolutionBlock(in_channels=num_channel, out_channels=3)) self.blocks = nn.Sequential(*tuple(model_body)) self.model_end = nn.Sequential( ConvolutionBlock(in_channels=num_channel, out_channels=3), MeanShift(rgb_range, rgb_mean, rgb_std, sign=1) )
def __init__(self, scale, num_blocks=32, num_channel=256): super().__init__() self.model_0 = nn.Sequential( MeanShift(sign=-1), ConvolutionBlock(in_channels=3, out_channels=num_channel) ) self.model_1 = nn.Sequential(*tuple([ResBlock(num_channel) for _ in range(num_blocks)])) self.model_2 = nn.Sequential( *tuple([PixelShuffleUpscale(channels=num_channel) for _ in range(int(log2(scale)))]), ConvolutionBlock(in_channels=num_channel, out_channels=3), MeanShift(sign=1) )
def __init__(self, scale=4, num_groups=12, num_channel=128, rep_pad=False): super().__init__() self.model_0 = nn.Sequential( MeanShift(sign=-1), ConvolutionBlock(in_channels=3, out_channels=num_channel, rep_pad=rep_pad) ) self.model_1 = nn.Sequential(*tuple([BasicGroup(num_channel, rep_pad=rep_pad) for _ in range(num_groups)])) upscaler = [PixelShuffleUpscale(channels=num_channel, basic_block=DepthSeparableConvBlock, rep_pad=rep_pad) for _ in range(int(log2(scale)))] self.model_2 = nn.Sequential( *tuple(upscaler), ConvolutionBlock(in_channels=num_channel, out_channels=3, rep_pad=rep_pad), MeanShift(sign=1) )
def __init__(self, scale=4, num_groups=20, num_channel=64): super().__init__() self.model_0 = nn.Sequential( MeanShift(sign=-1), ConvolutionBlock(in_channels=3, out_channels=num_channel)) self.model_1 = nn.Sequential( *tuple([BasicGroup(num_channel) for _ in range(num_groups)])) upscaler = [ PixelShuffleUpscale(channels=num_channel) for _ in range(int(log2(scale))) ] self.model_2 = nn.Sequential( *tuple(upscaler), ConvolutionBlock(in_channels=num_channel, out_channels=3), MeanShift(sign=1))
def __init__(self, num_groups=None, num_channel=128): super().__init__() if num_groups is None: num_groups = [6, 3] self.model_0 = nn.Sequential( MeanShift(sign=-1), ConvolutionBlock(in_channels=3, out_channels=num_channel) ) self.model_1 = nn.Sequential(*tuple([BasicGroup(num_channel) for _ in range(num_groups[0])])) self.upscale_1 = PixelShuffleUpscale(channels=num_channel) self.model_2 = nn.Sequential(*tuple([BasicGroup(num_channel) for _ in range(num_groups[1])])) self.upscale_2 = PixelShuffleUpscale(channels=num_channel) self.model_3 = nn.Sequential( ConvolutionBlock(in_channels=num_channel, out_channels=3), MeanShift(sign=1) )