import json from collections import OrderedDict from math import exp from .Common import * # +++++++++++++++++++++++++++++++++++++ # FP16 Training # ------------------------------------- # Modified from Nvidia/Apex # https://github.com/NVIDIA/apex/blob/master/apex/fp16_utils/fp16util.py class tofp16(nn.Module): def __init__(self): super(tofp16, self).__init__() def forward(self, input): if input.is_cuda: return input.half() else: # PyTorch 1.0 doesn't support fp16 in CPU return input.float() def BN_convert_float(module): if isinstance(module, torch.nn.modules.batchnorm._BatchNorm): module.float() for child in module.children(): BN_convert_float(child) return module def network_to_half(network): return nn.Sequential(tofp16(), BN_convert_float(network.half())) # warnings.simplefilter('ignore') # +++++++++++++++++++++++++++++++++++++ # DCSCN # ------------------------------------- class DCSCN(BaseModule): # https://github.com/jiny2001/dcscn-super-resolution def __init__(self, color_channel=3, up_scale=2, feature_layers=12, first_feature_filters=196, last_feature_filters=48, reconstruction_filters=128, up_sampler_filters=32 ): super(DCSCN, self).__init__() self.total_feature_channels = 0 self.total_reconstruct_filters = 0 self.upscale = up_scale self.act_fn = nn.SELU(inplace=False) self.feature_block = self.make_feature_extraction_block(color_channel, feature_layers, first_feature_filters, last_feature_filters) self.reconstruction_block = self.make_reconstruction_block(reconstruction_filters) self.up_sampler = self.make_upsampler(up_sampler_filters, color_channel) self.selu_init_params() def selu_init_params(self): for i in self.modules(): if isinstance(i, nn.Conv2d): i.weight.data.normal_(0.0, 1.0 / sqrt(i.weight.numel())) if i.bias is not None: i.bias.data.fill_(0) def conv_block(self, in_channel, out_channel, kernel_size): m = OrderedDict([ # ("Padding", nn.ReplicationPad2d((kernel_size - 1) // 2)), ('Conv2d', nn.Conv2d(in_channel, out_channel, kernel_size=kernel_size, padding=(kernel_size - 1) // 2)), ('Activation', self.act_fn) ]) return nn.Sequential(m) def make_feature_extraction_block(self, color_channel, num_layers, first_filters, last_filters): # input layer feature_block = [("Feature 1", self.conv_block(color_channel, first_filters, 3))] # exponential decay # rest layers alpha_rate = log(first_filters / last_filters) / (num_layers - 1) filter_nums = [round(first_filters * exp(-alpha_rate * i)) for i in range(num_layers)] self.total_feature_channels = sum(filter_nums) layer_filters = [[filter_nums[i], filter_nums[i + 1], 3] for i in range(num_layers - 1)] feature_block.extend([("Feature {}".format(index + 2), self.conv_block(*x)) for index, x in enumerate(layer_filters)]) return nn.Sequential(OrderedDict(feature_block)) def make_reconstruction_block(self, num_filters): B1 = self.conv_block(self.total_feature_channels, num_filters // 2, 1) B2 = self.conv_block(num_filters // 2, num_filters, 3) m = OrderedDict([ ("A", self.conv_block(self.total_feature_channels, num_filters, 1)), ("B", nn.Sequential(*[B1, B2])) ]) self.total_reconstruct_filters = num_filters * 2 return nn.Sequential(m) def make_upsampler(self, out_channel, color_channel): out = out_channel * self.upscale ** 2 m = OrderedDict([ ('Conv2d_block', self.conv_block(self.total_reconstruct_filters, out, kernel_size=3)), ('PixelShuffle', nn.PixelShuffle(self.upscale)), ("Conv2d", nn.Conv2d(out_channel, color_channel, kernel_size=3, padding=1, bias=False)) ]) return nn.Sequential(m) def forward(self, x): # residual learning lr, lr_up = x feature = [] for layer in self.feature_block.children(): lr = layer(lr) feature.append(lr) feature = torch.cat(feature, dim=1) reconstruction = [layer(feature) for layer in self.reconstruction_block.children()] reconstruction = torch.cat(reconstruction, dim=1) lr = self.up_sampler(reconstruction) return lr + lr_up # +++++++++++++++++++++++++++++++++++++ # CARN # ------------------------------------- class CARN_Block(BaseModule): def __init__(self, channels, kernel_size=3, padding=1, dilation=1, groups=1, activation=nn.SELU(), repeat=3, SEBlock=False, conv=nn.Conv2d, single_conv_size=1, single_conv_group=1): super(CARN_Block, self).__init__() m = [] for i in range(repeat): m.append(ResidualFixBlock(channels, channels, kernel_size=kernel_size, padding=padding, dilation=dilation, groups=groups, activation=activation, conv=conv)) if SEBlock: m.append(SpatialChannelSqueezeExcitation(channels, reduction=channels)) self.blocks = nn.Sequential(*m) self.singles = nn.Sequential( *[ConvBlock(channels * (i + 2), channels, kernel_size=single_conv_size, padding=(single_conv_size - 1) // 2, groups=single_conv_group, activation=activation, conv=conv) for i in range(repeat)]) def forward(self, x): c0 = x for block, single in zip(self.blocks, self.singles): b = block(x) c0 = c = torch.cat([c0, b], dim=1) x = single(c) return x class CARN(BaseModule): # Fast, Accurate, and Lightweight Super-Resolution with Cascading Residual Network # https://github.com/nmhkahn/CARN-pytorch def __init__(self, color_channels=3, mid_channels=64, scale=2, activation=nn.SELU(), num_blocks=3, conv=nn.Conv2d): super(CARN, self).__init__() self.color_channels = color_channels self.mid_channels = mid_channels self.scale = scale self.entry_block = ConvBlock(color_channels, mid_channels, kernel_size=3, padding=1, activation=activation, conv=conv) self.blocks = nn.Sequential( *[CARN_Block(mid_channels, kernel_size=3, padding=1, activation=activation, conv=conv, single_conv_size=1, single_conv_group=1) for _ in range(num_blocks)]) self.singles = nn.Sequential( *[ConvBlock(mid_channels * (i + 2), mid_channels, kernel_size=1, padding=0, activation=activation, conv=conv) for i in range(num_blocks)]) self.upsampler = UpSampleBlock(mid_channels, scale=scale, activation=activation, conv=conv) self.exit_conv = conv(mid_channels, color_channels, kernel_size=3, padding=1) def forward(self, x): x = self.entry_block(x) c0 = x for block, single in zip(self.blocks, self.singles): b = block(x) c0 = c = torch.cat([c0, b], dim=1) x = single(c) x = self.upsampler(x) out = self.exit_conv(x) return out class CARN_V2(CARN): def __init__(self, color_channels=3, mid_channels=64, scale=2, activation=nn.LeakyReLU(0.1), SEBlock=True, conv=nn.Conv2d, atrous=(1, 1, 1), repeat_blocks=3, single_conv_size=3, single_conv_group=1): super(CARN_V2, self).__init__(color_channels=color_channels, mid_channels=mid_channels, scale=scale, activation=activation, conv=conv) num_blocks = len(atrous) m = [] for i in range(num_blocks): m.append(CARN_Block(mid_channels, kernel_size=3, padding=1, dilation=1, activation=activation, SEBlock=SEBlock, conv=conv, repeat=repeat_blocks, single_conv_size=single_conv_size, single_conv_group=single_conv_group)) self.blocks = nn.Sequential(*m) self.singles = nn.Sequential( *[ConvBlock(mid_channels * (i + 2), mid_channels, kernel_size=single_conv_size, padding=(single_conv_size - 1) // 2, groups=single_conv_group, activation=activation, conv=conv) for i in range(num_blocks)]) def forward(self, x): x = self.entry_block(x) c0 = x res = x for block, single in zip(self.blocks, self.singles): b = block(x) c0 = c = torch.cat([c0, b], dim=1) x = single(c) x = x + res x = self.upsampler(x) out = self.exit_conv(x) return out # +++++++++++++++++++++++++++++++++++++ # original Waifu2x model # ------------------------------------- class UpConv_7(BaseModule): # https://github.com/nagadomi/waifu2x/blob/3c46906cb78895dbd5a25c3705994a1b2e873199/lib/srcnn.lua#L311 def __init__(self): super(UpConv_7, self).__init__() self.act_fn = nn.LeakyReLU(0.1, inplace=False) self.offset = 7 # because of 0 padding from torch.nn import ZeroPad2d self.pad = ZeroPad2d(self.offset) m = [nn.Conv2d(3, 16, 3, 1, 0), self.act_fn, nn.Conv2d(16, 32, 3, 1, 0), self.act_fn, nn.Conv2d(32, 64, 3, 1, 0), self.act_fn, nn.Conv2d(64, 128, 3, 1, 0), self.act_fn, nn.Conv2d(128, 128, 3, 1, 0), self.act_fn, nn.Conv2d(128, 256, 3, 1, 0), self.act_fn, # in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding= nn.ConvTranspose2d(256, 3, kernel_size=4, stride=2, padding=3, bias=False) ] self.Sequential = nn.Sequential(*m) def load_pre_train_weights(self, json_file): with open(json_file) as f: weights = json.load(f) box = [] for i in weights: box.append(i['weight']) box.append(i['bias']) own_state = self.state_dict() for index, (name, param) in enumerate(own_state.items()): own_state[name].copy_(torch.FloatTensor(box[index])) def forward(self, x): x = self.pad(x) return self.Sequential.forward(x) class Vgg_7(UpConv_7): def __init__(self): super(Vgg_7, self).__init__() self.act_fn = nn.LeakyReLU(0.1, inplace=False) self.offset = 7 m = [nn.Conv2d(3, 32, 3, 1, 0), self.act_fn, nn.Conv2d(32, 32, 3, 1, 0), self.act_fn, nn.Conv2d(32, 64, 3, 1, 0), self.act_fn, nn.Conv2d(64, 64, 3, 1, 0), self.act_fn, nn.Conv2d(64, 128, 3, 1, 0), self.act_fn, nn.Conv2d(128, 128, 3, 1, 0), self.act_fn, nn.Conv2d(128, 3, 3, 1, 0) ] self.Sequential = nn.Sequential(*m)