import torch import torch.nn as nn # from basicsr.models.archs import recons_video81 as recons_video # from basicsr.models.archs import flow_pwc82 as flow_pwc import numpy as np from torch.nn import functional as F import torch.utils.checkpoint as checkpoint from torch.cuda.amp import autocast as autocast # from .StudentImage_arch import StudentImage from torch.nn.parallel import DistributedDataParallel from collections import OrderedDict from .arch_util import DCNv2Pack from .common import ResList class LayerNormFunction(torch.autograd.Function): @staticmethod def forward(ctx, x, weight, bias, eps): ctx.eps = eps N, C, H, W = x.size() mu = x.mean(1, keepdim=True) var = (x - mu).pow(2).mean(1, keepdim=True) y = (x - mu) / (var + eps).sqrt() ctx.save_for_backward(y, var, weight) y = weight.view(1, C, 1, 1) * y + bias.view(1, C, 1, 1) return y @staticmethod def backward(ctx, grad_output): eps = ctx.eps N, C, H, W = grad_output.size() y, var, weight = ctx.saved_variables g = grad_output * weight.view(1, C, 1, 1) mean_g = g.mean(dim=1, keepdim=True) mean_gy = (g * y).mean(dim=1, keepdim=True) gx = 1. / torch.sqrt(var + eps) * (g - y * mean_gy - mean_g) return gx, (grad_output * y).sum(dim=3).sum(dim=2).sum(dim=0), grad_output.sum(dim=3).sum(dim=2).sum( dim=0), None class LayerNorm2d(nn.Module): def __init__(self, channels, eps=1e-6): super(LayerNorm2d, self).__init__() self.register_parameter('weight', nn.Parameter(torch.ones(channels))) self.register_parameter('bias', nn.Parameter(torch.zeros(channels))) self.eps = eps def forward(self, x): return LayerNormFunction.apply(x, self.weight, self.bias, self.eps) class SimpleGate(nn.Module): def forward(self, x): x1, x2 = x.chunk(2, dim=1) return x1 * x2 class NAFBlock(nn.Module): def __init__(self, c, DW_Expand=2, FFN_Expand=2, drop_out_rate=0.): super().__init__() dw_channel = c * DW_Expand self.conv1 = nn.Conv2d(in_channels=c, out_channels=dw_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True) self.conv2 = nn.Conv2d(in_channels=dw_channel, out_channels=dw_channel, kernel_size=3, padding=1, stride=1, groups=dw_channel, bias=True) self.conv3 = nn.Conv2d(in_channels=dw_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True) # Simplified Channel Attention self.sca = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels=dw_channel // 2, out_channels=dw_channel // 2, kernel_size=1, padding=0, stride=1, groups=1, bias=True), ) # SimpleGate self.sg = SimpleGate() ffn_channel = FFN_Expand * c self.conv4 = nn.Conv2d(in_channels=c, out_channels=ffn_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True) self.conv5 = nn.Conv2d(in_channels=ffn_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True) self.norm1 = LayerNorm2d(c) self.norm2 = LayerNorm2d(c) self.dropout1 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity() self.dropout2 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity() self.beta = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True) self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True) def forward(self, inp): x = inp x = self.norm1(x) x = self.conv1(x) x = self.conv2(x) x = self.sg(x) x = x * self.sca(x) x = self.conv3(x) x = self.dropout1(x) y = inp + x * self.beta x = self.conv4(self.norm2(y)) x = self.sg(x) x = self.conv5(x) x = self.dropout2(x) return y + x * self.gamma class NAF_Video(nn.Module): def __init__(self,args, img_channel=4, width=64, middle_blk_num=12, enc_blk_nums=[2, 2, 4, 8], dec_blk_nums=[2, 2, 2, 2]): super().__init__() self.lrelu = nn.LeakyReLU(0.2) self.convfist = nn.Conv2d(4, 64, 3, 1, 1) self.feature_extraction = ResList(5, 64) self.ending = nn.Conv2d(in_channels=width, out_channels=4, kernel_size=3, padding=1, stride=1, groups=1, bias=True) self.encoders = nn.ModuleList() self.decoders = nn.ModuleList() self.middle_blks = nn.ModuleList() self.ups = nn.ModuleList() self.downs = nn.ModuleList() chan = width for num in enc_blk_nums: self.encoders.append( nn.Sequential( *[NAFBlock(chan) for _ in range(num)] ) ) self.downs.append( nn.Conv2d(chan, 2*chan, 2, 2) ) chan = chan * 2 self.middle_blks = \ nn.Sequential( *[NAFBlock(chan) for _ in range(middle_blk_num)] ) for num in dec_blk_nums: self.ups.append( nn.Sequential( nn.Conv2d(chan, chan * 2, 1, bias=False), nn.PixelShuffle(2) ) ) chan = chan // 2 self.decoders.append( nn.Sequential( *[NAFBlock(chan) for _ in range(num)] ) ) self.padder_size = 2 ** len(self.encoders) # 16 def forward(self, x): center = x x = self.lrelu(self.convfist(x)) x = self.feature_extraction(x) encs = [] for encoder, down in zip(self.encoders, self.downs): x = encoder(x) encs.append(x) x = down(x) x = self.middle_blks(x) for decoder, up, enc_skip in zip(self.decoders, self.ups, encs[::-1]): x = up(x) x = x + enc_skip x = decoder(x) x = self.ending(x) x = x + center return x def load_networks(network, resume, strict=True): load_path = resume if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel): #会带有.module network = network.module #可以当前net的把.module 去掉 load_net = torch.load(load_path, map_location=torch.device('cuda')) load_net_clean = OrderedDict() # remove unnecessary 'module.' for k, v in load_net.items(): #可以把加载的net的把.module 去掉 if k.startswith('module.'): load_net_clean[k[7:]] = v else: load_net_clean[k] = v network.load_state_dict(load_net_clean, strict=True)