diff --git a/codes/models/archs/RRDBNet_arch.py b/codes/models/archs/RRDBNet_arch.py index 26ddee36..6c4ac41c 100644 --- a/codes/models/archs/RRDBNet_arch.py +++ b/codes/models/archs/RRDBNet_arch.py @@ -1,9 +1,12 @@ +import os + import torch import torch.nn as nn import torch.nn.functional as F +import torchvision from torch.utils.checkpoint import checkpoint_sequential -from models.archs.arch_util import make_layer, default_init_weights +from models.archs.arch_util import make_layer, default_init_weights, ConvGnSilu class ResidualDenseBlock(nn.Module): @@ -79,6 +82,44 @@ class RRDB(nn.Module): return out * 0.2 + x +class RRDBWithBypass(nn.Module): + """Residual in Residual Dense Block. + + Used in RRDB-Net in ESRGAN. + + Args: + mid_channels (int): Channel number of intermediate features. + growth_channels (int): Channels for each growth. + """ + + def __init__(self, mid_channels, growth_channels=32): + super(RRDBWithBypass, self).__init__() + self.rdb1 = ResidualDenseBlock(mid_channels, growth_channels) + self.rdb2 = ResidualDenseBlock(mid_channels, growth_channels) + self.rdb3 = ResidualDenseBlock(mid_channels, growth_channels) + self.bypass = nn.Sequential(ConvGnSilu(mid_channels*2, mid_channels, kernel_size=3, bias=True, activation=True, norm=True), + ConvGnSilu(mid_channels, mid_channels//2, kernel_size=3, bias=False, activation=True, norm=False), + ConvGnSilu(mid_channels//2, 1, kernel_size=3, bias=False, activation=False, norm=False), + nn.Sigmoid()) + + def forward(self, x): + """Forward function. + + Args: + x (Tensor): Input tensor with shape (n, c, h, w). + + Returns: + Tensor: Forward results. + """ + out = self.rdb1(x) + out = self.rdb2(out) + out = self.rdb3(out) + bypass = self.bypass(torch.cat([x, out], dim=1)) + self.bypass_map = bypass.detach().clone() + # Emperically, we use 0.2 to scale the residual for better performance + return out * 0.2 * bypass + x + + class RRDBNet(nn.Module): """Networks consisting of Residual in Residual Dense Block, which is used in ESRGAN. @@ -100,11 +141,15 @@ class RRDBNet(nn.Module): out_channels, mid_channels=64, num_blocks=23, - growth_channels=32): + growth_channels=32, + body_block=RRDB, + blocks_per_checkpoint=4): super(RRDBNet, self).__init__() + self.num_blocks = num_blocks + self.blocks_per_checkpoint = blocks_per_checkpoint self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, 1, 1) self.body = make_layer( - RRDB, + body_block, num_blocks, mid_channels=mid_channels, growth_channels=growth_channels) @@ -134,7 +179,7 @@ class RRDBNet(nn.Module): """ feat = self.conv_first(x) - body_feat = self.conv_body(checkpoint_sequential(self.body, 5, feat)) + body_feat = self.conv_body(checkpoint_sequential(self.body, self.num_blocks // self.blocks_per_checkpoint, feat)) feat = feat + body_feat # upsample feat = self.lrelu( @@ -142,4 +187,9 @@ class RRDBNet(nn.Module): feat = self.lrelu( self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest'))) out = self.conv_last(self.lrelu(self.conv_hr(feat))) - return out \ No newline at end of file + return out + + def visual_dbg(self, step, path): + for i, bm in enumerate(self.body): + torchvision.utils.save_image(bm.bypass_map.cpu().float(), os.path.join(path, "%i_bypass_%i.png" % (step, i+1))) + diff --git a/codes/models/networks.py b/codes/models/networks.py index e4a4e34f..cc2b6f31 100644 --- a/codes/models/networks.py +++ b/codes/models/networks.py @@ -39,6 +39,10 @@ def define_G(opt, net_key='network_G', scale=None): elif which_model == 'RRDBNet': 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']) + elif which_model == 'RRDBNetBypass': + 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']) elif which_model == 'rcan': #args: n_resgroups, n_resblocks, res_scale, reduction, scale, n_feats opt_net['rgb_range'] = 255 diff --git a/codes/train.py b/codes/train.py index 852621c3..eccc14ab 100644 --- a/codes/train.py +++ b/codes/train.py @@ -265,7 +265,7 @@ class Trainer: if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_prog_imgset_multifaceted_chained.yml') + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_mi1_rrdb4x_6bl_bypass.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() diff --git a/codes/train2.py b/codes/train2.py index 456b570c..194a6257 100644 --- a/codes/train2.py +++ b/codes/train2.py @@ -278,7 +278,7 @@ class Trainer: if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_mi1_tecogen.yml') + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_mi1_rrdb4x_10bl_bypass.yml') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') args = parser.parse_args() opt = option.parse(args.opt, is_train=True)