def resume_prev(self, checkpoint_dir, hyperparameters):
# Load generators
last_model_name = get_model_list(checkpoint_dir, "gen")
print('resume from generator checkpoint named:', last_model_name)
state_dict = torch.load(last_model_name)
self.gen_a.load_state_dict(state_dict['a'])
self.gen_b.load_state_dict(state_dict['b'])
try:
iterations = int(float(last_model_name[4:11]))
except:
iterations = int(float(last_model_name.split('/')[-1][4:11]))
# Load discriminators
last_model_name = get_model_list(checkpoint_dir, "dis")
print('resume from discriminator checkpoint named:', last_model_name)
state_dict = torch.load(last_model_name)
self.dis_a.load_state_dict(state_dict['a'])
self.dis_b.load_state_dict(state_dict['b'])
# Load optimizers
state_dict = torch.load(os.path.join(checkpoint_dir, 'optimizer.pt'))
self.dis_opt.load_state_dict(state_dict['dis'])
self.gen_opt.load_state_dict(state_dict['gen'])
# Reinitilize schedulers
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters,
iterations)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters,
iterations)
print('Resume from iteration %d' % iterations)
return iterations
def __init__(self, hps, use_global=False):
super(V2VModel, self).__init__()
# Model config
self.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
self.style_dim = hps.gen_style_dim
# Initialization
self.gen_a = Generator(hps.input_dim_a, hps, use_global) # auto-encoder for domain a
self.gen_b = Generator(hps.input_dim_b, hps, use_global) # auto-encoder for domain b
self.dis_a = Discriminator(hps.input_dim_a, hps) # discriminator for domain a
self.dis_b = Discriminator(hps.input_dim_b, hps) # discriminator for domain b
# Setup the optimizers
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_optimizer = optim.Adam([p for p in dis_params if p.requires_grad],
lr=hps.lr, betas=(hps.beta1, hps.beta2), weight_decay=hps.weight_decay)
self.gen_optimizer = optim.Adam([p for p in gen_params if p.requires_grad],
lr=hps.lr, betas=(hps.beta1, hps.beta2), weight_decay=hps.weight_decay)
self.dis_scheduler = get_scheduler(self.dis_optimizer, hps)
self.gen_scheduler = get_scheduler(self.gen_optimizer, hps)
if hps.g_comp > 0:
self.temp_loss = TemporalLoss()
# Network weight initialization
self.apply(weights_init(hps.init))
self.dis_a.apply(weights_init('gaussian'))
self.dis_b.apply(weights_init('gaussian'))
def resume(self, checkpoint_dir, hyperparameters):
# Load generators
last_model_name = get_model_list(checkpoint_dir, "gen")
state_dict = torch.load(last_model_name)
self.gen_a.load_state_dict(state_dict['a'])
self.gen_b = self.gen_a
iterations = int(last_model_name[-11:-3])
# Load discriminators
last_model_name = get_model_list(checkpoint_dir, "dis")
state_dict = torch.load(last_model_name)
self.dis_a.load_state_dict(state_dict['a'])
self.dis_b = self.dis_a
# Load ID dis
last_model_name = get_model_list(checkpoint_dir, "id")
state_dict = torch.load(last_model_name)
self.id_a.load_state_dict(state_dict['a'])
self.id_b = self.id_a
# Load optimizers
try:
state_dict = torch.load(os.path.join(checkpoint_dir, 'optimizer.pt'))
self.dis_opt.load_state_dict(state_dict['dis'])
self.gen_opt.load_state_dict(state_dict['gen'])
self.id_opt.load_state_dict(state_dict['id'])
except:
pass
# Reinitilize schedulers
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters, iterations)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters, iterations)
print('Resume from iteration %d' % iterations)
return iterations
def __init__(self, param):
super(Trainer, self).__init__()
lr_d = param['lr_d']
# Initiate the networks
self.generator = get_generator(param)
self.discriminator_bg = get_discriminator(param)
self.discriminator_rf = get_discriminator(param)
# ############################################################################
# from thop import profile
# from thop import clever_format
# input_i = torch.randn(1, 3, 224, 224)
# macs, params = profile(self.discriminator_bg, inputs=(input_i, ))
# print('========================')
# print('MACs: ', macs)
# print('PARAMs: ', params)
# print('------------------------')
# macs, params = clever_format([macs, params], "%.3f")
# print('Clever MACs: ', macs)
# print('Clever PARAMs: ', params)
# print('========================')
# ############################################################################
# Setup the optimizers
beta1 = param['beta1']
beta2 = param['beta2']
dis_params = list(self.discriminator_bg.parameters()) + list(
self.discriminator_rf.parameters())
self.dis_opt = torch.optim.Adam(dis_params,
lr=lr_d,
betas=(beta1, beta2),
weight_decay=param['weight_decay'])
self.gen_opt = torch.optim.SGD(params=[{
'params':
self.get_params(self.generator, key='1x'),
'lr':
param.lr_g
}, {
'params':
self.get_params(self.generator, key='10x'),
'lr':
10 * param.lr_g
}],
momentum=param.momentum)
self.dis_scheduler = get_scheduler(self.dis_opt, param)
self.gen_scheduler = get_scheduler(self.gen_opt, param)
# self.dis_scheduler = None
# self.gen_scheduler = None
# Network weight initialization
# self.apply(weights_init(param['init']))
self.discriminator_bg.apply(weights_init('gaussian'))
self.discriminator_rf.apply(weights_init('gaussian'))
self.best_result = float('inf')
self.perceptual_criterion = PerceptualLoss()
self.retina_criterion = RetinaLoss()
self.semantic_criterion = nn.CrossEntropyLoss(ignore_index=255)
self.best_result = 0
def __init__(self, hyperparameters):
super(MUNIT_Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks
self.gen_a = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain a
self.gen_b = AdaINGen(hyperparameters['input_dim_b'], hyperparameters['gen']) # auto-encoder for domain b
self.dis_a = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator for domain a
self.dis_b = MsImageDis(hyperparameters['input_dim_b'], hyperparameters['dis']) # discriminator for domain b
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.style_dim = hyperparameters['gen']['style_dim']
self.criterionGAN = GANLoss(hyperparameters['dis']['gan_type']).cuda()
self.featureLoss = nn.MSELoss(reduction='mean')
# fix the noise used in sampling
display_size = int(hyperparameters['display_size'])
self.s_a = torch.randn(display_size, self.style_dim, 1, 1).cuda()
self.s_b = torch.randn(display_size, self.style_dim, 1, 1).cuda()
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['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=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
self.apply(weights_init(hyperparameters['init']))
self.dis_a.apply(weights_init('gaussian'))
self.dis_b.apply(weights_init('gaussian'))
def __init__(self, hyperparameters):
super(UNIT_Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks
self.gen_a = VAEGen(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain a
self.gen_b = VAEGen(hyperparameters['input_dim_b'], hyperparameters['gen']) # auto-encoder for domain b
self.dis_a = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator for domain a
self.dis_b = MsImageDis(hyperparameters['input_dim_b'], hyperparameters['dis']) # discriminator for domain b
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['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=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
self.apply(weights_init(hyperparameters['init']))
self.dis_a.apply(weights_init('gaussian'))
self.dis_b.apply(weights_init('gaussian'))
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
def __init__(self, hyperparameters):
super(MUNIT_Trainer, self).__init__()
lr = hyperparameters['lr']
self.gen_a = AdaInGenerator(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain a
self.gen_b = AdaInGenerator(hyperparameters['input_dim_b'], hyperparameters['gen']) # auto-encoder for domain b
self.dis_a = ImageDiscriminator(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator for domain a
self.dis_b = ImageDiscriminator(hyperparameters['input_dim_b'], hyperparameters['dis']) # discriminator for domain b
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.style_dim = hyperparameters['gen']['style_dim']
display_size = int(hyperparameters['display_size'])
self.s_a = torch.randn(display_size, self.style_dim, 1, 1).cuda()
self.s_b = torch.randn(display_size, self.style_dim, 1, 1).cuda()
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['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=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
self.apply(weights_init(hyperparameters['init']))
self.dis_a.apply(weights_init('gaussian'))
self.dis_b.apply(weights_init('gaussian'))
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
def __init__(self, config):
#super(Trainer, self).__init__()
self.config = config
lr = config['lr']
# Initiate the networks
imgconf = config['image']
self.vae = VAE(imgconf,config['gen'],config['latent'])
self.dis = DiscriminatorVAE(config['dis'],imgconf['image_size'],imgconf['image_dim'])
'''
disconf = config['dis']
self.dis = DiscriminatorVAE(disconf['n_downsample'],disconf['n_res'],
imgconf['image_size'],imgconf['image_dim'],
disconf['dim'],disconf['norm'],disconf['activ'],disconf['pad_type'])
'''
self.vae_optim = optim.Adam(self.vae.parameters(),lr=lr)
self.dis_optim = optim.Adam(self.dis.parameters(),lr=lr)
self.vae_scheduler = get_scheduler(self.vae_optim, config)
self.dis_scheduler = get_scheduler(self.dis_optim, config)
'''
beta1 = config['beta1']
beta2 = config['beta2']
self.vae_optim = optim.Adam(self.vae.parameters(),
lr=lr, betas=(beta1, beta2), weight_decay=config['weight_decay'])
self.dis_optim = optim.Adam(self.dis.parameters(),
lr=lr, betas=(beta1, beta2), weight_decay=config['weight_decay'])
'''
self.mse_crit = nn.MSELoss()
self.bce_vae = nn.BCELoss()
self.bce_dis = nn.BCELoss()
'''
def resume(self, checkpoint_dir, param):
# Load generators
last_model_name = get_model_list(checkpoint_dir, "gen")
state_dict = torch.load(last_model_name)
self.gen_i.load_state_dict(state_dict['i'])
self.gen_r.load_state_dict(state_dict['r'])
self.gen_s.load_state_dict(state_dict['s'])
if self.with_mapping:
self.fea_m.load_state_dict(state_dict['fm'])
self.fea_s.load_state_dict(state_dict['fs'])
self.best_result = state_dict['best_result']
iterations = int(last_model_name[-11:-3])
# Load discriminators
last_model_name = get_model_list(checkpoint_dir, "dis")
state_dict = torch.load(last_model_name)
self.dis_r.load_state_dict(state_dict['r'])
self.dis_s.load_state_dict(state_dict['s'])
# Load optimizers
state_dict = torch.load(os.path.join(checkpoint_dir, 'optimizer.pt'))
self.dis_opt.load_state_dict(state_dict['dis'])
self.gen_opt.load_state_dict(state_dict['gen'])
# Reinitilize schedulers
self.dis_scheduler = get_scheduler(self.dis_opt, param, iterations)
self.gen_scheduler = get_scheduler(self.gen_opt, param, iterations)
print('Resume from iteration %d' % iterations)
return iterations
def resume(self, checkpoint_dir, hyperparameters):
# Load generators
last_model_name = get_model_list(checkpoint_dir, "gen")
state_dict = torch.load(last_model_name)
self.gen_a.load_state_dict(state_dict['a'])
self.gen_b.load_state_dict(state_dict['b'])
iterations = int(last_model_name[-11:-3])
# Load discriminators
last_model_name = get_model_list(checkpoint_dir, "dis")
state_dict = torch.load(last_model_name)
self.dis_a.load_state_dict(state_dict['a'])
self.dis_b.load_state_dict(state_dict['b'])
# Load content classifier
last_model_name = get_model_list(checkpoint_dir, "con_cla")
state_dict = torch.load(last_model_name)
self.content_classifier.load_state_dict(state_dict['con_cla'])
# Load optimizers
state_dict = torch.load(os.path.join(checkpoint_dir, 'optimizer.pt'))
self.dis_opt.load_state_dict(state_dict['dis'])
self.gen_opt.load_state_dict(state_dict['gen'])
self.cla_opt.load_state_dict(state_dict['con_cla'])
# Reinitilize schedulers
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters,
iterations)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters,
iterations)
self.cla_scheduler = get_scheduler(self.cla_opt, hyperparameters,
iterations)
print('Resume from iteration %d' % iterations)
return iterations
def resume(self, checkpoint_dir, configs):
# Load generators
last_model_name = get_model_list(checkpoint_dir, "gen")
state_dict = torch.load(last_model_name,
map_location=lambda storage, loc: storage)
self.gen.load_state_dict(state_dict['a'])
iterations = int(
last_model_name[-15:-7]) if 'avg' in last_model_name else int(
last_model_name[-11:-3])
# Load discriminators
last_model_name = get_model_list(checkpoint_dir, "dis")
state_dict = torch.load(last_model_name,
map_location=lambda storage, loc: storage)
self.dis.load_state_dict(state_dict['b'])
# Load optimizers
#state_dict = torch.load(os.path.join(checkpoint_dir, 'optimizer.pt'), map_location=lambda storage, loc: storage)
#self.dis_opt.load_state_dict(state_dict['dis'])
#self.gen_opt.load_state_dict(state_dict['gen'])
# Reinitilize schedulers
self.dis_scheduler = get_scheduler(self.dis_opt, configs, iterations)
self.gen_scheduler = get_scheduler(self.gen_opt, configs, iterations)
if torch.__version__ != '0.4.1':
for _ in range(iterations):
self.gen_scheduler.step()
self.dis_scheduler.step()
print('Resume from iteration %d' % iterations)
return iterations
def __init__(self, hyperparameters):
super(UNIT_Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks
self.gen_a = VAEGen(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain a
self.gen_b = VAEGen(hyperparameters['input_dim_b'], hyperparameters['gen']) # auto-encoder for domain b
self.dis_a = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator for domain a
self.dis_b = MsImageDis(hyperparameters['input_dim_b'], hyperparameters['dis']) # discriminator for domain b
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['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=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
self.apply(weights_init(hyperparameters['init']))
self.dis_a.apply(weights_init('gaussian'))
self.dis_b.apply(weights_init('gaussian'))
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
def __init__(self, hyperparameters, resume_epoch=-1, snapshot_dir=None):
super(UNIT_Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks.
self.gen = VAEGen(
hyperparameters['input_dim'] + hyperparameters['n_datasets'],
hyperparameters['gen'],
hyperparameters['n_datasets']) # Auto-encoder for domain a.
self.dis = MsImageDis(
hyperparameters['input_dim'] + hyperparameters['n_datasets'],
hyperparameters['dis']) # Discriminator for domain a.
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.sup = UNet(input_channels=hyperparameters['input_dim'],
num_classes=2).cuda()
# Setup the optimizers.
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['beta2']
dis_params = list(self.dis.parameters())
gen_params = list(self.gen.parameters()) + list(self.sup.parameters())
self.dis_opt = torch.optim.Adam(
[p for p in dis_params if p.requires_grad],
lr=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam(
[p for p in gen_params if p.requires_grad],
lr=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization.
self.apply(weights_init(hyperparameters['init']))
self.dis.apply(weights_init('gaussian'))
# Presetting one hot encoding vectors.
self.one_hot_img = torch.zeros(hyperparameters['n_datasets'],
hyperparameters['batch_size'],
hyperparameters['n_datasets'], 256,
256).cuda()
self.one_hot_h = torch.zeros(hyperparameters['n_datasets'],
hyperparameters['batch_size'],
hyperparameters['n_datasets'], 64,
64).cuda()
for i in range(hyperparameters['n_datasets']):
self.one_hot_img[i, :, i, :, :].fill_(1)
self.one_hot_h[i, :, i, :, :].fill_(1)
if resume_epoch != -1:
self.resume(snapshot_dir, hyperparameters)
def __init__(self, hyperparameters):
super(MUNIT_Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks
self.is_ganilla_gen = hyperparameters['gen']['ganilla_gen']
if self.is_ganilla_gen == False:
self.gen_a = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain a
self.gen_b = AdaINGen(hyperparameters['input_dim_b'], hyperparameters['gen']) # auto-encoder for domain b
else:
self.gen_a = AdaINGanilla(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain a with ganilla architecture
self.gen_b = AdaINGanilla(hyperparameters['input_dim_b'], hyperparameters['gen']) # auto-encoder for domain b with ganilla architecture
print(self.gen_a)
if hyperparameters['dis']['dis_type'] == 'patch':
if hyperparameters['dis']['use_patch_gan']:
self.dis_a = PatchDis(hyperparameters['input_dim_a'], hyperparameters['dis'])
self.dis_b = PatchDis(hyperparameters['input_dim_b'], hyperparameters['dis'])
else:
self.dis_a = MsImageDis(hyperparameters['input_dim_a'],
hyperparameters['dis']) # discriminator for domain a
self.dis_b = MsImageDis(hyperparameters['input_dim_b'],
hyperparameters['dis']) # discriminator for domain b
print(self.dis_a)
else:
self.dis_a = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator for domain a
self.dis_b = MsImageDis(hyperparameters['input_dim_b'], hyperparameters['dis']) # discriminator for domain b
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.style_dim = hyperparameters['gen']['style_dim']
# fix the noise used in sampling
display_size = int(hyperparameters['display_size'])
self.s_a = torch.randn(display_size, self.style_dim, 1, 1).cuda()
self.s_b = torch.randn(display_size, self.style_dim, 1, 1).cuda()
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['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=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
self.apply(weights_init(hyperparameters['init']))
self.dis_a.apply(weights_init('gaussian'))
self.dis_b.apply(weights_init('gaussian'))
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')
self.vgg.eval()
self.VggExtract = VggExtract(self.vgg)
for param in self.vgg.parameters():
param.requires_grad = False
def __init__(self, hyperparameters):
super(IPMNet_Trainer, self).__init__()
lr = hyperparameters['lr']
vgg_weight_file = hyperparameters['vgg_weight_file']
# Initiate the networks
self.gen_a = AdaINGen(
hyperparameters['input_dim_a'],
hyperparameters['gen']) # auto-encoder for domain a
self.gen_b = self.gen_a # AdaINGen(hyperparameters['input_dim_b'], hyperparameters['gen']) # auto-encoder for domain b
self.dis_a = MsImageDis(
hyperparameters['input_dim_a'],
hyperparameters['dis']) # discriminator for domain a
self.dis_b = MsImageDis(
hyperparameters['input_dim_b'],
hyperparameters['dis']) # discriminator for domain b
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.style_dim = hyperparameters['gen']['style_dim']
# fix the noise used in sampling
display_size = int(hyperparameters['display_size'])
self.s_a = torch.randn(display_size, self.style_dim, 1, 1).cuda()
self.s_b = torch.randn(display_size, self.style_dim, 1, 1).cuda()
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['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=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam(
[p for p in gen_params if p.requires_grad],
lr=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
self.apply(weights_init(hyperparameters['init']))
self.dis_a.apply(weights_init(hyperparameters['init']))
self.dis_b.apply(weights_init(hyperparameters['init']))
# Load VGGFace model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_resnet50(vgg_weight_file)
self.vgg.eval()
self.vgg.fc.reset_parameters()
for param in self.vgg.parameters():
param.requires_grad = False
def update_learning_rate(self):
if self.lr_policy == 'cosa':
if self.dis_opt.param_groups[0]['lr'] == self.configs['eta_min'] or \
self.gen_opt.param_groups[0]['lr'] == self.configs['eta_min']:
self.configs['step_size'] *= self.configs['t_mult']
self.dis_scheduler = get_scheduler(self.dis_opt, self.configs)
self.gen_scheduler = get_scheduler(self.gen_opt, self.configs)
if self.dis_scheduler is not None:
self.dis_scheduler.step()
if self.gen_scheduler is not None:
self.gen_scheduler.step()
def __init__(self, hyperparameters):
super(UNIT_Trainer, self).__init__()
lr = hyperparameters["lr"]
# Initiate the networks
self.gen_a = VAEGen(
hyperparameters["input_dim_a"],
hyperparameters["gen"]) # auto-encoder for domain a
self.gen_b = VAEGen(
hyperparameters["input_dim_b"],
hyperparameters["gen"]) # auto-encoder for domain b
self.dis_a = MsImageDis(
hyperparameters["input_dim_a"],
hyperparameters["dis"]) # discriminator for domain a
self.dis_b = MsImageDis(
hyperparameters["input_dim_b"],
hyperparameters["dis"]) # discriminator for domain b
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
# Setup the optimizers
beta1 = hyperparameters["beta1"]
beta2 = hyperparameters["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=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters["weight_decay"],
)
self.gen_opt = torch.optim.Adam(
[p for p in gen_params if p.requires_grad],
lr=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters["weight_decay"],
)
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
self.apply(weights_init(hyperparameters["init"]))
self.dis_a.apply(weights_init("gaussian"))
self.dis_b.apply(weights_init("gaussian"))
# Load VGG model if needed
if "vgg_w" in hyperparameters.keys() and hyperparameters["vgg_w"] > 0:
self.vgg = load_vgg16(hyperparameters["vgg_model_path"] +
"/models")
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
def __init__(self, hyperparameters):
super(Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks
# auto-encoder for domain a
self.trait_dim = hyperparameters['gen']['trait_dim']
self.gen_a = VAEGen(hyperparameters['input_dim'],
hyperparameters['basis_encoder_dims'],
hyperparameters['trait_encoder_dims'],
hyperparameters['decoder_dims'], self.trait_dim)
# auto-encoder for domain b
self.gen_b = VAEGen(hyperparameters['input_dim'],
hyperparameters['basis_encoder_dims'],
hyperparameters['trait_encoder_dims'],
hyperparameters['decoder_dims'], self.trait_dim)
# discriminator for domain a
self.dis_a = Discriminator(hyperparameters['input_dim'],
hyperparameters['dis_dims'], 1)
# discriminator for domain b
self.dis_b = Discriminator(hyperparameters['input_dim'],
hyperparameters['dis_dims'], 1)
# fix the noise used in sampling
self.trait_a = torch.randn(8, self.trait_dim, 1, 1)
self.trait_b = torch.randn(8, self.trait_dim, 1, 1)
# Setup the optimizers
dis_params = list(self.dis_a.parameters()) + \
list(self.dis_b.parameters())
gen_params = list(self.gen_a.parameters()) + \
list(self.gen_b.parameters())
for _p in gen_params:
print(_p.data.shape)
self.dis_opt = torch.optim.Adam(
[p for p in dis_params if p.requires_grad],
lr=lr,
weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam(
[p for p in gen_params if p.requires_grad],
lr=lr,
weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
self.apply(weights_init(hyperparameters['init']))
self.gen_a.apply(weights_init('gaussian'))
self.gen_b.apply(weights_init('gaussian'))
self.dis_a.apply(weights_init('gaussian'))
self.dis_b.apply(weights_init('gaussian'))
def resume(self, checkpoint_dir, hyperparameters):
"""
Resume the training loading the network parameters
Arguments:
checkpoint_dir {string} -- path to the directory where the checkpoints are saved
hyperparameters {dictionnary} -- dictionnary with all hyperparameters
Returns:
int -- number of iterations (used by the optimizer)
"""
# Load generators
last_model_name = get_model_list(checkpoint_dir, "gen")
state_dict = torch.load(last_model_name)
if self.gen_state == 0:
self.gen_a.load_state_dict(state_dict["a"])
self.gen_b.load_state_dict(state_dict["b"])
elif self.gen_state == 1:
self.gen.load_state_dict(state_dict["2"])
else:
print("self.gen_state unknown value:", self.gen_state)
# Load domain classifier
if self.domain_classif == 1:
last_model_name = get_model_list(checkpoint_dir, "domain_classif")
state_dict = torch.load(last_model_name)
self.domain_classifier.load_state_dict(state_dict["d"])
iterations = int(last_model_name[-11:-3])
# Load discriminators
last_model_name = get_model_list(checkpoint_dir, "dis")
state_dict = torch.load(last_model_name)
self.dis_a.load_state_dict(state_dict["a"])
self.dis_b.load_state_dict(state_dict["b"])
# Load optimizers
state_dict = torch.load(os.path.join(checkpoint_dir, "optimizer.pt"))
self.dis_opt.load_state_dict(state_dict["dis"])
self.gen_opt.load_state_dict(state_dict["gen"])
if self.domain_classif == 1:
self.dann_opt.load_state_dict(state_dict["dann"])
self.dann_scheduler = get_scheduler(self.dann_opt, hyperparameters,
iterations)
# Reinitilize schedulers
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters,
iterations)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters,
iterations)
print("Resume from iteration %d" % iterations)
return iterations
def __init__(self, config):
self.config = config
lr = config['lr']
# Initiate the networks
imgconf = config['image']
self.vae = CatVAE(imgconf,config['gen'],config['latent'])
self.dis = CatQNet(config['dis'],imgconf['image_size'],imgconf['image_dim'],config['latent'])
self.vae_optim = optim.Adam(self.vae.parameters(),lr=lr)
self.dis_optim = optim.Adam(self.dis.parameters(),lr=lr)
self.vae_scheduler = get_scheduler(self.vae_optim, config)
self.dis_scheduler = get_scheduler(self.dis_optim, config)
self.mse_crit = nn.MSELoss()
def __init__(self, param):
super(SemanticTrainer, self).__init__()
lr_d = param['lr_d']
# Initiate the networks
self.generator = get_generator(param)
# Setup the optimizers
beta1 = param['beta1']
beta2 = param['beta2']
self.gen_opt = torch.optim.SGD(params=[{
'params':
self.get_params(self.generator, key='1x'),
'lr':
param.lr_g
}, {
'params':
self.get_params(self.generator, key='10x'),
'lr':
10 * param.lr_g
}],
momentum=param.momentum)
self.gen_scheduler = get_scheduler(self.gen_opt, param)
# self.dis_scheduler = None
# self.gen_scheduler = None
# Network weight initialization
# self.apply(weights_init(param['init']))
self.best_result = 0
self.semantic_criterion = nn.CrossEntropyLoss(ignore_index=255)
def __init__(self, param):
super(Trainer, self).__init__()
lr = param['lr']
# Initiate the networks
self.student = get_student(param)
self.teacher = get_teacher(param)
self.load_network(self.teacher, param['pretrain_model_T'])
# self.discriminator = get_discriminator(param) # todo: add D to further boost the performance
# Setup the optimizers
beta1 = param['beta1']
beta2 = param['beta2']
# teacher_params = list(self.discriminator.parameters())
student_params = list(self.student.parameters())
self.stu_opt = torch.optim.Adam(student_params,
lr=lr,
betas=(beta1, beta2),
weight_decay=param['weight_decay'])
self.stu_scheduler = get_scheduler(self.stu_opt, param)
# self.dis_scheduler = None
# self.gen_scheduler = None
# Network weight initialization
# self.apply(weights_init(param['init']))
self.best_result = float('inf')
self.criterion_pair = CriterionPairWiseForWholeFeatAfterPool(
param.pairwise_scale)
self.criterion_pixel = nn.L1Loss()
self.criterion = nn.L1Loss()
def __init__(self, opt):
self.opt = opt
self.is_train = opt.is_train
self.gpu_ids = opt.gpu_ids
self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
self.save_dir = os.path.join(opt.checkpoints_dir, opt.dataset, opt.class_name, opt.timestamp)
self.pretrain_dir = os.path.join(opt.checkpoints_dir, opt.dataset, opt.class_name, opt.pretrain)
self.optimizer = None
self.loss_func = None
self.loss = None
self.confusion = None # confusion matrix
self.multi_confusion = None
self.net_name = opt.net_name
self.net = net.init_net(opt)
self.net.train(self.is_train)
self.loss_func = torch.nn.NLLLoss().to(self.device)
if self.is_train:
self.optimizer = torch.optim.Adam(self.net.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.scheduler = utils.get_scheduler(self.optimizer, opt)
if not self.is_train: #or opt.continue_train:
self.load_network(opt.which_epoch, mode='test')
if self.is_train and opt.pretrain != '-':
self.load_network(opt.which_epoch, mode='pretrain')
def setup(self):
"""Load and print networks; create schedulers.
"""
if self.configuration['load_checkpoint'] >= 0:
last_checkpoint = self.configuration['load_checkpoint']
else:
last_checkpoint = -1
if last_checkpoint >= 0:
# enable restarting training
self.load_networks(last_checkpoint)
if self.is_train:
self.load_optimizers(last_checkpoint)
for o in self.optimizers:
o.param_groups[0]['lr'] = o.param_groups[0][
'initial_lr'] # reset learning rate
self.schedulers = [
get_scheduler(optimizer, self.configuration)
for optimizer in self.optimizers
]
if last_checkpoint > 0:
for s in self.schedulers:
for _ in range(last_checkpoint):
s.step()
self.print_networks()
def __init__(self, hyperparameters):
super(UNIT_Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks
self.gen_a = VAEGen(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain a
self.gen_b = VAEGen(hyperparameters['input_dim_b'], hyperparameters['gen']) # auto-encoder for domain b
if not hyperparameters['origin']:
self.dis_a = MultiscaleDiscriminator(hyperparameters['input_dim_a'], # discriminator for a
ndf=64, n_layers=3, norm_layer=nn.InstanceNorm2d, use_sigmoid=False,
num_D=2, getIntermFeat=True
)
self.dis_b = MultiscaleDiscriminator(hyperparameters['input_dim_b'], # discriminator for b
ndf=64, n_layers=3, norm_layer=nn.InstanceNorm2d, use_sigmoid=False,
num_D=2, getIntermFeat=True
)
self.criterionGAN = GANLoss(use_lsgan=True, tensor=torch.cuda.FloatTensor)
else:
self.dis_a = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis'])
self.dis_b = MsImageDis(hyperparameters['input_dim_b'], hyperparameters['dis'])
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['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=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
self.apply(weights_init(hyperparameters['init']))
self.dis_a.apply(weights_init('gaussian'))
self.dis_b.apply(weights_init('gaussian'))
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
def __init__(self, hyperparameters):
super(MUNIT_Trainer, self).__init__() # super() 函数是用于调用父类(超类)的一个方法。
lr = hyperparameters['lr']
# Initiate the networks， 需要好好看看生成器和鉴别器到底是如何构造的
self.gen_a = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain a
self.gen_b = AdaINGen(hyperparameters['input_dim_b'], hyperparameters['gen']) # auto-encoder for domain b
self.dis_a = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator for domain a
self.dis_b = MsImageDis(hyperparameters['input_dim_b'], hyperparameters['dis']) # discriminator for domain b
# https://blog.csdn.net/liuxiao214/article/details/81037416
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.style_dim = hyperparameters['gen']['style_dim']
# fix the noise used in sampling
display_size = int(hyperparameters['display_size'])
# s_a , s_b 表示的是两个不同的style
self.s_a = torch.randn(display_size, self.style_dim, 1, 1).cuda() # 16*8*1*1
self.s_b = torch.randn(display_size, self.style_dim, 1, 1).cuda()
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['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=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
# 优化策略
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
# 解释 apply apply(lambda x,y : x+y, (1),{'y' : 2}) https://zhuanlan.zhihu.com/p/42756654
self.apply(weights_init(hyperparameters['init'])) # 初始化当前类
self.dis_a.apply(weights_init('gaussian')) # 初始化dis_a,是一个类对象
self.dis_b.apply(weights_init('gaussian'))
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
def __init__(self, hyperparameters):
super(aclgan_Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks
self.gen_AB = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain A
self.gen_BA = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen']) # auto-encoder for domain B
self.dis_A = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator for domain A
self.dis_B = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator for domain B
self.dis_2 = MsImageDis(hyperparameters['input_dim_b'], hyperparameters['dis']) # discriminator 2
# self.dis_2B = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator 2 for domain B
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.style_dim = hyperparameters['gen']['style_dim']
# fix the noise used in sampling
display_size = int(hyperparameters['display_size'])
self.z_1 = torch.randn(display_size, self.style_dim, 1, 1).cuda()
self.z_2 = torch.randn(display_size, self.style_dim, 1, 1).cuda()
self.z_3 = torch.randn(display_size, self.style_dim, 1, 1).cuda()
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['beta2']
dis_params = list(self.dis_A.parameters()) + list(self.dis_B.parameters()) + list(self.dis_2.parameters())
gen_params = list(self.gen_AB.parameters()) + list(self.gen_BA.parameters())
self.dis_opt = torch.optim.Adam([p for p in dis_params if p.requires_grad],
lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
self.alpha = hyperparameters['alpha']
self.focus_lam = hyperparameters['focus_loss']
# Network weight initialization
self.apply(weights_init(hyperparameters['init']))
self.dis_A.apply(weights_init('gaussian'))
self.dis_B.apply(weights_init('gaussian'))
self.dis_2.apply(weights_init('gaussian'))
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
def resume(self, checkpoint_dir, hyperparameters):
print("--> " + checkpoint_dir)
# Load generator.
last_model_name = get_model_list(checkpoint_dir, "gen")
state_dict = torch.load(last_model_name)
self.gen.load_state_dict(state_dict)
epochs = int(last_model_name[-11:-3])
# Load supervised model.
last_model_name = get_model_list(checkpoint_dir, "sup")
state_dict = torch.load(last_model_name)
self.sup.load_state_dict(state_dict)
# Load discriminator.
last_model_name = get_model_list(checkpoint_dir, "dis")
state_dict = torch.load(last_model_name)
self.dis.load_state_dict(state_dict)
# Load optimizers.
last_model_name = get_model_list(checkpoint_dir, "opt")
state_dict = torch.load(last_model_name)
self.dis_opt.load_state_dict(state_dict['dis'])
self.gen_opt.load_state_dict(state_dict['gen'])
for state in self.dis_opt.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
for state in self.gen_opt.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
# Reinitilize schedulers.
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters,
epochs)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters,
epochs)
print('Resume from epoch %d' % epochs)
return epochs
def train(_config, resume: bool = False, test: bool = False):
print(json.dumps(config, indent=4))
device = torch.device(_config['device'])
os.environ["CUDA_VISIBLE_DEVICES"] = str(device.index)
device = torch.device(0)
dataset = _config['data']['dataset']
model_name = _config['model']['name']
optimizer_name = _config['optimizer']['name']
scheduler_name = _config['scheduler']['name']
loss = utils.get_loss(_config['loss']['name'])
loss.to(device)
model = create_model(dataset, _config['model'][model_name],
_config['model']['stadaptor'], device)
optimizer = utils.get_optimizer(optimizer_name, model.parameters(),
**_config['optimizer'][optimizer_name])
scheduler = None
if scheduler_name is not None:
scheduler = utils.get_scheduler(scheduler_name, optimizer,
**_config['scheduler'][scheduler_name])
save_folder = os.path.join('saves', dataset, _config['name'])
if not resume and not test:
shutil.rmtree(save_folder, ignore_errors=True)
os.makedirs(save_folder)
with open(os.path.join(save_folder, 'config.yaml'), 'w+') as _f:
yaml.safe_dump(_config, _f)
datasets = utils.get_datasets(dataset, _config['data']['input_dim'],
_config['data']['output_dim'])
scaler = utils.ZScoreScaler(datasets['train'].mean, datasets['train'].std)
trainer = utils.OursTrainer(model, loss, scaler, device, optimizer,
**_config['trainer'])
if not test:
utils.train_model(datasets=datasets,
batch_size=_config['data']['batch-size'],
folder=save_folder,
trainer=trainer,
scheduler=scheduler,
epochs=config['epochs'],
early_stop_steps=config['early_stop_steps'])
utils.test_model(datasets=datasets,
batch_size=_config['data']['batch-size'],
trainer=trainer,
folder=save_folder)
def __init__(self, hyperparameters):
super(MUNIT_Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks
self.gen_b = AdaINGen(
hyperparameters['input_dim_b'],
hyperparameters['gen']) # auto-encoder for domain b
self.dis_b = MsImageDis(
hyperparameters['input_dim_b'], hyperparameters['new_size'],
hyperparameters['dis']) # discriminator for domain b
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.style_dim = hyperparameters['gen']['style_dim']
self.reg_param = hyperparameters['reg_param']
self.beta_step = hyperparameters['beta_step']
self.target_kl = hyperparameters['target_kl']
self.gan_type = hyperparameters['gan_type']
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['beta2']
dis_params = list(self.dis_b.parameters())
gen_params = list(self.gen_b.parameters())
self.dis_opt = torch.optim.Adam(
[p for p in dis_params if p.requires_grad],
lr=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam(
[p for p in gen_params if p.requires_grad],
lr=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
# Network weight initialization
self.gen_b.apply(weights_init(hyperparameters['init']))
self.dis_b.apply(weights_init('gaussian'))
# SSIM Loss
self.ssim_loss = pytorch_ssim.SSIM()
def resume(self, gen_dir, dis_dir, opt_dir, iterations, hyperparameters):
# Load generators
state_dict = torch.load(gen_dir)
self.gen_a.load_state_dict(state_dict['a'])
self.gen_b.load_state_dict(state_dict['b'])
print('hi')
# Load discriminators
state_dict = torch.load(dis_dir)
self.dis_a.load_state_dict(state_dict['a'])
self.dis_b.load_state_dict(state_dict['b'])
# Load optimizers
state_dict = torch.load(opt_dir)
self.dis_opt.load_state_dict(state_dict['dis'])
self.gen_opt.load_state_dict(state_dict['gen'])
# Reinitilize schedulers
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters,
iterations)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters,
iterations)
print('Resume from iteration %d' % iterations)
def resume(self, checkpoint_dir, hyperparameters):
# Load generators
last_model_name = get_model_list(checkpoint_dir, "gen")
state_dict = torch.load(last_model_name)
self.gen_a.load_state_dict(state_dict['a'])
self.gen_b.load_state_dict(state_dict['b'])
iterations = int(last_model_name[-11:-3])
# Load discriminators
last_model_name = get_model_list(checkpoint_dir, "dis")
state_dict = torch.load(last_model_name)
self.dis_a.load_state_dict(state_dict['a'])
self.dis_b.load_state_dict(state_dict['b'])
# Load optimizers
state_dict = torch.load(os.path.join(checkpoint_dir, 'optimizer.pt'))
self.dis_opt.load_state_dict(state_dict['dis'])
self.gen_opt.load_state_dict(state_dict['gen'])
# Reinitilize schedulers
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters, iterations)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters, iterations)
print('Resume from iteration %d' % iterations)
return iterations
#! /usr/bin/python
# AD server program
from SocketServer import *
from XAIF import XAIF_Server
from utils import get_scheduler, update_load, shutdown, daemonize, change_workdir, remove_workdir
import sys, os, time, signal, socket, string, getopt, shutil
global upd_pid
# default settings
server_host = os.environ["HOST"]
current_dir = os.environ["PWD"]
server_port = 3351
server_type = "r_xaif"
exe = "/home/derivs/xaif/0.1/linux/xaifbooster"
scheduler_host,scheduler_port,scheduler_client_port=get_scheduler()
basedir = current_dir
timeout = 60
stdin = '/dev/null'
stdout = basedir + '/' + server_type + '_daemon_' + str(os.getpid()) + '.log'
stderr = stdout
verbose = 1
###########################################################################
class AdMixinHandler:
def handle(self):
signal.alarm(timeout)
if verbose: print time.asctime(),"server connected from",self.client_address
host,port = self.client_address
self.req_id = host+"_"+str(port)+"_"+str(int(time.time())) # unique request id
userid = self.check_user(self.recv_header1())
