import torch import torch.nn as nn import torch.nn.functional as F import torch.nn.utils.spectral_norm as SpectralNorm from torch.autograd import Function class BlurFunctionBackward(Function): @staticmethod def forward(ctx, grad_output, kernel, kernel_flip): ctx.save_for_backward(kernel, kernel_flip) grad_input = F.conv2d(grad_output, kernel_flip, padding=1, groups=grad_output.shape[1]) return grad_input @staticmethod def backward(ctx, gradgrad_output): kernel, kernel_flip = ctx.saved_tensors grad_input = F.conv2d(gradgrad_output, kernel, padding=1, groups=gradgrad_output.shape[1]) return grad_input, None, None class BlurFunction(Function): @staticmethod def forward(ctx, x, kernel, kernel_flip): ctx.save_for_backward(kernel, kernel_flip) output = F.conv2d(x, kernel, padding=1, groups=x.shape[1]) return output @staticmethod def backward(ctx, grad_output): kernel, kernel_flip = ctx.saved_tensors grad_input = BlurFunctionBackward.apply(grad_output, kernel, kernel_flip) return grad_input, None, None blur = BlurFunction.apply class Blur(nn.Module): def __init__(self, channel): super().__init__() kernel = torch.tensor([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=torch.float32) kernel = kernel.view(1, 1, 3, 3) kernel = kernel / kernel.sum() kernel_flip = torch.flip(kernel, [2, 3]) self.kernel = kernel.repeat(channel, 1, 1, 1) self.kernel_flip = kernel_flip.repeat(channel, 1, 1, 1) def forward(self, x): return blur(x, self.kernel.type_as(x), self.kernel_flip.type_as(x)) def calc_mean_std(feat, eps=1e-5): """Calculate mean and std for adaptive_instance_normalization. Args: feat (Tensor): 4D tensor. eps (float): A small value added to the variance to avoid divide-by-zero. Default: 1e-5. """ size = feat.size() assert len(size) == 4, 'The input feature should be 4D tensor.' n, c = size[:2] feat_var = feat.view(n, c, -1).var(dim=2) + eps feat_std = feat_var.sqrt().view(n, c, 1, 1) feat_mean = feat.view(n, c, -1).mean(dim=2).view(n, c, 1, 1) return feat_mean, feat_std def adaptive_instance_normalization(content_feat, style_feat): """Adaptive instance normalization. Adjust the reference features to have the similar color and illuminations as those in the degradate features. Args: content_feat (Tensor): The reference feature. style_feat (Tensor): The degradate features. """ size = content_feat.size() style_mean, style_std = calc_mean_std(style_feat) content_mean, content_std = calc_mean_std(content_feat) normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size) return normalized_feat * style_std.expand(size) + style_mean.expand(size) def AttentionBlock(in_channel): return nn.Sequential( SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1)), nn.LeakyReLU(0.2, True), SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1))) def conv_block(in_channels, out_channels, kernel_size=3, stride=1, dilation=1, bias=True): """Conv block used in MSDilationBlock.""" return nn.Sequential( SpectralNorm( nn.Conv2d( in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) // 2) * dilation, bias=bias)), nn.LeakyReLU(0.2), SpectralNorm( nn.Conv2d( out_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) // 2) * dilation, bias=bias)), ) class MSDilationBlock(nn.Module): """Multi-scale dilation block.""" def __init__(self, in_channels, kernel_size=3, dilation=(1, 1, 1, 1), bias=True): super(MSDilationBlock, self).__init__() self.conv_blocks = nn.ModuleList() for i in range(4): self.conv_blocks.append(conv_block(in_channels, in_channels, kernel_size, dilation=dilation[i], bias=bias)) self.conv_fusion = SpectralNorm( nn.Conv2d( in_channels * 4, in_channels, kernel_size=kernel_size, stride=1, padding=(kernel_size - 1) // 2, bias=bias)) def forward(self, x): out = [] for i in range(4): out.append(self.conv_blocks[i](x)) out = torch.cat(out, 1) out = self.conv_fusion(out) + x return out class UpResBlock(nn.Module): def __init__(self, in_channel): super(UpResBlock, self).__init__() self.body = nn.Sequential( nn.Conv2d(in_channel, in_channel, 3, 1, 1), nn.LeakyReLU(0.2, True), nn.Conv2d(in_channel, in_channel, 3, 1, 1), ) def forward(self, x): out = x + self.body(x) return out