Add additive mode to rrdb
This commit is contained in:
parent
2a507987df
commit
8a19c9ae15
|
@ -144,7 +144,9 @@ class RRDBNet(nn.Module):
|
|||
growth_channels=32,
|
||||
body_block=RRDB,
|
||||
blocks_per_checkpoint=4,
|
||||
scale=4):
|
||||
scale=4,
|
||||
additive_mode="not_additive" # Options: "not_additive", "additive", "additive_enforced"
|
||||
):
|
||||
super(RRDBNet, self).__init__()
|
||||
self.num_blocks = num_blocks
|
||||
self.blocks_per_checkpoint = blocks_per_checkpoint
|
||||
|
@ -166,6 +168,10 @@ class RRDBNet(nn.Module):
|
|||
self.conv_hr = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
|
||||
self.conv_last = nn.Conv2d(mid_channels, out_channels, 3, 1, 1)
|
||||
|
||||
self.additive_mode = additive_mode
|
||||
if additive_mode == "additive_enforced":
|
||||
self.add_enforced_pool = nn.AvgPool2d(kernel_size=scale, stride=scale)
|
||||
|
||||
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
||||
|
||||
for m in [
|
||||
|
@ -202,6 +208,14 @@ class RRDBNet(nn.Module):
|
|||
else:
|
||||
feat = self.lrelu(self.conv_up2(feat))
|
||||
out = self.conv_last(self.lrelu(self.conv_hr(feat)))
|
||||
if "additive" in self.additive_mode:
|
||||
x_interp = F.interpolate(x, scale_factor=self.scale, mode='bilinear')
|
||||
if self.additive_mode == 'additive':
|
||||
out = out + x_interp
|
||||
elif self.additive_mode == 'additive_enforced':
|
||||
out_pooled = self.add_enforced_pool(out)
|
||||
out = out - F.interpolate(out_pooled, scale_factor=self.scale, mode='nearest')
|
||||
out = out + x_interp
|
||||
return out
|
||||
|
||||
def visual_dbg(self, step, path):
|
||||
|
|
|
@ -43,12 +43,15 @@ def define_G(opt, net_key='network_G', scale=None):
|
|||
netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
|
||||
nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale'])
|
||||
elif which_model == 'RRDBNet':
|
||||
additive_mode = opt_net['additive_mode'] if 'additive_mode' in opt_net.keys() else 'not_additive'
|
||||
netG = RRDBNet_arch.RRDBNet(in_channels=opt_net['in_nc'], out_channels=opt_net['out_nc'],
|
||||
mid_channels=opt_net['nf'], num_blocks=opt_net['nb'])
|
||||
mid_channels=opt_net['nf'], num_blocks=opt_net['nb'], additive_mode=additive_mode)
|
||||
elif which_model == 'RRDBNetBypass':
|
||||
additive_mode = opt_net['additive_mode'] if 'additive_mode' in opt_net.keys() else 'not_additive'
|
||||
netG = RRDBNet_arch.RRDBNet(in_channels=opt_net['in_nc'], out_channels=opt_net['out_nc'],
|
||||
mid_channels=opt_net['nf'], num_blocks=opt_net['nb'], body_block=RRDBNet_arch.RRDBWithBypass,
|
||||
blocks_per_checkpoint=opt_net['blocks_per_checkpoint'], scale=opt_net['scale'])
|
||||
blocks_per_checkpoint=opt_net['blocks_per_checkpoint'], scale=opt_net['scale'],
|
||||
additive_mode=additive_mode)
|
||||
elif which_model == 'rcan':
|
||||
#args: n_resgroups, n_resblocks, res_scale, reduction, scale, n_feats
|
||||
opt_net['rgb_range'] = 255
|
||||
|
|
|
@ -19,11 +19,7 @@ def main():
|
|||
# compression time. If read raw images during training, use 0 for faster IO speed.
|
||||
|
||||
opt['dest'] = 'file'
|
||||
opt['input_folder'] = ['F:\\4k6k\\datasets\\images\\div2k\\DIV2K_train_HR',
|
||||
'F:\\4k6k\\datasets\\images\\flickr\\flickr2k\\Flickr2K_HR',
|
||||
'F:\\4k6k\\datasets\\images\\flickr\\flickr-scrape\\filtered',
|
||||
'F:\\4k6k\\datasets\\images\\goodeats\\hq\\new_season\\images',
|
||||
'F:\\4k6k\datasets\\images\\youtube\\images']
|
||||
opt['input_folder'] = ['F:\\4k6k\datasets\\images\\youtube\\videos\\4k_quote_unquote\\images']
|
||||
opt['save_folder'] = 'F:\\4k6k\\datasets\\images\\ge_full_1024'
|
||||
opt['imgsize'] = 1024
|
||||
|
||||
|
|
|
@ -291,14 +291,14 @@ class Trainer:
|
|||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_stylegan2_celebA_separated_disc.yml')
|
||||
parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_mi1_rrdb4x_6bl_corrected_disc.yml')
|
||||
parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher')
|
||||
parser.add_argument('--local_rank', type=int, default=0)
|
||||
args = parser.parse_args()
|
||||
opt = option.parse(args.opt, is_train=True)
|
||||
trainer = Trainer()
|
||||
|
||||
#### distributed training settings
|
||||
#### distributed training settings
|
||||
if args.launcher == 'none': # disabled distributed training
|
||||
opt['dist'] = False
|
||||
trainer.rank = -1
|
||||
|
|
Loading…
Reference in New Issue
Block a user