def define_pairD(opt):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_D']
which_model = opt_net['which_model_pairD']
if which_model == 'discriminator_vgg_128':
netD = arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'],
nf=opt_net['nf'])
elif which_model == 'dis_acd': # sft-gan, Auxiliary Classifier Discriminator
netD = sft_arch.ACD_VGG_BN_96()
elif which_model == 'discriminator_vgg_96':
netD = arch.Discriminator_VGG_96(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_192':
netD = arch.Discriminator_VGG_192(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_48':
netD = arch.Discriminator_VGG_48(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_128_SN':
netD = arch.Discriminator_VGG_128_SN()
elif which_model == 'discriminator_vgg_128':
netD = arch.Discriminator_VGG_128()
elif which_model == 'discriminator_patch':
netD = arch.NLayerDiscriminator(opt_net['in_nc'], opt_net['nf'],
opt_net['n_layers'])
else:
raise NotImplementedError(
'Discriminator model [{:s}] not recognized'.format(which_model))
init_weights(netD, init_type='kaiming', scale=1)
if gpu_ids:
netD = nn.DataParallel(netD)
return netD
def define_D_grad(opt):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_D']
which_model = opt_net['which_model_D']
if which_model == 'discriminator_vgg_128':
netD = arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_96':
netD = arch.Discriminator_VGG_96(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_192':
netD = arch.Discriminator_VGG_192(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_128_SN':
netD = arch.Discriminator_VGG_128_SN()
else:
raise NotImplementedError(
'Discriminator model [{:s}] not recognized'.format(which_model))
init_weights(netD, init_type='kaiming', scale=1)
netD = nn.DataParallel(netD)
return netD
def define_D(opt):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_D']
which_model = opt_net['which_model_D']
which_model_G = opt_net['which_model_G']
if which_model_G == 'ppon':
model_G = 'PPON'
else:
model_G = 'ESRGAN'
if which_model == 'discriminator_vgg_128':
netD = arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'], convtype=opt_net['convtype'], arch=model_G)
elif which_model == 'dis_acd': # sft-gan, Auxiliary Classifier Discriminator
netD = sft_arch.ACD_VGG_BN_96()
elif which_model == 'discriminator_vgg_96':
netD = arch.Discriminator_VGG_96(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'], convtype=opt_net['convtype'], arch=model_G)
elif which_model == 'discriminator_vgg_192' or which_model == 'discriminator_192':
netD = arch.Discriminator_VGG_192(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'], convtype=opt_net['convtype'], arch=model_G)
elif which_model == 'discriminator_vgg_256' or which_model == 'discriminator_256':
netD = arch.Discriminator_VGG_256(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'], convtype=opt_net['convtype'], arch=model_G)
elif which_model == 'discriminator_vgg_128_SN':
netD = arch.Discriminator_VGG_128_SN()
else:
raise NotImplementedError('Discriminator model [{:s}] not recognized'.format(which_model))
init_weights(netD, init_type='kaiming', scale=1)
if gpu_ids:
netD = nn.DataParallel(netD)
return netD
def define_D(opt):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_D']
which_model = opt_net['which_model_D']
if which_model == 'discriminator_vgg_128':
netD = arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_96':
netD = arch.Discriminator_VGG_96(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_192':
netD = arch.Discriminator_VGG_192(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_feature':
netD = arch.Discriminator_VGG_Feature(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_128_SN':
netD = arch.Discriminator_VGG_128_SN(in_nc=opt_net['in_nc'])
elif which_model == 'discriminator_vgg_96_SN':
netD = arch.Discriminator_VGG_96_SN(in_nc=opt_net['in_nc'])
elif which_model == 'discriminator_vgg_feature_SN':
netD = arch.Discriminator_VGG_Feature_SN(in_nc=opt_net['in_nc'])
elif which_model == 'discriminator_SNResNetProjection_96':
netD = arch.Discriminator_SNResNetProjection_96(
in_nc=opt_net['in_nc'], num_features=opt_net['nf'])
elif which_model == 'discriminator_SNResNetProjection_feature':
netD = arch.Discriminator_SNResNetProjection_Feature(
in_nc=opt_net['in_nc'], num_features=opt_net['nf'])
elif which_model == 'discriminator_feature':
netD = arch.Discriminator_Feature(nb=opt_net['nb'])
elif which_model == 'discriminator_Sphere20a':
netD = arch.Discriminator_Sphere20a(model_path=opt_net['sphere_path'])
else:
raise NotImplementedError(
'Discriminator model [{:s}] not recognized'.format(which_model))
if which_model != 'discriminator_Sphere20a':
init_weights(netD, init_type='kaiming', scale=1)
if opt['distributed']:
assert torch.cuda.is_available()
netD = nn.parallel.DistributedDataParallel(netD)
elif gpu_ids:
assert torch.cuda.is_available()
netD = nn.DataParallel(netD)
return netD
def define_D(opt):
gpu_ids = opt["gpu_ids"]
opt_net = opt["network_D"]
which_model = opt_net["which_model_D"]
if which_model == "discriminator_vgg_128":
netD = arch.Discriminator_VGG_128(
in_nc=opt_net["in_nc"],
base_nf=opt_net["nf"],
norm_type=opt_net["norm_type"],
mode=opt_net["mode"],
act_type=opt_net["act_type"],
)
elif which_model == "dis_acd": # sft-gan, Auxiliary Classifier Discriminator
netD = sft_arch.ACD_VGG_BN_96()
elif which_model == "discriminator_vgg_96":
netD = arch.Discriminator_VGG_96(
in_nc=opt_net["in_nc"],
base_nf=opt_net["nf"],
norm_type=opt_net["norm_type"],
mode=opt_net["mode"],
act_type=opt_net["act_type"],
)
elif which_model == "discriminator_vgg_192":
netD = arch.Discriminator_VGG_192(
in_nc=opt_net["in_nc"],
base_nf=opt_net["nf"],
norm_type=opt_net["norm_type"],
mode=opt_net["mode"],
act_type=opt_net["act_type"],
)
elif which_model == "discriminator_vgg_128_SN":
netD = arch.Discriminator_VGG_128_SN()
else:
raise NotImplementedError(
"Discriminator model [{:s}] not recognized".format(which_model))
init_weights(netD, init_type="kaiming", scale=1)
if gpu_ids:
netD = nn.DataParallel(netD)
return netD
def define_D(opt,CEM=None):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_D']
which_model = opt_net['which_model_D']
input_patch_size = opt['datasets']['train']['patch_size']
# in_nc = opt_net['in_nc']*(2 if opt['network_D']['decomposed_input'] else 1)
in_nc = opt_net['in_nc']
assert not ((opt_net['pre_clipping'] or opt_net['decomposed_input']) and which_model!='PatchGAN'),'Unsupported yet'
if CEM is not None:
input_patch_size -= 2*CEM.invalidity_margins_HR
if which_model == 'discriminator_vgg_128':
kwargs = {}
if 'num_2_strides' in opt_net:
kwargs['num_2_strides'] = opt_net['num_2_strides']
netD = arch.Discriminator_VGG_128(in_nc=in_nc, base_nf=opt_net['nf'], nb=opt_net['n_layers'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'],input_patch_size=input_patch_size,**kwargs)
else:
raise NotImplementedError('Discriminator model [{:s}] not recognized'.format(which_model))
init_weights(netD, init_type='kaiming', scale=1)
if gpu_ids:
netD = nn.DataParallel(netD)
return netD
def define_D(opt, CEM=None, **kwargs):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_D']
which_model = opt_net['which_model_D']
input_patch_size = opt['datasets']['train']['patch_size']
# in_nc = opt_net['in_nc']*(2 if opt['network_D']['decomposed_input'] else 1)
in_nc = opt_net['in_nc']
assert not ((opt_net['pre_clipping'] or opt_net['decomposed_input'])
and which_model != 'PatchGAN'), 'Unsupported yet'
if CEM is not None:
input_patch_size -= 2 * CEM.invalidity_margins_HR
if which_model == 'discriminator_vgg_128':
kwargs = {}
if 'num_2_strides' in opt_net:
kwargs['num_2_strides'] = opt_net['num_2_strides']
netD = arch.Discriminator_VGG_128(in_nc=in_nc,
base_nf=opt_net['nf'],
nb=opt_net['n_layers'],
norm_type=opt_net['norm_type'],
mode=opt_net['mode'],
act_type=opt_net['act_type'],
input_patch_size=input_patch_size,
**kwargs)
elif which_model == 'discriminator_vgg_128_nonModified':
netD = arch.Discriminator_VGG_128_nonModified(in_nc=in_nc,
nf=opt_net['nf'])
elif which_model == 'dis_acd': # sft-gan, Auxiliary Classifier Discriminator
netD = sft_arch.ACD_VGG_BN_96()
elif which_model == 'PatchGAN':
norm_layer = functools.partial(nn.InstanceNorm2d,
affine=False,
track_running_stats=False)
netD = arch.PatchGAN_Discriminator(input_nc=in_nc,
opt_net=opt_net,
ndf=opt_net['nf'],
n_layers=opt_net['n_layers'],
norm_layer=norm_layer)
elif which_model == 'discriminator_vgg_96':
netD = arch.Discriminator_VGG_96(in_nc=in_nc, base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_192':
netD = arch.Discriminator_VGG_192(in_nc=in_nc, base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_128_SN':
netD = arch.Discriminator_VGG_128_SN()
elif 'DnCNN_D' in which_model:
chroma_mode = kwargs['chroma_mode'] if 'chroma_mode' in kwargs.keys(
) else False
opt_net_G = opt['network_G']
assert opt_net['DCT_D'] == 1
G_in_nc = opt['scale']**2 + 2 * 64 if chroma_mode else 64
G_out_nc = 2 * (opt['scale']**2) if chroma_mode else 64
# Even when not in concat_inpiut mode, I'm supplying D with channel Y, so it does not need to determine realness based only on the chroma channels
D_input_channels = G_in_nc + G_out_nc if opt_net['concat_input'] else (
opt['scale']**2 + G_out_nc if chroma_mode else G_out_nc)
num_latent_channels = None
if opt_net['inject_Z']:
num_latent_channels = opt_net_G['latent_channels']
# D_input_channels += num_latent_channels
netD = arch.DnCNN(
n_channels=opt_net_G['nf']
if opt_net['nf'] is None else opt_net['nf'],
depth=opt_net_G['nb'],
in_nc=D_input_channels,
norm_type='layer' if (opt['train']['gan_type'] == 'wgan-gp'
and opt_net_G['norm_type'] == 'batch') else
opt_net_G['norm_type'],
discriminator=True,
expected_input_size=opt['datasets']['train']['patch_size'] //
opt['scale'],
latent_input=opt_net_G['latent_input'],
num_latent_channels=num_latent_channels,
chroma_generator=False,
spectral_norm='sn' in opt['train']['gan_type'],
pooling_no_FC=opt_net['pooling_no_fc'])
else:
raise NotImplementedError(
'Discriminator model [{:s}] not recognized'.format(which_model))
init_weights(netD, init_type='kaiming', scale=1)
if gpu_ids:
netD = nn.DataParallel(netD)
return netD
