def __init__(self, opt): super(NET, self).__init__() n_resblocks = opt.n_resblocks n_feats = opt.channels bias = opt.bias norm_type = opt.norm_type act_type = opt.act_type block_type = opt.block_type head = [common.ConvBlock(4, n_feats, 5, act_type=act_type, bias=True)] if block_type.lower() == 'rrdb': resblock = [common.RRDB(n_feats, n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(n_resblocks)] elif block_type.lower() == 'res': resblock = [common.ResBlock(n_feats, 3, norm_type, act_type, res_scale=1, bias=bias) for _ in range(n_resblocks)] else: raise RuntimeError('block_type is not supported') resblock += [common.ConvBlock(n_feats, n_feats, 3, bias=True)] tail = [common.ConvBlock(n_feats, 3, 3, bias=True)] self.model = nn.Sequential(*head, common.ShortcutBlock(nn.Sequential(*resblock)), *tail) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_normal_(m.weight) m.weight.requires_grad = True if m.bias is not None: m.bias.data.zero_() m.bias.requires_grad = True
def __init__(self, opt): super(green_res, self).__init__() sr_n_resblocks = opt.sr_n_resblocks dm_n_resblocks = opt.dm_n_resblocks sr_n_feats = opt.channels dm_n_feats = opt.channels scale = opt.scale denoise = opt.denoise block_type = opt.block_type act_type = opt.act_type bias = opt.bias norm_type = opt.norm_type self.head = common.ConvBlock(2, dm_n_feats, 5, act_type=act_type, bias=True) self.r1 = common.RRDB(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) self.r2 = common.RRDB(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) #self.r3 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) #self.r4 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) #self.r5 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) #self.r6 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) self.final = common.ConvBlock(dm_n_feats, dm_n_feats, 3, bias=bias) self.up = nn.Sequential( common.Upsampler(2, dm_n_feats, norm_type, act_type, bias=bias), common.ConvBlock(dm_n_feats, 1, 3, bias=True), nn.LeakyReLU(0.2, inplace=True))
def __init__(self, n_resblock=24, n_feats=256, scale=2, bias=True, norm_type=False, act_type='prelu'): super(NET, self).__init__() self.scale = scale m = [common.default_conv(1, n_feats, 3, stride=2)] m += [nn.PixelShuffle(2), common.ConvBlock(n_feats//4, n_feats, bias=True, act_type=act_type) ] m += [common.ResBlock(n_feats, 3, norm_type, act_type, res_scale=1, bias=bias) for _ in range(n_resblock)] for _ in range(int(math.log(scale, 2))): m += [nn.PixelShuffle(2), common.ConvBlock(n_feats//4, n_feats, bias=True, act_type=act_type) ] m += [common.default_conv(n_feats, 3, 3)] self.model = nn.Sequential(*m) for m in self.modules(): # pdb.set_trace() if isinstance(m, nn.Conv2d): # Xavier # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') nn.init.xavier_normal_(m.weight) m.weight.requires_grad = True if m.bias is not None: m.bias.data.zero_() m.bias.requires_grad = True
def __init__(self, opt): super(NET, self).__init__() # parameter denoise = opt.denoise block_type = opt.block_type n_feats = opt.channels act_type = opt.act_type bias = opt.bias norm_type = opt.norm_type n_resblocks = opt.n_resblocks # architecture if denoise: dm_head = [ common.ConvBlock(5, n_feats, 5, act_type=act_type, bias=True) ] else: dm_head = [ common.ConvBlock(4, n_feats, 5, act_type=act_type, bias=True) ] if block_type.lower() == 'rrdb': dm_resblock = [ common.RRDB(n_feats, n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(n_resblocks) ] elif block_type.lower() == 'res': dm_resblock = [ common.ResBlock(n_feats, 3, norm_type, act_type, res_scale=1, bias=bias) for _ in range(n_resblocks) ] else: raise RuntimeError('block_type is not supported') dm_resblock += [common.ConvBlock(n_feats, n_feats, 3, bias=True)] m_dm_up = [ common.Upsampler(2, n_feats, norm_type, act_type, bias=bias), common.ConvBlock(n_feats, 3, 3, bias=True) ] self.model_dm = nn.Sequential( *dm_head, common.ShortcutBlock(nn.Sequential(*dm_resblock)), *m_dm_up) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_normal_(m.weight) m.weight.requires_grad = True if m.bias is not None: m.bias.data.zero_() m.bias.requires_grad = True
def __init__(self, opt): super(m_res, self).__init__() sr_n_resblocks = opt.sr_n_resblocks dm_n_resblocks = opt.dm_n_resblocks sr_n_feats = opt.channels dm_n_feats = opt.channels scale = opt.scale denoise = opt.denoise block_type = opt.block_type act_type = opt.act_type bias = opt.bias norm_type = opt.norm_type self.r1 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) self.r2 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) self.r3 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) self.r4 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) self.r5 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) self.r6 = common.RRDB2(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) self.final = common.ConvBlock(dm_n_feats, dm_n_feats, 3, bias=bias)
def __init__(self, opt): super(NET, self).__init__() sr_n_resblocks = opt.sr_n_resblocks dm_n_resblocks = opt.dm_n_resblocks sr_n_feats = opt.channels dm_n_feats = opt.channels scale = opt.scale denoise = opt.denoise block_type = opt.block_type act_type = opt.act_type bias = opt.bias norm_type = opt.norm_type # define sr module if denoise: m_sr_head = [ common.ConvBlock(5, sr_n_feats, 5, act_type=act_type, bias=True) ] else: m_sr_head = [ common.ConvBlock(4, sr_n_feats, 5, act_type=act_type, bias=True) ] if block_type.lower() == 'rrdb': m_sr_resblock = [ common.RRDB(sr_n_feats, sr_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(sr_n_resblocks) ] elif block_type.lower() == 'dudb': m_sr_resblock = [ common.DUDB(sr_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(sr_n_resblocks) ] elif block_type.lower() == 'res': m_sr_resblock = [ common.ResBlock(sr_n_feats, 3, norm_type, act_type, res_scale=1, bias=bias) for _ in range(sr_n_resblocks) ] else: raise RuntimeError('block_type is not supported') m_sr_resblock += [ common.ConvBlock(sr_n_feats, sr_n_feats, 3, bias=bias) ] m_sr_up = [ common.Upsampler(scale, sr_n_feats, norm_type, act_type, bias=bias), common.ConvBlock(sr_n_feats, 4, 3, bias=True) ] # branch for sr_raw output m_sr_tail = [nn.PixelShuffle(2)] # define demosaick module m_dm_head = [ common.ConvBlock(4, dm_n_feats, 5, act_type=act_type, bias=True) ] if block_type.lower() == 'rrdb': m_dm_resblock = [ common.RRDB(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(dm_n_resblocks) ] elif block_type.lower() == 'dudb': m_dm_resblock = [ common.DUDB(dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(dm_n_resblocks) ] elif block_type.lower() == 'res': m_dm_resblock = [ common.ResBlock(dm_n_feats, 3, norm_type, act_type, res_scale=1, bias=bias) for _ in range(dm_n_resblocks) ] else: raise RuntimeError('block_type is not supported') m_dm_resblock += [ common.ConvBlock(dm_n_feats, dm_n_feats, 3, bias=bias) ] m_dm_up = [ common.Upsampler(2, dm_n_feats, norm_type, act_type, bias=bias), common.ConvBlock(dm_n_feats, 3, 3, bias=True) ] self.model_sr = nn.Sequential( *m_sr_head, common.ShortcutBlock(nn.Sequential(*m_sr_resblock)), *m_sr_up) self.sr_output = nn.Sequential(*m_sr_tail) self.model_dm = nn.Sequential( *m_dm_head, common.ShortcutBlock(nn.Sequential(*m_dm_resblock)), *m_dm_up) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_normal_(m.weight) m.weight.requires_grad = True if m.bias is not None: m.bias.data.zero_() m.bias.requires_grad = True
def __init__(self, name="CSPDarkNet53", **kwargs): super(CSPDarkNet53, self).__init__(self, name=name, **kwargs) self.conv1_1 = cm.ConvBlock((3, 3, 3, 32), activate_type="mish") self.conv1_2 = cm.ConvBlock((3, 3, 32, 64), downsample=True, activate_type="mish") self.conv2_1 = cm.ConvBlock((1, 1, 64, 64), activate_type="mish") self.conv2_2 = cm.ConvBlock((1, 1, 64, 64), activate_type="mish") self.res_blocks1 = [] for _ in range(1): self.res_blocks1.append(cm.ResBlock(64, 32, 64, activate_type="mish")) self.conv2_3 = cm.ConvBlock((1, 1, 64, 64), activate_type="mish") self.conv3_1 = cm.ConvBlock((1, 1, 128, 64), activate_type="mish") self.conv3_2 = cm.ConvBlock((3, 3, 64, 128), downsample=True, activate_type="mish") self.conv4_1 = cm.ConvBlock((1, 1, 128, 64), activate_type="mish") self.conv4_2 = cm.ConvBlock((1, 1, 128, 64), activate_type="mish") self.res_blocks2 = [] for _ in range(2): self.res_blocks2.append(cm.ResBlock(64, 64, 64, activate_type="mish")) self.conv4_3 = cm.ConvBlock((1, 1, 64, 64), activate_type="mish") self.conv5_1 = cm.ConvBlock((1, 1, 128, 128), activate_type="mish") self.conv5_2 = cm.ConvBlock((3, 3, 128, 256), downsample=True, activate_type="mish") self.conv6_1 = cm.ConvBlock((1, 1, 256, 128), activate_type="mish") self.conv6_2 = cm.ConvBlock((1, 1, 256, 128), activate_type="mish") self.res_blocks3 = [] for _ in range(8): self.res_blocks3.append(cm.ResBlock(128, 128, 128, activate_type="mish")) self.conv6_3 = cm.ConvBlock((1, 1, 128, 128), activate_type="mish") self.conv7_1 = cm.ConvBlock((1, 1, 256, 256), activate_type="mish") self.conv7_2 = cm.ConvBlock((3, 3, 256, 512), downsample=True, activate_type="mish") self.conv8_1 = cm.ConvBlock((1, 1, 512, 256), activate_type="mish") self.conv8_2 = cm.ConvBlock((1, 1, 512, 256), activate_type="mish") self.res_blocks4 = [] for _ in range(8): self.res_blocks4.append(cm.ResBlock(256, 256, 256, activate_type="mish")) self.conv8_3 = cm.ConvBlock((1, 1, 256, 256), activate_type="mish") self.conv9_1 = cm.ConvBlock((1, 1, 512, 512), activate_type="mish") self.conv9_2 = cm.ConvBlock((3, 3, 512, 1024), downsample=True, activate_type="mish") self.conv10_1 = cm.ConvBlock((1, 1, 1024, 512), activate_type="mish") self.conv10_2 = cm.ConvBlock((1, 1, 1024, 512), activate_type="mish") self.res_blocks5 = [] for _ in range(4): self.res_blocks5.append(cm.ResBlock(512, 512, 512, activate_type="mish")) self.conv10_3 = cm.ConvBlock((1, 1, 512, 512), activate_type="mish") self.conv_last = cm.ConvBlock((1, 1, 1024, 1024), activate_type="mish")
def YOLOv4_head(input_data, NUM_CLASS, bbox_per_head=2): shape1, shape2, shape3 = input_data route1 = tf.keras.layers.Input(shape=shape1[1:], name="head1") route2 = tf.keras.layers.Input(shape=shape2[1:], name="head2") route3 = tf.keras.layers.Input(shape=shape3[1:], name="head3") conv_small_box = cm.ConvBlock((3, 3, 128, 256))(route1) confidence = cm.ConvBlock((3, 3, 256, bbox_per_head), activate=True, activate_type="sigmoid", bn=False)(conv_small_box) box_params = cm.ConvBlock((3, 3, 256, bbox_per_head * 4), activate=True, activate_type="partial_sigmoid", bn=False)(conv_small_box) classes = [] for i in range(bbox_per_head): classes.append( cm.ConvBlock((3, 3, 256, NUM_CLASS), activate=True, activate_type="softmax", bn=False)(conv_small_box)) conv_small_box = tf.keras.layers.concatenate( [confidence, box_params, tf.concat(classes, axis=-1)], axis=-1, name="output_1") x = cm.ConvBlock((3, 3, 128, 256), downsample=True)(route1) x = tf.concat([x, route2], axis=-1) x = cm.ConvBlock((1, 1, 512, 256))(x) x = cm.ConvBlock((3, 3, 256, 512))(x) x = cm.ConvBlock((1, 1, 512, 256))(x) x = cm.ConvBlock((3, 3, 256, 512))(x) x = cm.ConvBlock((1, 1, 512, 256))(x) skip = x conv_medium_box = cm.ConvBlock((3, 3, 256, 512))(x) confidence = cm.ConvBlock((3, 3, 512, bbox_per_head), activate=True, activate_type="sigmoid", bn=False)(conv_medium_box) box_params = cm.ConvBlock((3, 3, 512, bbox_per_head * 4), activate=True, activate_type="partial_sigmoid", bn=False)(conv_medium_box) classes = [] for i in range(bbox_per_head): classes.append( cm.ConvBlock((3, 3, 512, NUM_CLASS), activate=True, activate_type="softmax", bn=False)(conv_medium_box)) conv_medium_box = tf.keras.layers.concatenate( [confidence, box_params, tf.concat(classes, axis=-1)], axis=-1, name="output_2") x = cm.ConvBlock((3, 3, 256, 512), downsample=True)(skip) x = tf.concat([x, route3], axis=-1) x = cm.ConvBlock((1, 1, 1024, 512))(x) x = cm.ConvBlock((3, 3, 512, 1024))(x) x = cm.ConvBlock((1, 1, 1024, 512))(x) x = cm.ConvBlock((3, 3, 512, 1024))(x) x = cm.ConvBlock((1, 1, 1024, 512))(x) conv_big_box = cm.ConvBlock((3, 3, 512, 1024))(x) confidence = cm.ConvBlock((3, 3, 1024, bbox_per_head), activate=True, activate_type="sigmoid", bn=False)(conv_big_box) box_params = cm.ConvBlock((3, 3, 1024, bbox_per_head * 4), activate=True, activate_type="partial_sigmoid", bn=False)(conv_big_box) classes = [] for i in range(bbox_per_head): classes.append( cm.ConvBlock((3, 3, 1024, NUM_CLASS), activate=True, activate_type="softmax", bn=False)(conv_big_box)) conv_big_box = tf.keras.layers.concatenate( [confidence, box_params, tf.concat(classes, axis=-1)], axis=-1, name="output_3") return tf.keras.Model(inputs=[route1, route2, route3], outputs=(conv_small_box, conv_medium_box, conv_big_box), name="Heads")
def __init__(self, sr_n_resblocks=6, dm_n_resblock=6, sr_n_feats=64, dm_n_feats=64, scale=2, denoise=True, bias=True, norm_type=False, act_type='relu', block_type='rrdb'): super(NET, self).__init__() # act = nn.LeakyReLU(negative_slope=0.1, inplace=False) # act = nn.PReLU(n_feats) # define sr module if denoise: m_sr_head = [ common.ConvBlock(5, sr_n_feats, 5, act_type=act_type, bias=True) ] else: m_sr_head = [ common.ConvBlock(4, sr_n_feats, 5, act_type=act_type, bias=True) ] if block_type.lower() == 'rrdb': m_sr_resblock = [ common.RRDB(sr_n_feats, sr_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(sr_n_resblocks) ] elif block_type.lower() == 'dudb': m_sr_resblock = [ common.DUDB(sr_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(sr_n_resblocks) ] elif block_type.lower() == 'res': m_sr_resblock = [ common.ResBlock(sr_n_feats, 3, norm_type, act_type, res_scale=1, bias=bias) for _ in range(sr_n_resblocks) ] else: raise RuntimeError('block_type is not supported') m_sr_resblock += [ common.ConvBlock(sr_n_feats, sr_n_feats, 3, bias=bias) ] m_sr_up = [ common.Upsampler(scale, sr_n_feats, norm_type, act_type, bias=bias), common.ConvBlock(sr_n_feats, 4, 3, bias=True) ] # branch for sr_raw output m_sr_tail = [nn.PixelShuffle(2)] # define demosaick module m_dm_head = [ common.ConvBlock(4, dm_n_feats, 5, act_type=act_type, bias=True) ] if block_type.lower() == 'rrdb': m_dm_resblock = [ common.RRDB(dm_n_feats, dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(dm_n_resblock) ] elif block_type.lower() == 'dudb': m_dm_resblock = [ common.DUDB(dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(dm_n_resblock) ] elif block_type.lower() == 'res': m_dm_resblock = [ common.ResBlock(dm_n_feats, 3, norm_type, act_type, res_scale=1, bias=bias) for _ in range(dm_n_resblock) ] else: raise RuntimeError('block_type is not supported') m_dm_resblock += [ common.ConvBlock(dm_n_feats, dm_n_feats, 3, bias=bias) ] m_dm_up = [ common.Upsampler(2, dm_n_feats, norm_type, act_type, bias=bias), common.ConvBlock(dm_n_feats, 3, 3, bias=True) ] self.model_sr = nn.Sequential( *m_sr_head, common.ShortcutBlock(nn.Sequential(*m_sr_resblock)), *m_sr_up) self.sr_output = nn.Sequential(*m_sr_tail) self.model_dm = nn.Sequential( *m_dm_head, common.ShortcutBlock(nn.Sequential(*m_dm_resblock)), *m_dm_up) for m in self.modules(): # pdb.set_trace() if isinstance(m, nn.Conv2d): # Xavier # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') nn.init.xavier_normal_(m.weight) m.weight.requires_grad = True if m.bias is not None: m.bias.data.zero_() m.bias.requires_grad = True
def YOLOv4_body(input_shape): inp = tf.keras.layers.Input(shape=input_shape) route1, route2, x = back.CSPDarkNet53()(inp) x = cm.ConvBlock((1, 1, 1024, 512))(x) x = cm.ConvBlock((3, 3, 512, 1024))(x) x = cm.ConvBlock((1, 1, 1024, 512))(x) x = SPP()(x) x = cm.ConvBlock((1, 1, 2048, 512))(x) x = cm.ConvBlock((3, 3, 512, 1024))(x) x = cm.ConvBlock((1, 1, 1024, 512), name="neck3")(x) route3 = x x = cm.ConvBlock((1, 1, 512, 256))(route3) upsampled = L.UpSampling2D()(x) x = cm.ConvBlock((1, 1, 512, 256))(route2) x = tf.concat([x, upsampled], axis=-1) x = cm.ConvBlock((1, 1, 512, 256))(x) x = cm.ConvBlock((3, 3, 256, 512))(x) x = cm.ConvBlock((1, 1, 512, 256))(x) x = cm.ConvBlock((3, 3, 256, 512))(x) x = cm.ConvBlock((1, 1, 512, 256), name="neck2")(x) route2 = x x = cm.ConvBlock((1, 1, 256, 128))(route2) upsampled = L.UpSampling2D()(x) x = cm.ConvBlock((1, 1, 256, 128))(route1) x = tf.concat([x, upsampled], axis=-1) x = cm.ConvBlock((1, 1, 256, 128))(x) x = cm.ConvBlock((3, 3, 128, 256))(x) x = cm.ConvBlock((1, 1, 256, 128))(x) x = cm.ConvBlock((3, 3, 128, 256))(x) x = cm.ConvBlock((1, 1, 256, 128), name="neck1")(x) route1 = x return keras.Model(inputs=[inp], outputs=(route1, route2, route3), name="Body")
def __init__(self, opt): super(NET, self).__init__() sr_n_resblocks = opt.sr_n_resblocks dm_n_resblocks = opt.dm_n_resblocks sr_n_feats = opt.channels dm_n_feats = opt.channels scale = opt.scale denoise = opt.denoise block_type = opt.block_type act_type = opt.act_type bias = opt.bias norm_type = opt.norm_type # define sr module if denoise: m_sr_head = [ common.ConvBlock(6, sr_n_feats, 5, act_type=act_type, bias=True) ] else: m_sr_head = [ common.ConvBlock(4, sr_n_feats, 5, act_type=act_type, bias=True) ] if block_type.lower() == 'rrdb': m_sr_resblock = [ common.RRDB(sr_n_feats, sr_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(sr_n_resblocks) ] elif block_type.lower() == 'dudb': m_sr_resblock = [ common.DUDB(sr_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(sr_n_resblocks) ] elif block_type.lower() == 'res': m_sr_resblock = [ common.ResBlock(sr_n_feats, 3, norm_type, act_type, res_scale=1, bias=bias) for _ in range(sr_n_resblocks) ] else: raise RuntimeError('block_type is not supported') m_sr_resblock += [ common.ConvBlock(sr_n_feats, sr_n_feats, 3, bias=bias) ] m_sr_up = [ common.Upsampler(scale, sr_n_feats, norm_type, act_type, bias=bias), common.ConvBlock(sr_n_feats, 4, 3, bias=True) ] # branch for sr_raw output m_sr_tail = [nn.PixelShuffle(2)] # define demosaick module m_dm_head = [ common.ConvBlock(4, dm_n_feats, 5, act_type=act_type, bias=True) ] if block_type.lower() == 'rrdb': m_dm_resblock = m_res( opt) #[common.RRDB(dm_n_feats, dm_n_feats, 3, #1, bias, norm_type, act_type, 0.2) #for _ in range(dm_n_resblocks)] elif block_type.lower() == 'dudb': m_dm_resblock = [ common.DUDB(dm_n_feats, 3, 1, bias, norm_type, act_type, 0.2) for _ in range(dm_n_resblocks) ] elif block_type.lower() == 'res': m_dm_resblock = [ common.ResBlock(dm_n_feats, 3, norm_type, act_type, res_scale=1, bias=bias) for _ in range(dm_n_resblocks) ] else: raise RuntimeError('block_type is not supported') #m_dm_resblock += [common.ConvBlock(dm_n_feats, dm_n_feats, 3, bias=bias)] m_dm_up = [ common.Upsampler(2, dm_n_feats, norm_type, act_type, bias=bias) ] #common.ConvBlock(dm_n_feats, 3, 3, bias=True)] self.model_sr = nn.Sequential( *m_sr_head, common.ShortcutBlock(nn.Sequential(*m_sr_resblock)), *m_sr_up) self.sr_output = nn.Sequential(*m_sr_tail) self.model_dm1 = nn.Sequential(*m_dm_head) self.model_dm2 = m_dm_resblock self.model_dm3 = nn.Sequential(*m_dm_up) greenresblock = green_res(opt) self.green = greenresblock for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_normal_(m.weight) m.weight.requires_grad = True if m.bias is not None: m.bias.data.zero_() m.bias.requires_grad = True #self.sft = SFTLayer() self.combine = nn.Sequential( common.ConvBlock(dm_n_feats + 1, dm_n_feats, 1, bias=True), nn.LeakyReLU(0.2, inplace=True)) self.greenup = nn.Sequential(common.ConvBlock(1, 4, 1, bias=True), nn.LeakyReLU(0.2, inplace=True), common.ConvBlock(4, 8, 1, bias=True), nn.LeakyReLU(0.2, inplace=True)) self.pac = PacConvTranspose2d(64, 64, kernel_size=5, stride=2, padding=2, output_padding=1) self.final = common.ConvBlock(dm_n_feats, 3, 3, bias=True) self.norm = nn.InstanceNorm2d(1)
def __init__(self, opt): super(NET, self).__init__() act_type = opt.act_type bias = opt.bias norm_type = opt.norm_type self.conv1 = nn.Sequential( common.ConvBlock(5, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 12, 3, act_type=act_type, bias=True), ) self.conv2 = nn.Sequential( common.ConvBlock(6, 64, 3, act_type=act_type, bias=True), common.ConvBlock(64, 3, 3, act_type=act_type, bias=True)) self.ps = nn.PixelShuffle(2) self.size = 64 s = self.size mask1 = torch.zeros(s, s) mask2 = torch.zeros(s, s) mask3 = torch.zeros(s, s) for i in range(0, s): for j in range(0, s): if i % 2 == 1 and j % 2 == 0: mask1[i, j] = 1 for i in range(0, s): for j in range(0, s): if (i % 2 == 0 and j % 2 == 0) or (i % 2 == 1 and j % 2 == 1): mask2[i, j] = 1 for i in range(0, s): for j in range(0, s): if i % 2 == 0 and j % 2 == 1: mask3[i, j] = 1 self.mask1 = mask1.cuda() self.mask2 = mask2.cuda() self.mask3 = mask3.cuda() self.srcnn = srcnn() self.carn = carn() self.edsr = EDSR()