def get_nets_dis(path, which_epoch='latest'):
gpu_ids = [0]
Tensor = torch.cuda.FloatTensor
opt = util.load_opt(path)
# assume caffe style model
opt.caffe = False
netD_A = networks.define_D(opt.output_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
False,
opt.init_type,
gpu_ids,
opt=opt)
netD_B = networks.define_D(opt.input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
False,
opt.init_type,
gpu_ids,
opt=opt)
load_network_with_path(netD_A, 'D_A', which_epoch, path)
load_network_with_path(netD_B, 'D_B', which_epoch, path)
netD_A.cuda()
netD_B.cuda()
return {'A': netD_A, 'B': netD_B}
def initialize_networks(self, opt):
self.netGA = networks.define_G(opt, opt['netGA'])
self.netGB = networks.define_G(opt, opt['netGB'])
self.netDA = networks.define_D(opt, opt['netDA'])
self.netDB = networks.define_D(opt, opt['netDB'])
self.netEA, self.netHairA = networks.define_RES(
opt, opt['input_nc_A'], opt['netEDA'])
self.netEB, self.netHairB = networks.define_RES(
opt, opt['input_nc_B'], opt['netEDB'])
if self.opt['pretrain']:
self.train_nets = [
self.netGA, self.netGB, self.netDA, self.netDB, self.netEA,
self.netHairA, self.netEB, self.netHairB
]
else:
self.train_nets = [self.netEA, self.netHairA]
# set require gradients
if self.isTrain:
self.set_requires_grad(self.train_nets, True)
else:
self.set_requires_grad(self.train_nets, False)
if self.use_gpu:
for i in range(len(self.train_nets)):
self.train_nets[i] = DataParallelWithCallback(
self.train_nets[i], device_ids=opt['gpu_ids'])
if self.opt['pretrain']:
self.netGA, self.netGB, self.netDA, self.netDB, self.netEA, \
self.netHairA, self.netEB, self.netHairB = self.train_nets
else:
self.netEA, self.netHairA = self.train_nets
def initialize_networks(self, opt):
netG = networks.define_G(opt)
# netD = networks.define_D(opt) if opt.isTrain else None
if opt.isTrain:
opt.label_nc = opt.label_nc-1
netD = networks.define_D(opt)
else:
netD_fine = None
netE = networks.define_E(opt) if opt.use_vae else None
if opt.isTrain:
opt.label_nc = (opt.label_nc+1)
netD_fine = networks.define_D(opt)
else:
netD_fine = None
if not opt.isTrain or opt.continue_train:
netG = util.load_network(netG, 'G', opt.which_epoch, opt)
if opt.isTrain:
netD = util.load_network(netD, 'D', opt.which_epoch, opt)
netD_fine = util.load_network(netD_fine, 'D', opt.which_epoch, opt)
else:
netD = None
netD_fine = None
if opt.use_vae:
netE = util.load_network(netE, 'E', opt.which_epoch, opt)
return netG, netD, netE, netD_fine
def initialize_networks(self, opt):
netG = networks.define_G(opt)
netD = networks.define_D(opt) if opt.isTrain else None
netD_rotate = networks.define_D(opt) if opt.isTrain else None
netE = networks.define_E(opt) if opt.use_vae else None
pretrained_path = ''
if not opt.isTrain or opt.continue_train:
self.load_network(netG, 'G', opt.which_epoch, pretrained_path)
if opt.isTrain and not opt.noload_D:
self.load_network(netD, 'D', opt.which_epoch, pretrained_path)
self.load_network(netD_rotate, 'D_rotate', opt.which_epoch,
pretrained_path)
if opt.use_vae:
self.load_network(netE, 'E', opt.which_epoch, pretrained_path)
else:
if opt.load_separately:
netG = self.load_separately(netG, 'G', opt)
if not opt.noload_D:
netD = self.load_separately(netD, 'D', opt)
netD_rotate = self.load_separately(netD_rotate, 'D_rotate',
opt)
if opt.use_vae:
netE = self.load_separately(netE, 'E', opt)
return netG, netD, netE, netD_rotate
def initialize_networks(self, opt, end2end=False, triple=False):
if opt.end2endtri:
netG_1 = networks.define_G(opt, triple)
netD_1 = networks.define_D(opt, triple)
netG = networks.define_G(opt)
netD = networks.define_D(opt) if opt.isTrain else None
netE = networks.define_E(opt) if opt.use_vae else None
if not opt.isTrain or opt.continue_train:
if opt.end2endtri:
netG_1 = util.load_network(netG_1, 'G', opt.which_triple_epoch,
opt, triple)
netG = util.load_network(netG, 'G', opt.which_epoch, opt)
if opt.isTrain: # and end2end:
netD = util.load_network(netD, 'D', opt.which_epoch, opt)
if opt.end2endtri:
netD_1 = util.load_network(netD_1, 'D',
opt.which_triple_epoch, opt,
triple)
if opt.use_vae:
netE = util.load_network(netE, 'E', opt.which_epoch, opt)
if not opt.end2endtri:
netG_1 = None
netD_1 = None
return netG, netD, netE, netG_1, netD_1
def __init__(self, opt, cur_stage):
BaseModel.__init__(self, opt)
self.loss_names = ["A", "B",
'D_A', 'D_B',
'IS_A', 'IS_B',
'reward_A', 'reward_B',
"adv_A", "adv_B",
"entropy_A", "entropy_B"]
visual_names_A = ['real_A', 'fake_B']
visual_names_B = ['real_B', 'fake_A']
self.visual_names = visual_names_A + visual_names_B
self.cur_stage = cur_stage
self.model_names = ["C_A", "C_B"]
self.ctrl_sample_batch = opt.ctrl_sample_batch
self.netC_A = Controller(opt, self.cur_stage)
self.netC_B = Controller(opt, self.cur_stage)
self.netD_A = networks.define_D(3, 64, "basic", norm='instance')
self.netD_B = networks.define_D(3, 64, "basic", norm='instance')
load_saves(self.netD_A, "res", "D_A", os.path.join(opt.path, "pre_mod"))
load_saves(self.netD_B, "res", "D_B", os.path.join(opt.path, "pre_mod"))
self.loss = networks.GANLoss("lsgan")
self.prev_hiddens_A = None
self.prev_archs_A = None
self.prev_hiddens_B = None
self.prev_archs_B = None
if len(self.gpu_ids) != 0:
self.cuda()
networks.init_weights(self.netC_A, opt.init_type, opt.init_gain)
networks.init_weights(self.netC_A, opt.init_type, opt.init_gain)
self.optimizers_names = ["A", "B"]
self.optimizerA = torch.optim.Adam(filter(lambda p: p.requires_grad, self.netC_A.parameters()),
opt.ctrl_lr, (0.0, 0.9))
self.optimizerB = torch.optim.Adam(filter(lambda p: p.requires_grad, self.netC_A.parameters()),
opt.ctrl_lr, (0.0, 0.9))
self.optimizers.append(self.optimizerA)
self.optimizers.append(self.optimizerB)
self.valid_dataloader = create_dataset(opt, valid=True)
self.baseline_decay = opt.baseline_decay
self.entropy_coeff = opt.entropy_coeff
self.netG_A = None
self.netG_B = None
self.baseline_A = None
self.baseline_B = None
def __init__(self, opt, device):
super(CycleGAN, self).__init__()
self.device = device
self.opt = opt
self.netG_A = networks.define_G(self.opt.input_nc, self.opt.output_nc,
self.opt.ngf, self.opt.netG,
self.opt.norm, self.opt.dropout,
self.opt.init_type, self.opt.init_gain,
self.opt.task_num,
self.opt.netG_A_filter_list)
self.netG_B = networks.define_G(self.opt.input_nc, self.opt.output_nc,
self.opt.ngf, self.opt.netG,
self.opt.norm, self.opt.dropout,
self.opt.init_type, self.opt.init_gain,
self.opt.task_num,
self.opt.netG_B_filter_list)
if opt.train:
self.netD_A = networks.define_D(self.opt.input_nc, self.opt.ndf,
self.opt.netD, self.opt.norm,
self.opt.init_type,
self.opt.init_gain)
self.netD_B = networks.define_D(self.opt.input_nc, self.opt.ndf,
self.opt.netD, self.opt.norm,
self.opt.init_type,
self.opt.init_gain)
self.fake_A_pool = ImageBuffer(
self.opt.pool_size
) # create image buffer to store previously generated images
self.fake_B_pool = ImageBuffer(
self.opt.pool_size
) # create image buffer to store previously generated images
self.criterionGAN = networks.GANLoss(self.opt.gan_mode).to(
self.device) # define GAN loss.
self.criterionCycle = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
self.optimizer_G = torch.optim.Adam(itertools.chain(
self.netG_A.parameters(), self.netG_B.parameters()),
lr=self.opt.lr,
betas=(self.opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(itertools.chain(
self.netD_A.parameters(), self.netD_B.parameters()),
lr=self.opt.lr,
betas=(self.opt.beta1, 0.999))
self.optimizers = []
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
self.schedulers = [
networks.get_scheduler(optimizer, opt)
for optimizer in self.optimizers
]
def __init__(self, opt):
"""Initialize the CycleGAN class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
visual_names_A = ['real_A', 'fake_B', 'rec_A']
visual_names_B = ['real_B', 'fake_A', 'rec_B']
if self.isTrain and self.opt.lambda_identity > 0.0: # if identity loss is used, we also visualize idt_B=G_A(B) ad idt_A=G_A(B)
visual_names_A.append('idt_B')
visual_names_B.append('idt_A')
self.visual_names = visual_names_A + visual_names_B # combine visualizations for A and B
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
if self.isTrain:
self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
else: # during test time, only load Gs
self.model_names = ['G_A', 'G_B']
# define networks (both Generators and discriminators)
# The naming is different from those used in the paper.
# Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
self.netG_B = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain: # define discriminators
self.netD_A = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
self.netD_B = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain:
if opt.lambda_identity > 0.0: # only works when input and output images have the same number of channels
assert(opt.input_nc == opt.output_nc)
self.fake_A_pool = ImagePool(opt.pool_size) # create image buffer to store previously generated images
self.fake_B_pool = ImagePool(opt.pool_size) # create image buffer to store previously generated images
# define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device) # define GAN loss.
self.criterionCycle = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
def define_networks(self, start_epoch):
opt = self.opt
# Generator network
input_nc = opt.label_nc if (opt.label_nc != 0 and not self.pose) else opt.input_nc
netG_input_nc = input_nc
opt.for_face = False
self.netG = networks.define_G(opt)
if self.refine_face:
opt_face = copy.deepcopy(opt)
opt_face.n_downsample_G -= 1
if opt_face.n_adaptive_layers > 0: opt_face.n_adaptive_layers -= 1
opt_face.input_nc = opt.output_nc
opt_face.fineSize = self.faceRefiner.face_size
opt_face.aspect_ratio = 1
opt_face.for_face = True
self.netGf = networks.define_G(opt_face)
# Discriminator network
if self.isTrain or opt.finetune:
netD_input_nc = input_nc + opt.output_nc + (1 if self.concat_fg_mask_for_D else 0)
if self.concat_ref_for_D:
netD_input_nc *= 2
self.netD = networks.define_D(opt, netD_input_nc, opt.ndf, opt.n_layers_D, opt.norm_D, opt.netD_subarch,
opt.num_D, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids)
if self.add_face_D:
self.netDf = networks.define_D(opt, opt.output_nc * 2, opt.ndf, opt.n_layers_D, opt.norm_D, 'n_layers',
1, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids)
else:
self.netDf = None
self.temporal = False
self.netDT = None
print('---------- Networks initialized -------------')
# initialize optimizers
if self.isTrain:
# optimizer G
params = list(self.netG.parameters())
if self.refine_face: params += list(self.netGf.parameters())
self.optimizer_G = self.get_optimizer(params, for_discriminator=False)
# optimizer D
params = list(self.netD.parameters())
if self.add_face_D: params += list(self.netDf.parameters())
self.optimizer_D = self.get_optimizer(params, for_discriminator=True)
print('---------- Optimizers initialized -------------')
# make model temporal by generating multiple frames
if (not opt.isTrain or start_epoch > opt.niter_single) and opt.n_frames_G > 1:
self.make_temporal_model()
def initialize_networks(self, opt):
netG_for_CT = networks.define_G(opt)
netD_aligned = networks.define_D(opt) if opt.isTrain else None
netG_for_MR = networks.define_G(opt)
netD_unaligned = networks.define_D(opt) if opt.isTrain else None
if not opt.isTrain or opt.continue_train:
netG_for_CT = util.load_network(netG_for_CT, 'G_for_CT', opt.which_epoch, opt)
netG_for_MR = util.load_network(netG_for_MR, 'G_for_MR', opt.which_epoch, opt)
if opt.isTrain:
netD_aligned = util.load_network(netD_aligned, 'D_aligned', opt.which_epoch, opt)
netD_unaligned = util.load_network(netD_unaligned, 'D_unaligned', opt.which_epoch, opt)
return netG_for_CT, netD_aligned, netG_for_MR, netD_unaligned
def __init__(self, args):
self.args = args
Tensor = torch.cuda.FloatTensor if args.gpu_ids else torch.Tensor
use_sigmoid = args.no_lsgan
# Global discriminator
self.netD = networks.define_D(args.input_nc, args.ndf,
args.which_model_netD, args.n_layers_D,
args.norm, use_sigmoid, args.gpu_ids)
# Local discriminator
self.netD_local = networks.define_D(args.input_nc, args.ndf,
args.which_model_netD,
args.n_layers_D, args.norm,
use_sigmoid, args.gpu_ids)
# Generator
self.netG = TransformerNet(args.norm, args.affine_state)
self.gan_loss = networks.GANLoss(use_lsgan=not args.no_lsgan,
tensor=Tensor)
self.identity_criterion = torch.nn.L1Loss()
# Resume
if args.resume_netG != '':
self.netG.load_state_dict(torch.load(args.resume_netG))
if args.resume_netD != '':
self.netD.load_state_dict(torch.load(args.resume_netD))
if args.resume_netD_local != '':
self.netD_local.load_state_dict(torch.load(args.resume_netD_local))
# optimizer
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
self.optimizer_D_local = torch.optim.Adam(self.netD_local.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
if self.args.cuda:
self.netD = self.netD.cuda()
self.netG = self.netG.cuda()
def init_temporal_model(self):
opt = self.opt
self.temporal = True
self.netG.init_temporal_network()
self.netG.cuda()
if opt.isTrain:
self.lossCollector.tD = min(opt.n_frames_D, opt.n_frames_G)
params = list(self.netG.parameters())
if self.refine_face: params += list(self.netGf.parameters())
self.optimizer_G = self.get_optimizer(params,
for_discriminator=False)
# temporal discriminator
self.netDT = networks.define_D(opt,
opt.output_nc *
self.lossCollector.tD,
opt.ndf,
opt.n_layers_D,
opt.norm_D,
'n_layers',
1,
not opt.no_ganFeat_loss,
gpu_ids=self.gpu_ids)
# optimizer D
params = list(self.netD.parameters()) + list(
self.netDT.parameters())
if self.add_face_D: params += list(self.netDf.parameters())
self.optimizer_D = self.get_optimizer(params,
for_discriminator=True)
Visualizer.vis_print(
self.opt,
'---------- Now start training multiple frames -------------')
def __init__(self, opt):
BaseModel.__init__(self, opt)
# specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake']
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
self.visual_names = ['real_A', 'fake_B', 'real_B']
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
if self.isTrain:
self.model_names = ['G', 'D']
else: # during test time, only load G
self.model_names = ['G']
# define networks (both generator and discriminator)
self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain: # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain:
# define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
self.criterionL1 = torch.nn.L1Loss()
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
def make_temporal_model(self):
opt = self.opt
self.temporal = True
self.netG.set_flow_prev()
self.netG.cuda()
if opt.isTrain:
self.lossCollector.tD = min(opt.n_frames_D, opt.n_frames_G)
if opt.finetune_all:
params = list(self.netG.parameters())
else:
train_names = ['flow_network_temp']
if opt.spade_combine:
train_names += ['img_warp_embedding', 'mlp_gamma3', 'mlp_beta3']
params, _ = self.get_train_params(self.netG, train_names)
if self.refine_face: params += list(self.netGf.parameters())
self.optimizer_G = self.get_optimizer(params, for_discriminator=False)
# temporal discriminator
self.netDT = networks.define_D(opt, opt.output_nc * self.lossCollector.tD, opt.ndf, opt.n_layers_D, opt.norm_D, 'n_layers',
1, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids)
# optimizer D
params = list(self.netD.parameters()) + list(self.netDT.parameters())
if self.add_face_D: params += list(self.netDf.parameters())
self.optimizer_D = self.get_optimizer(params, for_discriminator=True)
print('---------- Now start training multiple frames -------------')
def initialize_networks(self, opt):
net = {}
net['netG'] = networks.define_G(opt)
net['netD'] = networks.define_D(opt) if opt.isTrain else None
net['netCorr'] = networks.define_Corr(opt)
net['netDomainClassifier'] = networks.define_DomainClassifier(
opt) if opt.weight_domainC > 0 and opt.domain_rela else None
if not opt.isTrain or opt.continue_train:
net['netG'] = util.load_network(net['netG'], 'G', opt.which_epoch,
opt)
if opt.isTrain:
net['netD'] = util.load_network(net['netD'], 'D',
opt.which_epoch, opt)
net['netCorr'] = util.load_network(net['netCorr'], 'Corr',
opt.which_epoch, opt)
if opt.weight_domainC > 0 and opt.domain_rela:
net['netDomainClassifier'] = util.load_network(
net['netDomainClassifier'], 'DomainClassifier',
opt.which_epoch, opt)
if (not opt.isTrain) and opt.use_ema:
net['netG'] = util.load_network(net['netG'], 'G_ema',
opt.which_epoch, opt)
net['netCorr'] = util.load_network(net['netCorr'],
'netCorr_ema',
opt.which_epoch, opt)
return net
def __init__(self, args, logger):
super().__init__(args, logger)
# specify the training losses you want to print out. The program will call base_model.get_current_losses
self.loss_names = ['loss_G', 'loss_D']
# specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
self.model_names = ['G', 'D']
self.sample_names = ['fake_B', 'real_A', 'real_B']
# load/define networks
self.G = networks.define_G(args.input_nc, args.output_nc, args.ngf,
args.which_model_netG, args.norm, not args.no_dropout, args.init_type, args.init_gain, self.gpu_ids)
if not 'continue_train' in args:
use_sigmoid = args.no_lsgan
self.D = networks.define_D(args.input_nc + args.output_nc, args.ndf,
args.which_model_netD,
args.n_layers_D, args.norm, use_sigmoid, args.init_type, args.init_gain, self.gpu_ids)
self.fake_AB_pool = ImagePool(args.pool_size)
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not args.no_lsgan).to(self.device)
self.criterionL1 = torch.nn.L1Loss()
# initialize optimizers
self.optimizers = []
self.optimizer_G = torch.optim.Adam(self.G.parameters(),
lr=args.g_lr, betas=(args.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.D.parameters(),
lr=args.d_lr, betas=(args.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
def __init__(self, opt):
BaseModel.__init__(self, opt)
self.netP = networks.define_P(opt, weight_path=opt.parse_net_weight)
self.netG = networks.define_G(opt, use_norm='spectral_norm')
if self.isTrain:
self.netD = networks.define_D(opt, opt.Dinput_nc, use_norm='spectral_norm')
self.vgg_model = loss.PCPFeat(weight_path='./pretrain_models/vgg19-dcbb9e9d.pth').to(opt.device)
if len(opt.gpu_ids) > 0:
self.vgg_model = torch.nn.DataParallel(self.vgg_model, opt.gpu_ids, output_device=opt.device)
self.model_names = ['G']
self.loss_names = ['Pix', 'PCP', 'G', 'FM', 'D', 'SS'] # Generator loss, fm loss, parsing loss, discriminator loss
self.visual_names = ['img_LR', 'img_HR', 'img_SR', 'ref_Parse', 'hr_mask']
self.fm_weights = [1**x for x in range(opt.D_num)]
if self.isTrain:
self.model_names = ['G', 'D']
self.load_model_names = ['G', 'D']
self.criterionParse = torch.nn.CrossEntropyLoss().to(opt.device)
self.criterionFM = loss.FMLoss().to(opt.device)
self.criterionGAN = loss.GANLoss(opt.gan_mode).to(opt.device)
self.criterionPCP = loss.PCPLoss(opt)
self.criterionPix= nn.L1Loss()
self.criterionRS = loss.RegionStyleLoss()
self.optimizer_G = optim.Adam([p for p in self.netG.parameters() if p.requires_grad], lr=opt.g_lr, betas=(opt.beta1, 0.999))
self.optimizer_D = optim.Adam([p for p in self.netD.parameters() if p.requires_grad], lr=opt.d_lr, betas=(opt.beta1, 0.999))
self.optimizers = [self.optimizer_G, self.optimizer_D]
def initialize_networks(self, opt):
netG = networks.define_G(opt)
opt.input_nc = 2
netD1 = networks.define_D(opt) if opt.isTrain else None
opt.input_nc = 7
netD2 = networks.define_D(opt) if opt.isTrain else None
netE = networks.define_E(opt) if opt.use_vae else None
if not opt.isTrain or opt.continue_train:
netG = util.load_network(netG, 'G', opt.which_epoch, opt)
if opt.isTrain:
netD1 = util.load_network(netD1, 'D1', opt.which_epoch, opt)
netD2 = util.load_network(netD2, 'D2', opt.which_epoch, opt)
if opt.use_vae:
netE = util.load_network(netE, 'E', opt.which_epoch, opt)
return netG, netD1, netD2, netE
def initialize_networks(self, opt):
netG = networks.define_G(opt)
netD = networks.define_D(opt) if opt.isTrain else None
netD2 = networks.define_D(
opt) if opt.isTrain and opt.unpairTrain else None
netE = networks.define_E(
opt) if opt.use_vae else None # this is for original spade network
netIG = networks.define_IG(
opt
) if opt.use_ig else None # this is the orient inpainting network
netSIG = networks.define_SIG(
opt
) if opt.use_stroke else None # this is the stroke orient inpainting network
netFE = networks.define_FE(
opt
) if opt.use_instance_feat else None # this is the feat encoder from pix2pixHD
netB = networks.define_B(opt) if opt.use_blender else None
if not opt.isTrain or opt.continue_train:
# if the pth exist
save_filename = '%s_net_%s.pth' % (opt.which_epoch, 'G')
save_dir = os.path.join(opt.checkpoints_dir, opt.name)
G_path = os.path.join(save_dir, save_filename)
if os.path.exists(G_path):
netG = util.load_network(netG, 'G', opt.which_epoch, opt)
if opt.fix_netG:
netG.eval()
if opt.use_blender:
netB = util.load_blend_network(netB, 'B', opt.which_epoch,
opt)
if opt.isTrain:
netD = util.load_network(netD, 'D', opt.which_epoch, opt)
if opt.unpairTrain:
netD2 = util.load_network(netD2, 'D', opt.which_epoch,
opt)
if opt.use_vae:
netE = util.load_network(netE, 'E', opt.which_epoch, opt)
if opt.use_ig:
netIG = util.load_inpainting_network(netIG, opt)
netIG.eval()
if opt.use_stroke:
netSIG = util.load_sinpainting_network(netSIG, opt)
netSIG.eval()
return netG, netD, netE, netIG, netFE, netB, netD2, netSIG
def discriminate_sep(disc_model_path, image_a_path, image_b_path):
# create the discriminator
netD = networks.define_D(input_nc=6,
ndf=64,
netD="basic",
n_layers_D=3,
norm="batch",
init_type="normal",
init_gain=0.02,
gpu_ids=[0])
print('loading the model from %s' % disc_model_path)
device = torch.device('cuda:0')
state_dict = torch.load(disc_model_path, map_location=str(device))
if hasattr(state_dict, '_metadata'):
del state_dict._metadata
netD.module.load_state_dict(state_dict)
transform_list = []
method=Image.BICUBIC
# bring image to certain size
load_size = 286
osize = [load_size, load_size]
transform_list.append(transforms.Resize(osize, method))
# crop image to right dimensions
crop_size = 256
transform_list.append(transforms.RandomCrop(crop_size))
# transform image to tensor
transform_list += [transforms.ToTensor()]
# normalize image
transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
transform_data = transforms.Compose(transform_list)
img_A_tensor = transform_data(Image.open(image_a_path).convert('RGB'))
img_B_tensor = transform_data(Image.open(image_b_path).convert('RGB'))
# adding dimension
img_A_tensor = img_A_tensor.unsqueeze(0)
img_B_tensor = img_B_tensor.unsqueeze(0)
real_A = img_A_tensor.to(device)
real_B = img_B_tensor.to(device)
real_AB = torch.cat((real_A, real_B), 1)
pred = netD(real_AB.detach())
print("Image {}: discriminator: {}".format(image_a_path, pred.mean()))
def initialize_networks(self, opt):
netG = networks.define_G(opt)
netD = networks.define_D(opt) if opt.isTrain else None
if not opt.isTrain or opt.continue_train:
netG = util.load_network(netG, 'G', opt.which_epoch, opt)
if opt.isTrain:
netD = util.load_network(netD, 'D', opt.which_epoch, opt)
return netG, netD
def initialize(self, opt, net):
BaseModel.initialize(self, opt)
self.net = net.to(self.device)
self.edge_map = EdgeMap(scale=1).to(self.device)
if self.isTrain:
# define loss functions
self.vgg = losses.Vgg19(requires_grad=False).to(self.device)
self.loss_dic = losses.init_loss(opt, self.Tensor)
vggloss = losses.ContentLoss()
vggloss.initialize(losses.VGGLoss(self.vgg))
self.loss_dic['t_vgg'] = vggloss
cxloss = losses.ContentLoss()
if opt.unaligned_loss == 'vgg':
cxloss.initialize(
losses.VGGLoss(self.vgg, weights=[0.1], indices=[31]))
elif opt.unaligned_loss == 'ctx':
cxloss.initialize(
losses.CXLoss(self.vgg,
weights=[0.1, 0.1, 0.1],
indices=[8, 13, 22]))
elif opt.unaligned_loss == 'mse':
cxloss.initialize(nn.MSELoss())
elif opt.unaligned_loss == 'ctx_vgg':
cxloss.initialize(
losses.CXLoss(self.vgg,
weights=[0.1, 0.1, 0.1, 0.1],
indices=[8, 13, 22, 31],
criterions=[losses.CX_loss] * 3 +
[nn.L1Loss()]))
else:
raise NotImplementedError
self.loss_dic['t_cx'] = cxloss
# initialize optimizers
self.optimizer_G = torch.optim.Adam(self.net.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999),
weight_decay=opt.wd)
self._init_optimizer([self.optimizer_G])
# define discriminator
# if self.opt.lambda_gan > 0:
self.netD = networks.define_D(opt, 3)
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self._init_optimizer([self.optimizer_D])
if opt.no_verbose is False:
self.print_network()
def initialize_networks(self, opt):
netG = None
netD = None
netE = None
netV = None
netA = None
netA_sync = None
if opt.train_recognition:
netV = networks.define_V(opt)
elif opt.train_sync:
netA_sync = networks.define_A_sync(opt) if opt.use_audio else None
netE = networks.define_E(opt)
else:
netG = networks.define_G(opt)
netA = networks.define_A(
opt) if opt.use_audio and opt.use_audio_id else None
netA_sync = networks.define_A_sync(opt) if opt.use_audio else None
netE = networks.define_E(opt)
netV = networks.define_V(opt)
if opt.isTrain:
netD = networks.define_D(opt)
if not opt.isTrain or opt.continue_train:
self.load_network(netG, 'G', opt.which_epoch)
self.load_network(netV, 'V', opt.which_epoch)
self.load_network(netE, 'E', opt.which_epoch)
if opt.use_audio:
if opt.use_audio_id:
self.load_network(netA, 'A', opt.which_epoch)
self.load_network(netA_sync, 'A_sync', opt.which_epoch)
if opt.isTrain and not opt.noload_D:
self.load_network(netD, 'D', opt.which_epoch)
# self.load_network(netD_rotate, 'D_rotate', opt.which_epoch, pretrained_path)
else:
if self.opt.pretrain:
if opt.netE == 'fan':
netE.load_pretrain()
netV.load_pretrain()
if opt.load_separately:
netG = self.load_separately(netG, 'G', opt)
netA = self.load_separately(
netA, 'A',
opt) if opt.use_audio and opt.use_audio_id else None
netA_sync = self.load_separately(
netA_sync, 'A_sync', opt) if opt.use_audio else None
netV = self.load_separately(netV, 'V', opt)
netE = self.load_separately(netE, 'E', opt)
if not opt.noload_D:
netD = self.load_separately(netD, 'D', opt)
return netG, netD, netA, netA_sync, netV, netE
def load_checkpoint(opt):
"""Loads the generator and discriminator models from checkpoints.
"""
use_dropout = not opt.no_dropout
netG_A = define_G(opt.input_nc, opt.output_nc, opt.ndf, opt.which_model_netG, opt.norm, use_dropout)
netG_B = define_G(opt.output_nc, opt.input_nc, opt.ndf, opt.which_model_netG, opt.norm, use_dropout)
use_sigmoid = opt.no_lsgan
netD_A = define_D(opt.output_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, opt.norm, use_sigmoid)
netD_B = define_D(opt.input_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, opt.norm, use_sigmoid)
checkpoint_file = "%d_checkpoint_ep%d" % (opt.checkpoint_epoch, opt.checkpoint_epoch)
checkpoint = torch.load(os.path.join(opt.save_dir,checkpoint_file))
netG_A.load_state_dict(checkpoint['netG_A'])
netG_B.load_state_dict(checkpoint['netG_B'])
netD_A.load_state_dict(checkpoint['netD_A'])
netD_B.load_state_dict(checkpoint['netD_B'])
# start_epoch = checkpoint['epoch'] + 1
return netG_A, netG_B, netD_A, netD_B
def load_weight(self, pathlist: dict):
self.net_Gs = []
self.net_Ds = []
for weight in pathlist['net_G']:
net_G = define_G(self.opt).to(self.device)
net_G.load_state_dict(torch.load(weight, map_location=self.device))
self.net_Gs.append(net_G)
for weight in pathlist['net_D']:
net_D = define_D(self.opt).to(self.device)
net_D.load_state_dict(torch.load(weight, map_location=self.device))
self.net_Ds.append(net_D)
def __init__(self, opt):
assert opt.isTrain
opt = copy.deepcopy(opt)
if len(opt.gpu_ids) > 0:
opt.gpu_ids = opt.gpu_ids[:1]
self.gpu_ids = opt.gpu_ids
super(SPADEModelModules, self).__init__()
self.opt = opt
self.model_names = ['G_student', 'G_teacher', 'D']
teacher_opt = self.create_option('teacher')
self.netG_teacher = networks.define_G(opt.teacher_netG,
gpu_ids=self.gpu_ids,
opt=teacher_opt)
student_opt = self.create_option('student')
self.netG_student = networks.define_G(opt.student_netG,
init_type=opt.init_type,
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=student_opt)
if hasattr(opt, 'distiller'):
pretrained_opt = self.create_option('pretrained')
self.netG_pretrained = networks.define_G(opt.pretrained_netG,
gpu_ids=self.gpu_ids,
opt=pretrained_opt)
self.netD = networks.define_D(opt.netD,
init_type=opt.init_type,
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=opt)
self.mapping_layers = ['head_0', 'G_middle_1', 'up_1']
self.netAs = nn.ModuleList()
for i, mapping_layer in enumerate(self.mapping_layers):
if mapping_layer != 'up_1':
fs, ft = opt.student_ngf * 16, opt.teacher_ngf * 16
else:
fs, ft = opt.student_ngf * 4, opt.teacher_ngf * 4
if hasattr(opt, 'distiller'):
netA = nn.Conv2d(in_channels=fs,
out_channels=ft,
kernel_size=1)
else:
netA = SuperConv2d(in_channels=fs,
out_channels=ft,
kernel_size=1)
networks.init_net(netA, opt.init_type, opt.init_gain, self.gpu_ids)
self.netAs.append(netA)
self.criterionGAN = GANLoss(opt.gan_mode)
self.criterionFeat = nn.L1Loss()
self.criterionVGG = VGGLoss()
self.optimizers = []
self.netG_teacher.eval()
self.config = None
def initialize_networks(self, opt):
netG = networks.define_G(opt)
if opt.isTrain:
netD = networks.define_D(opt)
netD_uncond = networks.define_D(opt, True)
else:
netD = None
netD_uncond = None
netE = networks.define_E(opt) if opt.use_vae else None
if not opt.isTrain or opt.continue_train:
netG = util.load_network(netG, 'G', opt.which_epoch, opt)
if opt.isTrain:
netD = util.load_network(netD, 'D', opt.which_epoch, opt)
netD_uncond = util.load_network(netD_uncond, 'D_uncond',
opt.which_epoch, opt)
if opt.use_vae:
netE = util.load_network(netE, 'E', opt.which_epoch, opt)
return netG, netD, netD_uncond, netE
def initialize_networks(self, opt, load_weights=True):
netG = networks.define_G(opt)
netD = networks.define_D(opt) if opt.isTrain else None
netE = networks.define_E(opt) if opt.use_vae else None
if (not opt.isTrain or opt.continue_train) and load_weights:
netG = util.load_network(netG, 'G', opt.which_epoch, opt)
if opt.isTrain:
netD = util.load_network(netD, 'D', opt.which_epoch, opt)
if opt.use_vae:
netE = util.load_network(netE, 'E', opt.which_epoch, opt)
return netG, netD, netE
def initialize_networks(self, opt):
netG = networks.define_G(opt)
netD = networks.define_D(opt) if opt.isTrain else None
netE = networks.define_E(opt) if opt.use_vae else None
if not opt.isTrain or opt.continue_train:
netG = util.load_network(netG, "G", opt.which_epoch, opt)
if opt.isTrain:
netD = util.load_network(netD, "D", opt.which_epoch, opt)
if opt.use_vae:
netE = util.load_network(netE, "E", opt.which_epoch, opt)
return netG, netD, netE
def initialize_networks(self, opt):
netG2 = networks.define_G(opt)
netD2 = networks.define_D(opt) if opt.isTrain else None
netE2 = networks.define_E(opt) if opt.use_vae else None
if not opt.isTrain or opt.continue_train:
netG2 = util.load_network2(netG2, 'G', opt.which_epoch, opt)
if opt.isTrain:
netD2 = util.load_network2(netD2, 'D', opt.which_epoch, opt)
if opt.use_vae:
netE2 = util.load_network2(netE2, 'E', opt.which_epoch, opt)
elif opt.use_vae and opt.pretrain_vae:
netE2 = util.load_network2(netE2, 'E', opt.which_epoch, opt)
if opt.edge_cat:
opt.label_nc -= 1
opt.semantic_nc -= 1
netG = networks.define_G(opt)
netD = networks.define_D(opt) if opt.isTrain else None
netE = networks.define_E(opt) if opt.use_vae else None
if not opt.isTrain or opt.continue_train:
netG = util.load_network(netG, 'G', opt.which_epoch, opt)
if opt.isTrain:
netD = util.load_network(netD, 'D', opt.which_epoch, opt)
if opt.use_vae:
netE = util.load_network(netE, 'E', opt.which_epoch, opt)
elif opt.use_vae and opt.pretrain_vae:
netE = util.load_network(netE, 'E', opt.which_epoch, opt)
print('Load fixed netE.')
if opt.edge_cat:
opt.label_nc += 1
opt.semantic_nc += 1
return netG, netD, netE, netG2, netD2, netE2
