mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add additive mode to rrdb

Browse Source
This commit is contained in:
James Betker 2020-11-16 20:45:09 -07:00
parent 2a507987df
commit 8a19c9ae15
4 changed files with 23 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
codes/models/archs/RRDBNet_arch.py
View File

@ -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):

7
codes/models/networks.py
View File

@ -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

6
codes/scripts/extract_square_images.py
View File

@ -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

2
codes/train.py
View File

@ -291,7 +291,7 @@ 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()