def __init__(self, nf=64, nframes=3, groups=8, front_RBs=5, back_RBs=10):
super(OnlyNonLocalNet, self).__init__()
self.nf = nf
self.in_frames = 1 + nframes // 2
self.ot_frames = nframes
self.ratio = 2 # 缩小比例
p_size = 48 # a place holder, not so useful
patch_size = (p_size, p_size)
n_layers = 1
hidden_dim = []
for i in range(n_layers):
hidden_dim.append(nf)
ResidualBlock_noBN_f = functools.partial(mutil.ResidualBlock_noBN,
nf=nf)
self.non_local = NONLocalBlock2D(3 * nframes * (self.ratio**2),
sub_sample=False,
bn_layer=False)
self.conv_first = nn.Conv2d(3, nf, 3, 1, 1, bias=True)
self.feature_extraction = mutil.make_layer(ResidualBlock_noBN_f,
front_RBs)
# reconstruction
self.recon_trunk = mutil.make_layer(ResidualBlock_noBN_f, back_RBs)
# upsampling
self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
self.upconv2 = nn.Conv2d(nf, 64 * 4, 3, 1, 1, bias=True)
self.pixel_shuffle = nn.PixelShuffle(2)
self.HRconv = nn.Conv2d(64, 64, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(64, 3, 3, 1, 1, bias=True)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
def __init__(self, in_nc=3, nf=64):
super(ResNet_alpha_beta_sconv, self).__init__()
self.conv_1 = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True) # 64f
basic_block_64 = functools.partial(mutil.ResidualBlock_noBN, nf=nf)
self.encoder1 = mutil.make_layer(basic_block_64, 2) # 64f
self.down_conv_1 = nn.Conv2d(nf, nf, 4, 2, 1, bias=False)
self.conv_2 = nn.Conv2d(nf, nf * 2, 3, 1, 1, bias=True) # 128f
basic_block_128 = functools.partial(mutil.ResidualBlock_noBN,
nf=nf * 2)
self.encoder2 = mutil.make_layer(basic_block_128, 2) # 128f
self.down_conv_2 = nn.Conv2d(nf * 2, nf * 2, 4, 2, 1, bias=False)
self.conv_3 = nn.Conv2d(nf * 2, nf * 4, 3, 1, 1, bias=True) # 256f
basic_block_256 = functools.partial(mutil.ResidualBlock_noBN,
nf=nf * 4)
self.encoder3 = mutil.make_layer(basic_block_256, 2) # 256f
self.down_conv_3 = nn.Conv2d(nf * 4, nf * 4, 4, 2, 1, bias=False)
self.conv_4 = nn.Conv2d(nf * 4, nf, 3, 1, 1, bias=True) # 64f
self.conv_5 = nn.Conv2d(nf, 6, 3, 1, 1, bias=True) # 6f
# pooling
#self.avg_pool = nn.AvgPool2d(2)
self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
def __init__(self, in_nc=3, nf=64):
super(UDC_dense_4, self).__init__()
# encoder
self.conv_1 = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True) # 64f
basic_block_64 = functools.partial(ResidualDenseBlock_5C, nf=nf)
self.encoder1 = mutil.make_layer(basic_block_64, 2) # 64f
self.conv_2 = nn.Conv2d(nf, nf * 2, 3, 1, 1, bias=True) # 128f
basic_block_128 = functools.partial(ResidualDenseBlock_5C, nf=nf * 2)
self.encoder2 = mutil.make_layer(basic_block_128, 2) # 128f
self.conv_3 = nn.Conv2d(nf * 2, nf * 4, 3, 1, 1, bias=True) # 256f
basic_block_256 = functools.partial(ResidualDenseBlock_5C, nf=nf * 4)
self.encoder3 = mutil.make_layer(basic_block_256, 2) # 256f
self.conv_4 = nn.Conv2d(nf * 4, nf * 8, 3, 1, 1, bias=True) # 512f
basic_block_512 = functools.partial(ResidualDenseBlock_5C, nf=nf * 8)
self.encoder4 = mutil.make_layer(basic_block_512, 2) # 512f
self.conv_5 = nn.Conv2d(nf * 8, nf * 8, 3, 1, 1, bias=True) # 64f
self.conv_6 = nn.Conv2d(nf * 8, 512, 3, 1, 1, bias=True) # 64f
# pooling
self.avg_pool = nn.AvgPool2d(2)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# decoder
self.conv_7 = nn.Conv2d(512, 256, 3, 1, padding=1, bias=True)
self.conv_8 = nn.Conv2d(768, 512, 3, 1, padding=1, bias=True)
self.conv_9 = nn.Conv2d(512, 512, 3, 1, padding=1, bias=True)
self.conv_10 = nn.Conv2d(640, 512, 3, 1, padding=1, bias=True)
self.conv_11 = nn.Conv2d(512, 512, 3, 1, padding=1, bias=True)
self.conv_12 = nn.Conv2d(384, 256, 3, 1, padding=1, bias=True)
self.conv_13 = nn.Conv2d(256, 256, 3, 1, padding=1, bias=True)
self.conv_14 = nn.Conv2d(192, 128, 3, 1, padding=1, bias=True)
self.conv_15 = nn.Conv2d(128, 128, 3, 1, padding=1, bias=True)
self.conv_16 = nn.Conv2d(96, 64, 3, 1, padding=1, bias=True)
self.conv_17 = nn.Conv2d(64, 3, 3, 1, padding=1, bias=True)
self.pixshuffle_1 = nn.PixelShuffle(upscale_factor=2)
self.pixshuffle_2 = nn.PixelShuffle(upscale_factor=2)
self.pixshuffle_3 = nn.PixelShuffle(upscale_factor=2)
self.pixshuffle_4 = nn.PixelShuffle(upscale_factor=2)
def __init__(self,
nf=64,
nframes=5,
groups=8,
front_RBs=5,
back_RBs=10,
center=None,
predeblur=False,
HR_in=False,
w_TSA=True):
super(EDVR, self).__init__()
self.nf = nf
self.center = nframes // 2 if center is None else center
self.is_predeblur = True if predeblur else False
self.HR_in = True if HR_in else False
self.w_TSA = w_TSA
ResidualBlock_noBN_f = functools.partial(mutil.ResidualBlock_noBN,
nf=nf)
#### extract features (for each frame)
if self.is_predeblur:
self.pre_deblur = Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in)
self.conv_1x1 = nn.Conv2d(nf, nf, 1, 1, bias=True)
else:
if self.HR_in:
self.conv_first_1 = nn.Conv2d(3, nf, 3, 1, 1, bias=True)
self.conv_first_2 = nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
self.conv_first_3 = nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
else:
self.conv_first = nn.Conv2d(3, nf, 3, 1, 1, bias=True)
self.feature_extraction = mutil.make_layer(ResidualBlock_noBN_f,
front_RBs)
self.fea_L2_conv1 = nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
self.fea_L2_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.fea_L3_conv1 = nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
self.fea_L3_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.pcd_align = PCD_Align(nf=nf, groups=groups)
if self.w_TSA:
self.tsa_fusion = TSA_Fusion(nf=nf,
nframes=nframes,
center=self.center)
else:
self.tsa_fusion = nn.Conv2d(nframes * nf, nf, 1, 1, bias=True)
#### reconstruction
self.recon_trunk = mutil.make_layer(ResidualBlock_noBN_f, back_RBs)
#### upsampling
self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
self.upconv2 = nn.Conv2d(nf, 64 * 4, 3, 1, 1, bias=True)
self.pixel_shuffle = nn.PixelShuffle(2)
self.HRconv = nn.Conv2d(64, 64, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(64, 3, 3, 1, 1, bias=True)
#### activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1)
def __init__(self, in_nc=3, out_nc=3, nf=64, nb=16, upscale=4):
super(MSRResNet, self).__init__()
self.upscale = upscale
self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
basic_block = functools.partial(mutil.ResidualBlock_noBN, nf=nf)
self.recon_trunk = mutil.make_layer(basic_block, nb)
# upsampling
if self.upscale == 2:
self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
self.pixel_shuffle = nn.PixelShuffle(2)
elif self.upscale == 3:
self.upconv1 = nn.Conv2d(nf, nf * 9, 3, 1, 1, bias=True)
self.pixel_shuffle = nn.PixelShuffle(3)
elif self.upscale == 4:
self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
self.upconv2 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
self.pixel_shuffle = nn.PixelShuffle(2)
self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# initialization
mutil.initialize_weights(
[self.conv_first, self.upconv1, self.HRconv, self.conv_last], 0.1)
if self.upscale == 4:
mutil.initialize_weights(self.upconv2, 0.1)
def __init__(self, in_nc, out_nc, nf, nb, gc=32, differential=None, time_dependent=False, adjoint=False, sb=5):
super(RRDBNet, self).__init__()
self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
if differential == "checkpointed":
self.conv_trunk = SRTrunk(nf, nb, make_odeblock(5, 'RK4'))
mutil.initialize_weights(self.conv_trunk.odefunc.convs)
elif differential == "standard":
self.conv_trunk = ODEBlock(ODEfunc(nf, nb=nb, normalization=False, time_dependent=time_dependent), adjoint=adjoint)
mutil.initialize_weights(self.conv_trunk.odefunc.convs)
elif differential == "sequential":
self.conv_trunk = nn.Sequential(*[ODEBlock(ODEfunc(nf, nb=sb, normalization=False, time_dependent=time_dependent), adjoint=adjoint) for _ in range(nb)])
for block in self.conv_trunk:
mutil.initialize_weights(block.odefunc.convs)
elif differential == "augmented":
augment_dim = nf//4
method = 'dopri5'
if method == 'euler':
warnings.warn("euler mode")
self.conv_trunk = AugBlock(AugFunc(nf=nf, nb=nb, augment_dim=augment_dim, time_dependent=time_dependent), adjoint=adjoint, is_conv=True, method=method)
self.trunk_conv = nn.Conv2d(nf+augment_dim, nf, 3, 1, 1, bias=True)
mutil.initialize_weights(self.conv_trunk.odefunc.convs)
elif differential is None or differential == "nodiff":
RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
self.conv_trunk = mutil.make_layer(RRDB_block_f, nb)
else:
raise NotImplementedError("unrecognized differential system passed")
#### upsampling
self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def __init__(self, in_nc=3, nf=64):
super(ResNet_alpha_beta_decoder_3x3_IN_encoder, self).__init__()
# encoder
self.conv_1 = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True) # 64f
basic_block_64 = functools.partial(mutil.ResidualBlock_IN, nf=nf)
self.encoder1 = mutil.make_layer(basic_block_64, 2) # 64f
self.conv_2 = nn.Conv2d(nf, nf * 2, 3, 1, 1, bias=True) # 128f
basic_block_128 = functools.partial(mutil.ResidualBlock_IN, nf=nf * 2)
self.encoder2 = mutil.make_layer(basic_block_128, 2) # 128f
self.conv_3 = nn.Conv2d(nf * 2, nf * 4, 3, 1, 1, bias=True) # 256f
basic_block_256 = functools.partial(mutil.ResidualBlock_IN, nf=nf * 4)
self.encoder3 = mutil.make_layer(basic_block_256, 2) # 256f
self.conv_4 = nn.Conv2d(nf * 4, nf, 3, 1, 1, bias=True) # 64f
self.bn_4 = nn.InstanceNorm2d(nf, affine=True)
self.conv_5 = nn.Conv2d(nf, 6, 3, 1, 1, bias=True) # 6f
# pooling
self.avg_pool = nn.AvgPool2d(2)
self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# decoder
self.conv_6 = nn.Conv2d(6, 32, 3, 1, padding=1, bias=True)
#self.bn_6 = nn.InstanceNorm2d(32, affine=True)
self.conv_7 = nn.Conv2d(96, 32, 3, 1, padding=1, bias=True)
#self.bn_7 = nn.InstanceNorm2d(32, affine=True)
self.conv_8 = nn.Conv2d(32, 32, 3, 1, padding=1, bias=True)
#self.bn_8 = nn.InstanceNorm2d(32, affine=True)
self.conv_9 = nn.Conv2d(288, 32, 3, 1, padding=1, bias=True)
#self.bn_9 = nn.InstanceNorm2d(32, affine=True)
self.conv_10 = nn.Conv2d(32, 32, 3, 1, padding=1, bias=True)
#self.bn_10 = nn.InstanceNorm2d(32, affine=True)
self.conv_11 = nn.Conv2d(160, 32, 3, 1, padding=1, bias=True)
#self.bn_11 = nn.InstanceNorm2d(32, affine=True)
self.conv_12 = nn.Conv2d(32, 32, 3, 1, padding=1, bias=True)
#self.bn_12 = nn.InstanceNorm2d(32, affine=True)
self.conv_13 = nn.Conv2d(96, 32, 3, 1, padding=1, bias=True)
#self.bn_13 = nn.InstanceNorm2d(32, affine=True)
self.conv_14 = nn.Conv2d(32, 6, 3, 1, padding=1, bias=True)
def __init__(self, nf=64, nframes=3, groups=8, front_RBs=5, back_RBs=10):
super(NonLocalNet, self).__init__()
self.nf = nf
self.in_frames = 1 + nframes // 2
self.ot_frames = nframes
p_size = 48 # a place holder, not so useful
patch_size = (p_size, p_size)
n_layers = 1
hidden_dim = []
for i in range(n_layers):
hidden_dim.append(nf)
ResidualBlock_noBN_f = functools.partial(mutil.ResidualBlock_noBN,
nf=nf)
self.non_local = NONLocalBlock3D(3)
self.conv_first = nn.Conv2d(3, nf, 3, 1, 1, bias=True)
self.feature_extraction = mutil.make_layer(ResidualBlock_noBN_f,
front_RBs)
# self.ConvBRNN = BiDeformableConvLSTM(input_size=patch_size, input_dim=nf, hidden_dim=hidden_dim,
# kernel_size=(3, 3), num_layers=1, batch_first=True, front_RBs=front_RBs,
# groups=groups)
self.ConvBRNN = BiDeformableConvGRU(input_size=patch_size,
input_dim=nf,
hidden_dim=hidden_dim,
kernel_size=[(3, 3)],
num_layers=n_layers,
batch_first=True,
groups=groups)
# reconstruction
self.recon_trunk = mutil.make_layer(ResidualBlock_noBN_f, back_RBs)
# upsampling
self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
self.upconv2 = nn.Conv2d(nf, 64 * 4, 3, 1, 1, bias=True)
self.pixel_shuffle = nn.PixelShuffle(2)
self.HRconv = nn.Conv2d(64, 64, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(64, 3, 3, 1, 1, bias=True)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
def __init__(self, in_nc=3, nf=64, nb=3):
super(ResNet_plain, self).__init__()
# directly output the corrected image
self.conv_1 = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True) # 64f
basic_block_64 = functools.partial(mutil.ResidualBlock_IN, nf=nf)
self.res_blocks = mutil.make_layer(basic_block_64, nb) # 64f
self.conv_2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(nf, 3, 3, 1, 1, bias=True)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
def __init__(self, in_nc, out_nc, nf, nb, gc=32):
super(RRDBNet, self).__init__()
RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
self.RRDB_trunk = mutil.make_layer(RRDB_block_f, nb)
self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
#### upsampling
self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
def __init__(self,
in_nc,
out_nc,
nf,
nb,
nframes=5,
groups=8,
front_RBs=3,
gc=128,
use_snorm=False,
center=None):
super(DenseMWNet, self).__init__()
ResidualBlock_noBN_f = functools.partial(ResidualBlock_noBN, nf=nf)
act = nn.LeakyReLU(True)
conv = common.default_conv
self.center = nframes // 2 if center is None else center
self.DWT = common.DWT()
self.IWT = common.IWT()
self.motion_align = MWP_Align(nf=nf, groups=groups)
self.attention_fusion = nn.Conv2d(nframes * nf, nf, 1, 1, bias=True)
self.conv_first_1 = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
self.conv_first_2 = nn.Conv2d(nf * 4, nf, 3, 1, 1, bias=True)
self.conv_first_3 = nn.Conv2d(nf * 4, nf, 3, 1, 1, bias=True)
self.feature_extraction = mutil.make_layer(ResidualBlock_noBN_f,
front_RBs)
self.multi_fea_l1 = Multi_extfea(nf=nf, gc=nf, use_snorm=use_snorm)
self.fea_L2_conv1 = nn.Conv2d(nf * 4, nf, 3, 1, 1, bias=True)
self.fea_L2_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.fea_L3_conv1 = nn.Conv2d(nf * 4, nf, 3, 1, 1, bias=True)
self.fea_L3_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
# self.motion_align = DWT_motion_Pyramid(gc=nf)
ResidualBlock = []
for i in range(nb):
ResidualBlock.append(WDRB(nf=nf * 4, gc=gc, use_snorm=use_snorm))
self.ResidualBlock = nn.Sequential(*ResidualBlock)
### upsampling
self.upconv1 = nn.Conv2d(nf // 4, nf, 3, 1, 1, bias=True)
self.upconv2 = nn.Conv2d(nf // 4, nf, 3, 1, 1, bias=True)
# self.HRconv = nn.Conv2d(in_nc*4, in_nc, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def __init__(self, in_nc=3, out_nc=3, nf=64, nb=16):
super(MSRResNet, self).__init__()
self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
basic_block = functools.partial(mutil.ResidualBlock_noBN, nf=nf)
self.recon_trunk = mutil.make_layer(basic_block, nb)
self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# initialization
mutil.initialize_weights(
[self.conv_first, self.HRconv, self.conv_last], 0.1)
def __init__(self,
in_nc=3,
out_nc=3,
nf=64,
nb=16,
upscale=4,
differential=False):
super(MSRResNet, self).__init__()
self.upscale = upscale
self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
if differential == "checkpointed":
self.conv_trunk = SRTrunk(nf, nb, make_odeblock(2, 'RK4'))
mutil.initialize_weights(self.conv_trunk.odefunc.convs)
elif differential == "standard":
self.conv_trunk = ODEBlock(ODEfunc(nf, nb=nb, normalization=False))
mutil.initialize_weights(self.conv_trunk.odefunc.convs)
elif differential is None:
basic_block = functools.partial(mutil.ResidualBlock_noBN, nf=nf)
self.conv_trunk = mutil.make_layer(basic_block, nb)
else:
raise NotImplementedError(
"unrecognized differential system passed")
# upsampling
if self.upscale == 2:
self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
self.pixel_shuffle = nn.PixelShuffle(2)
elif self.upscale == 3:
self.upconv1 = nn.Conv2d(nf, nf * 9, 3, 1, 1, bias=True)
self.pixel_shuffle = nn.PixelShuffle(3)
elif self.upscale == 4:
self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
self.upconv2 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
self.pixel_shuffle = nn.PixelShuffle(2)
self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=False)
# initialization
mutil.initialize_weights(
[self.conv_first, self.upconv1, self.HRconv, self.conv_last], 0.1)
if self.upscale == 4:
mutil.initialize_weights(self.upconv2, 0.1)
def __init__(self, in_nc, out_nc, nf, nb, gc=32, scale=4, opt=None):
self.opt = opt
super(RRDBNet, self).__init__()
RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
self.scale = scale
# self.downhaar = flow.HaarDownsampling(channel_in=in_nc)
self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
self.RRDB_trunk = mutil.make_layer(RRDB_block_f, nb)
self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
#### upsampling
self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
if self.scale >= 8:
self.upconv3 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
if self.scale >= 16:
self.upconv4 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
if self.scale >= 32:
self.upconv5 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def __init__(self, in_nc=3, nf=64):
super(UDC_res_6, self).__init__()
#CA
modules_body_1 = []
modules_body_1.append(CALayer(nf, reduction=8))
self.body_1 = nn.Sequential(*modules_body_1)
# encoder
self.conv_1 = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True) # 64f
basic_block_64 = functools.partial(mutil.ResidualBlock_noBN, nf=nf)
self.encoder1 = mutil.make_layer(basic_block_64, 2) # 64f
self.conv_2 = nn.Conv2d(nf, nf * 2, 3, 1, 1, bias=True) # 128f
basic_block_128 = functools.partial(mutil.ResidualBlock_noBN,
nf=nf * 2)
self.encoder2 = mutil.make_layer(basic_block_128, 2) # 128f
self.conv_3 = nn.Conv2d(nf * 2, nf * 4, 3, 1, 1, bias=True) # 256f
basic_block_256 = functools.partial(mutil.ResidualBlock_noBN,
nf=nf * 4)
self.encoder3 = mutil.make_layer(basic_block_256, 2) # 256f
self.conv_4 = nn.Conv2d(nf * 4, nf * 8, 3, 1, 1, bias=True) # 512f
basic_block_512 = functools.partial(mutil.ResidualBlock_noBN,
nf=nf * 8)
self.encoder4 = mutil.make_layer(basic_block_512, 2) # 512f
self.conv_5 = nn.Conv2d(nf * 8, nf * 16, 3, 1, 1, bias=True) # 1024f
basic_block_1024 = functools.partial(mutil.ResidualBlock_noBN,
nf=nf * 16)
self.encoder5 = mutil.make_layer(basic_block_1024, 2) # 1024f
self.conv_6 = nn.Conv2d(nf * 16, nf * 32, 3, 1, 1, bias=True) # 2048f
basic_block_2048 = functools.partial(mutil.ResidualBlock_noBN,
nf=nf * 32)
self.encoder6 = mutil.make_layer(basic_block_2048, 2) # 2048f
self.conv_7 = nn.Conv2d(nf * 32, nf * 16, 3, 1, 1, bias=True) # 1024f
self.conv_8 = nn.Conv2d(nf * 16, nf * 16, 3, 1, 1, bias=True) # 512f
# pooling
self.avg_pool = nn.AvgPool2d(2)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
# decoder
self.conv_9 = nn.Conv2d(1024, 512, 3, 1, padding=1, bias=True)
self.conv_10 = nn.Conv2d(1536, 1024, 3, 1, padding=1, bias=True)
self.conv_11 = nn.Conv2d(1024, 1024, 3, 1, padding=1, bias=True)
self.conv_12 = nn.Conv2d(2304, 1024, 3, 1, padding=1, bias=True)
self.conv_13 = nn.Conv2d(1024, 1024, 3, 1, padding=1, bias=True)
self.conv_14 = nn.Conv2d(1280, 1024, 3, 1, padding=1, bias=True)
self.conv_15 = nn.Conv2d(1024, 1024, 3, 1, padding=1, bias=True)
self.conv_16 = nn.Conv2d(768, 512, 3, 1, padding=1, bias=True)
self.conv_17 = nn.Conv2d(512, 512, 3, 1, padding=1, bias=True)
self.conv_18 = nn.Conv2d(384, 256, 3, 1, padding=1, bias=True)
self.conv_19 = nn.Conv2d(256, 256, 3, 1, padding=1, bias=True)
self.conv_20 = nn.Conv2d(192, 160, 3, 1, padding=1, bias=True)
self.conv_21 = nn.Conv2d(160, 160, 3, 1, padding=1, bias=True)
self.conv_22 = nn.Conv2d(104, 96, 3, 1, padding=1, bias=True)
self.conv_23 = nn.Conv2d(96, 3, 3, 1, padding=1, bias=True)
self.pixshuffle_1 = nn.PixelShuffle(upscale_factor=2)
self.pixshuffle_2 = nn.PixelShuffle(upscale_factor=2)
self.pixshuffle_3 = nn.PixelShuffle(upscale_factor=2)
self.pixshuffle_4 = nn.PixelShuffle(upscale_factor=2)
self.pixshuffle_5 = nn.PixelShuffle(upscale_factor=2)
self.pixshuffle_6 = nn.PixelShuffle(upscale_factor=2)
