def _init_net(self, net, net_name, ckpt):
net.to(self.device)
net_optimizer = self.define_optimizer(net)
if ckpt is not None:
net.load_state_dict(ckpt[net_name]['weight'])
net_optimizer.load_state_dict(ckpt[net_name]['optimizer'])
lr_scheduler = util.get_scheduler(net_optimizer, self.opt,
ckpt['epoch'])
else:
net.apply(network.weights_init(self.opt.init_type))
lr_scheduler = util.get_scheduler(net_optimizer, self.opt, -1)
net.train()
self.nets[net_name] = net
self.optimizers[net_name] = net_optimizer
self.lr_schedulers[net_name] = lr_scheduler
def initialize(self, opt):
BaseModel.initialize(self, opt)
self.config = get_config(opt.config)
nb = opt.batchSize
size = opt.fineSize
self.input_A = self.Tensor(nb, opt.input_nc, size, size)
self.input_B = self.Tensor(nb, opt.output_nc, size, size)
# load/define networks
# The naming conversion is different from those used in the paper
# Code (paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
self.gen_a = VAEGen(self.config['input_dim_a'], self.config['gen'])
self.gen_b = VAEGen(self.config['input_dim_a'], self.config['gen'])
if self.isTrain:
self.dis_a = MsImageDis(self.config['input_dim_a'], self.config['dis']) # discriminator for domain a
self.dis_b = MsImageDis(self.config['input_dim_b'], self.config['dis']) # discriminator for domain b
if not self.isTrain or opt.continue_train:
which_epoch = opt.which_epoch
self.load_network(self.gen_a, 'G_A', which_epoch)
self.load_network(self.gen_b, 'G_B', which_epoch)
if self.isTrain:
self.load_network(self.dis_a, 'D_A', which_epoch)
self.load_network(self.dis_b, 'D_B', which_epoch)
if self.isTrain:
self.old_lr = self.config['lr']
self.fake_A_pool = ImagePool(opt.pool_size)
self.fake_B_pool = ImagePool(opt.pool_size)
# define loss functions
# Setup the optimizers
beta1 = self.config['beta1']
beta2 = self.config['beta2']
dis_params = list(self.dis_a.parameters()) + list(self.dis_b.parameters())
gen_params = list(self.gen_a.parameters()) + list(self.gen_b.parameters())
self.dis_opt = torch.optim.Adam([p for p in dis_params if p.requires_grad],
lr=self.config['lr'], betas=(beta1, beta2), weight_decay=self.config['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=self.config['lr'], betas=(beta1, beta2), weight_decay=self.config['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, self.config)
self.gen_scheduler = get_scheduler(self.gen_opt, self.config)
# Network weight initialization
# self.apply(weights_init(self.config['init']))
self.dis_a.apply(weights_init('gaussian'))
self.dis_b.apply(weights_init('gaussian'))
# Load VGG model if needed
if 'vgg_w' in self.config.keys() and self.config['vgg_w'] > 0:
self.vgg = load_vgg16(self.config['vgg_model_path'] + '/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
self.gen_a.cuda()
self.gen_b.cuda()
self.dis_a.cuda()
self.dis_b.cuda()
print('---------- Networks initialized -------------')
networks.print_network(self.gen_a)
networks.print_network(self.gen_b)
if self.isTrain:
networks.print_network(self.dis_a)
networks.print_network(self.dis_b)
print('-----------------------------------------------')