def setup(self, args, parser=None):
if not 'continue_train' in args:
self.schedulers = [
get_scheduler(optimizer, args) for optimizer in self.optimizers
]
else:
self.load_networks(args.which_epoch)
self.print_networks(args.verbose)
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, _isTrain=False):
self.initialize(opt)
self.mode = opt.mode
if opt.input == 'single_view':
self.num_input = 3
elif opt.input == 'two_view':
self.num_input = 6
elif opt.input == 'two_view_k':
self.num_input = 7
else:
raise ValueError("Unknown input type %s" % opt.input)
if self.mode == 'Ours_Bilinear':
# print(
# '====================================== DIW NETWORK TRAIN FROM %s======================='
# % self.mode)
new_model = hourglass.HourglassModel(self.num_input)
# print(
# '===================Loading Pretrained Model OURS ==================================='
# )
if not _isTrain:
if self.num_input == 7:
model_parameters = self.load_network(
new_model, 'G', 'best_depth_Ours_Bilinear_inc_7')
elif self.num_input == 3:
model_parameters = self.load_network(
new_model, 'G', 'best_depth_Ours_Bilinear_inc_3')
elif self.num_input == 6:
model_parameters = self.load_network(
new_model, 'G', 'best_depth_Ours_Bilinear_inc_6')
else:
print('Something Wrong')
sys.exit()
new_model.load_state_dict(model_parameters)
# new_model = torch.nn.parallel.DataParallel(
# new_model.cuda(), device_ids=range(torch.cuda.device_count()))
self.netG = new_model
else:
print('ONLY SUPPORT Ours_Bilinear')
sys.exit()
self.old_lr = opt.lr
self.netG.train()
if True:
self.criterion_joint = networks.JointLoss(opt)
# initialize optimizers
self.optimizer_G = torch.optim.Adam(
self.netG.parameters(), lr=opt.lr, betas=(0.9, 0.999))
self.scheduler = networks.get_scheduler(self.optimizer_G, opt)
def setup(self, opt, verbose=True):
self.schedulers = [
networks.get_scheduler(optimizer, opt)
for optimizer in self.optimizers
]
self.load_networks(verbose)
if verbose:
self.print_networks()
self.add_mapping_hook()
def setup(self, opt):
if self.isTrain:
self.schedulers = [
networks.get_scheduler(optimizer, opt)
for optimizer in self.optimizers
]
if not self.isTrain or opt.continue_train:
self.load_networks(opt.which_epoch)
self.print_networks(opt.verbose)
def setup(self, args, parser=None):
if self.isTrain:
self.schedulers = [
networks.get_scheduler(optimizer, args)
for optimizer in self.optimizers
]
# if not self.isTrain or args.continue_train:
if not self.isTrain:
self.load_networks(args.epoch)
self.print_networks(args.verbose)
def setup(self, opt):
"""Load and print networks; create schedulers
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
if self.isTrain:
self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
if not self.isTrain or opt.continue_train:
load_suffix = 'iter_%d' % opt.load_iter if opt.load_iter > 0 else opt.epoch
self.load_networks(load_suffix)
self.print_networks(opt.verbose)
def test_scheduler():
from options.gan_options import TrainGANOptions
from models.networks import get_scheduler
model = nn.Linear(10, 10)
optim = torch.optim.Adam(model.parameters(), lr=1e-4)
opt = TrainGANOptions().parse('--gpu_ids -1 --benchmark debug', False,
False, False)
scheduler = get_scheduler(optim, opt)
for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
scheduler.step()
print('epoch: %d, lr: %f' % (epoch, optim.param_groups[0]['lr']))
def setup(self, opt, verbose=True):
"""Load and print networks; create schedulers
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
self.load_networks(verbose=verbose)
if self.isTrain:
self.schedulers = [
networks.get_scheduler(optimizer, opt)
for optimizer in self.optimizers
]
if verbose:
self.print_networks()
def setup(self):
"""Load and print networks; create schedulers
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
load_prefix = self.opt.load_flag if self.opt.load_flag else 'latest'
epoch = self.load_networks(load_prefix)
if self.isTrain:
self.schedulers = [
get_scheduler(optimizer, self.opt)
for optimizer in self.optimizers
]
self.print_networks(self.opt.print_network)
return epoch
def setup(self, opt, verbose=True):
self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
self.load_networks(verbose)
if verbose:
self.print_networks()
if self.opt.lambda_distill > 0:
def get_activation(mem, name):
def get_output_hook(module, input, output):
mem[name] = output
return get_output_hook
def add_hook(net, mem, mapping_layers):
for n, m in net.named_modules():
if n in mapping_layers:
m.register_forward_hook(get_activation(mem, n))
add_hook(self.netG_teacher, self.Tacts, self.mapping_layers)
add_hook(self.netG_student, self.Sacts, self.mapping_layers)
def initialize(self, opt):
BaseModel.initialize(self, opt)
self.isTrain = opt.isTrain
# load/define networks
self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf,
opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids)
if self.isTrain:
use_sigmoid = opt.no_lsgan
self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf,
opt.which_model_netD,
opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, self.gpu_ids)
if not self.isTrain or opt.continue_train:
self.load_network(self.netG, 'G', opt.which_epoch)
if self.isTrain:
self.load_network(self.netD, 'D', opt.which_epoch)
if self.isTrain:
self.fake_AB_pool = ImagePool(opt.pool_size)
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)
self.criterionL1 = torch.nn.L1Loss()
# initialize optimizers
self.schedulers = []
self.optimizers = []
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)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
if self.isTrain:
networks.print_network(self.netD)
print('-----------------------------------------------')
def initialize(self, opt):
BaseModel.initialize(self, opt)
self.opt = opt
self.isTrain = opt.isTrain
# specify the training losses you want to print out. The program will call base_model.get_current_losses
self.loss_names = [
'G_GAN', 'G_L1', 'D', 'style', 'content', 'tv', 'hole', 'valid'
]
# specify the images you want to save/display. The program will call base_model.get_current_visuals
if self.opt.show_flow:
self.visual_names = ['real_A', 'fake_B', 'real_B', 'flow_srcs']
else:
self.visual_names = [
'real_input', 'fake_B', 'real_GTimg', 'mask_global',
'output_comp'
]
# specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
if self.isTrain:
self.model_names = ['G', 'D']
else: # during test time, only load Gs
self.model_names = ['G']
#
# self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf,
# opt.which_model_netG, opt, opt.norm, opt.use_spectral_norm_G, opt.init_type, self.gpu_ids, opt.init_gain)
self.netG = PConvUNet().to(self.device)
print(self.netG)
if self.isTrain:
use_sigmoid = False
if opt.gan_type == 'vanilla':
use_sigmoid = True # only vanilla GAN using BCECriterion
# don't use cGAN
self.netD = networks.define_D(opt.input_nc, opt.ndf,
opt.which_model_netD, opt.n_layers_D,
opt.norm, use_sigmoid,
opt.use_spectral_norm_D,
opt.init_type, self.gpu_ids,
opt.init_gain)
# add style extractor
self.vgg16_extractor = util.VGG16FeatureExtractor().to(self.gpu_ids[0])
self.vgg16_extractor = torch.nn.DataParallel(self.vgg16_extractor,
self.gpu_ids)
if self.isTrain:
self.old_lr = opt.lr
# define loss functions
self.criterionGAN = networks.GANLoss(gan_type=opt.gan_type).to(
self.device)
self.criterionL1 = torch.nn.L1Loss()
self.criterionL1_mask = networks.Discounted_L1(opt).to(
self.device
) # make weights/buffers transfer to the correct device
# VGG loss
self.criterionL2_style_loss = torch.nn.MSELoss()
self.criterionL2_content_loss = torch.nn.MSELoss()
# TV loss
self.tv_criterion = networks.TVLoss(self.opt.tv_weight)
# initialize optimizers
self.schedulers = []
self.optimizers = []
if self.opt.gan_type == 'wgan_gp':
opt.beta1 = 0
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.9))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.9))
else:
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)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
if not self.isTrain or opt.continue_train:
self.load_networks(opt.which_epoch)
self.print_networks(opt.verbose)
dataloader = DataLoader(opt)
model = RegresserModel(opt)
if not opt.no_cuda:
if len(opt.gpu_ids) > 1:
model = model.cuda()
model = nn.DataParallel(model, device_ids=opt.gpu_ids)
net_dict = model.state_dict()
else:
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
net_dict = model.state_dict()
optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.001)
scheduler = get_scheduler(optimizer, opt)
if opt.resume_path:
if os.path.isfile(opt.resume_path):
print("=> loading checkpoint '{}'".format(opt.resume_path))
checkpoint = torch.load(opt.resume_path)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(opt.resume_path, checkpoint['epoch']))
train(dataloader=dataloader,
model=model,
optimizer=optimizer,
scheduler=scheduler,
total_epochs=opt.nepoch,
def initialize(self, opt):
BaseModel.initialize(self, opt)
nb = opt.batchSize
size = opt.fineSize
self.target_weight = []
self.input_A = self.Tensor(nb, opt.input_nc, size, size)
self.input_B = self.Tensor(nb, opt.output_nc, size, size)
self.input_C = self.Tensor(nb, opt.output_nc, size, size)
self.input_C_sr = self.Tensor(nb, opt.output_nc, size, size)
if opt.aux:
self.A_aux = self.Tensor(nb, opt.input_nc, size, size)
self.B_aux = self.Tensor(nb, opt.output_nc, size, size)
self.C_aux = self.Tensor(nb, opt.output_nc, size, size)
self.netE_A = networks.define_G(opt.input_nc,
opt.output_nc,
opt.ngf,
'ResnetEncoder_my',
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
opt=opt,
n_downsampling=2)
mult = self.netE_A.get_mult()
self.netE_C = networks.define_G(opt.input_nc,
opt.output_nc,
64,
'ResnetEncoder_my',
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
opt=opt,
n_downsampling=3)
self.net_D = networks.define_G(opt.input_nc,
opt.output_nc,
opt.ngf,
'ResnetDecoder_my',
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
opt=opt,
mult=mult)
mult = self.net_D.get_mult()
self.net_Dc = networks.define_G(opt.input_nc,
opt.output_nc,
opt.ngf,
'ResnetDecoder_my',
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
opt=opt,
mult=mult,
n_upsampling=1)
self.netG_A = networks.define_G(opt.input_nc,
opt.output_nc,
opt.ngf,
'GeneratorLL',
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
opt=opt,
mult=mult)
mult = self.net_Dc.get_mult()
self.netG_C = networks.define_G(opt.input_nc,
opt.output_nc,
opt.ngf,
'GeneratorLL',
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
opt=opt,
mult=mult)
# self.netG_A_running = networks.define_G(opt.input_nc, opt.output_nc,
# opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt)
# set_eval(self.netG_A_running)
# accumulate(self.netG_A_running, self.netG_A, 0)
# self.netG_B = networks.define_G(opt.output_nc, opt.input_nc,
# opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt)
# self.netG_B_running = networks.define_G(opt.output_nc, opt.input_nc,
# opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt)
# set_eval(self.netG_B_running)
# accumulate(self.netG_B_running, self.netG_B, 0)
if self.isTrain:
use_sigmoid = opt.no_lsgan
self.netD_A = networks.define_D(opt.output_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
opt=opt)
# self.netD_B = networks.define_D(opt.input_nc, opt.ndf,
# opt.which_model_netD,
# opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, self.gpu_ids, opt=opt)
print('---------- Networks initialized -------------')
# networks.print_network(self.netG_B, opt, (opt.input_nc, opt.fineSize, opt.fineSize))
networks.print_network(self.netE_C, opt,
(opt.input_nc, opt.fineSize, opt.fineSize))
networks.print_network(
self.net_D, opt, (opt.ngf * 4, opt.fineSize / 4, opt.fineSize / 4))
networks.print_network(self.net_Dc, opt,
(opt.ngf, opt.CfineSize / 2, opt.CfineSize / 2))
# networks.print_network(self.netG_B, opt)
if self.isTrain:
networks.print_network(self.netD_A, opt)
# networks.print_network(self.netD_B, opt)
print('-----------------------------------------------')
if not self.isTrain or opt.continue_train:
print('Loaded model')
which_epoch = opt.which_epoch
self.load_network(self.netG_A, 'G_A', which_epoch)
self.load_network(self.netG_B, 'G_B', which_epoch)
if self.isTrain:
self.load_network(self.netG_A_running, 'G_A', which_epoch)
self.load_network(self.netG_B_running, 'G_B', which_epoch)
self.load_network(self.netD_A, 'D_A', which_epoch)
self.load_network(self.netD_B, 'D_B', which_epoch)
if self.isTrain and opt.load_path != '':
print('Loaded model from load_path')
which_epoch = opt.which_epoch
load_network_with_path(self.netG_A,
'G_A',
opt.load_path,
epoch_label=which_epoch)
load_network_with_path(self.netG_B,
'G_B',
opt.load_path,
epoch_label=which_epoch)
load_network_with_path(self.netD_A,
'D_A',
opt.load_path,
epoch_label=which_epoch)
load_network_with_path(self.netD_B,
'D_B',
opt.load_path,
epoch_label=which_epoch)
if self.isTrain:
self.old_lr = opt.lr
self.fake_A_pool = ImagePool(opt.pool_size)
self.fake_B_pool = ImagePool(opt.pool_size)
self.fake_C_pool = ImagePool(opt.pool_size)
# define loss functions
if len(self.target_weight) == opt.num_D:
print(self.target_weight)
self.criterionGAN = networks.GANLoss(
use_lsgan=not opt.no_lsgan,
tensor=self.Tensor,
target_weight=self.target_weight,
gan=opt.gan)
else:
self.criterionGAN = networks.GANLoss(
use_lsgan=not opt.no_lsgan,
tensor=self.Tensor,
gan=opt.gan)
self.criterionCycle = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
self.criterionColor = networks.ColorLoss()
# initialize optimizers
self.optimizer_G = torch.optim.Adam(itertools.chain(
self.netE_A.parameters(), self.net_D.parameters(),
self.netG_A.parameters(), self.net_Dc.parameters(),
self.netG_C.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizer_AE = torch.optim.Adam(itertools.chain(
self.netE_C.parameters(), self.net_D.parameters(),
self.net_Dc.parameters(), self.netG_C.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizer_G_A_sr = torch.optim.Adam(itertools.chain(
self.netE_A.parameters(), self.net_D.parameters(),
self.net_Dc.parameters(), self.netG_C.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizer_AE_sr = torch.optim.Adam(itertools.chain(
self.netE_C.parameters(), self.net_D.parameters(),
self.netG_A.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizer_D_A = torch.optim.Adam(self.netD_A.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
# self.optimizer_D_B = torch.optim.Adam(self.netD_B.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_AE)
# self.optimizers.append(self.optimizer_G_A_sr)
self.optimizers.append(self.optimizer_AE_sr)
self.optimizers.append(self.optimizer_D_A)
# self.optimizers.append(self.optimizer_D_B)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
def __init__(self, opt, data=None):
super(Skipganomaly, self).__init__(opt, data)
##
# -- Misc attributes
self.add_noise = True
self.epoch = 0
self.times = []
self.total_steps = 0
##
# Create and initialize networks.
self.netg = define_G(self.opt,
norm='batch',
use_dropout=False,
init_type='normal')
self.netd = define_D(self.opt,
norm='batch',
use_sigmoid=False,
init_type='normal')
##
if self.opt.resume != '':
print("\nLoading pre-trained networks.")
self.opt.iter = torch.load(
os.path.join(self.opt.resume, 'netG.pth'))['epoch']
self.netg.load_state_dict(
torch.load(os.path.join(self.opt.resume,
'netG.pth'))['state_dict'])
self.netd.load_state_dict(
torch.load(os.path.join(self.opt.resume,
'netD.pth'))['state_dict'])
print("\tDone.\n")
if self.opt.verbose:
print(self.netg)
print(self.netd)
##
# Loss Functions
self.l_adv = nn.BCELoss()
self.l_con = nn.L1Loss()
self.l_lat = l2_loss
##
# Initialize input tensors.
self.input = torch.empty(size=(self.opt.batchsize, 3, self.opt.isize,
self.opt.isize),
dtype=torch.float32,
device=self.device)
self.noise = torch.empty(size=(self.opt.batchsize, 3, self.opt.isize,
self.opt.isize),
dtype=torch.float32,
device=self.device)
self.label = torch.empty(size=(self.opt.batchsize, ),
dtype=torch.float32,
device=self.device)
self.gt = torch.empty(size=(opt.batchsize, ),
dtype=torch.long,
device=self.device)
self.fixed_input = torch.empty(size=(self.opt.batchsize, 3,
self.opt.isize, self.opt.isize),
dtype=torch.float32,
device=self.device)
self.real_label = torch.ones(size=(self.opt.batchsize, ),
dtype=torch.float32,
device=self.device)
self.fake_label = torch.zeros(size=(self.opt.batchsize, ),
dtype=torch.float32,
device=self.device)
##
# Setup optimizer
if self.opt.isTrain:
self.netg.train()
self.netd.train()
self.optimizers = []
self.optimizer_d = optim.Adam(self.netd.parameters(),
lr=self.opt.lr,
betas=(self.opt.beta1, 0.999))
self.optimizer_g = optim.Adam(self.netg.parameters(),
lr=self.opt.lr,
betas=(self.opt.beta1, 0.999))
self.optimizers.append(self.optimizer_d)
self.optimizers.append(self.optimizer_g)
self.schedulers = [
get_scheduler(optimizer, opt) for optimizer in self.optimizers
]
def set_scheduler(self, opts, ep):
assert (self.isTrain)
self.schedulers = [
networks.get_scheduler(optimizer, opts, ep)
for optimizer in self.optimizers
]
def __init__(self, opts):
BaseModel.__init__(self, opts)
lr = self.opt.lr
self.model_names = ['G_A', 'G_B']
self.loss_names = [
'd_total', 'g_total', 'g_rec_x_a', 'g_rec_x_b', 'g_rec_s_a',
'g_rec_s_b', 'g_rec_c_a', 'g_rec_c_b', 'g_adv_a', 'g_adv_b'
]
self.visual_names = []
# Initiate the networks
if self.isTrain:
self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
self.netD_A = init_net(MsImageDis(self.opt.input_dim_a,
self.opt.dis),
init_type=self.opt.init,
gpu_ids=self.gpu_ids)
self.netD_B = init_net(MsImageDis(self.opt.input_dim_b,
self.opt.dis),
init_type=self.opt.init,
gpu_ids=self.gpu_ids)
self.netG_A = init_net(AdaINGen(self.opt.input_dim_a, self.opt.gen),
init_type=self.opt.init,
gpu_ids=self.gpu_ids)
self.netG_B = init_net(AdaINGen(self.opt.input_dim_b, self.opt.gen),
init_type=self.opt.init,
gpu_ids=self.gpu_ids)
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.style_dim = self.opt.gen['style_dim']
# fix the noise used in sampling
display_size = self.opt.display_size
self.s_a_fixed = torch.randn(display_size, self.style_dim, 1,
1).to(self.device)
self.s_b_fixed = torch.randn(display_size, self.style_dim, 1,
1).to(self.device)
if self.isTrain:
# Setup the optimizers
d_params = list(self.netD_A.parameters()) + list(
self.netD_B.parameters())
g_params = list(self.netG_A.parameters()) + list(
self.netG_B.parameters())
self.optimizer_D = torch.optim.Adam(
[p for p in d_params if p.requires_grad],
lr=lr,
betas=(self.opt.beta1, self.opt.beta2),
weight_decay=self.opt.weight_decay)
self.optimizer_G = torch.optim.Adam(
[p for p in g_params if p.requires_grad],
lr=lr,
betas=(self.opt.beta1, self.opt.beta2),
weight_decay=self.opt.weight_decay)
self.optimizer_names = ['optimizer_D', 'optimizer_G']
self.optimizers.append(self.optimizer_D)
self.optimizers.append(self.optimizer_G)
self.scheduler_D = get_scheduler(self.optimizer_D, self.opt)
self.scheduler_G = get_scheduler(self.optimizer_G, self.opt)
self.schedulers = [self.scheduler_D, self.scheduler_G]
# Load VGG model if needed
if (self.opt.vgg_w is not None) and self.opt.vgg_w > 0:
self.vgg = load_vgg16(self.opt.vgg_model_path + '/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
def initialize(self, opt):
BaseModel.initialize(self, opt)
self.count = 0
input_depth = opt.input_nc
output_depth = opt.output_nc
self.net_shared = skip(input_depth,
num_channels_down=[64, 128, 256, 256, 256],
num_channels_up=[64, 128, 256, 256, 256],
num_channels_skip=[4, 4, 4, 4, 4],
upsample_mode=[
'nearest', 'nearest', 'bilinear',
'bilinear', 'bilinear'
],
need_sigmoid=True,
need_bias=True,
pad='reflection')
self.netDec_a = ResNet_decoders(opt.ngf, output_depth)
self.netDec_b = ResNet_decoders(opt.ngf, output_depth)
self.net_input = self.get_noise(input_depth, 'noise',
(self.opt.fineSize, self.opt.fineSize))
self.net_input_saved = self.net_input.detach().clone()
self.noise = self.net_input.detach().clone()
use_sigmoid = opt.no_lsgan
self.netD_b = networks.define_D(opt.output_nc, opt.ndf,
opt.which_model_netD, opt.n_layers_D,
opt.norm, use_sigmoid, opt.init_type,
self.gpu_ids)
if not opt.dont_load_pretrained_autoencoder:
which_epoch = opt.which_epoch
self.load_network(self.netDec_b, 'Dec_b', which_epoch)
self.load_network(self.net_shared, 'Net_shared', which_epoch)
self.load_network(self.netD_b, 'D', which_epoch)
if len(self.gpu_ids) > 0:
dtype = torch.cuda.FloatTensor
self.net_input = self.net_input.type(dtype).detach()
self.net_shared = self.net_shared.type(dtype)
self.netDec_a = self.netDec_a.type(dtype)
self.netDec_b = self.netDec_b.type(dtype)
self.netD_b = self.netD_b.type(dtype)
self.fake_A_pool = ImagePool(opt.pool_size)
self.fake_B_pool = ImagePool(opt.pool_size)
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
self.mse = torch.nn.MSELoss()
# initialize optimizers
self.optimizer_Net = torch.optim.Adam(
itertools.chain(self.net_shared.parameters(),
self.netDec_a.parameters()),
lr=0.007,
betas=(opt.beta1,
0.999)) # skip 0.01 # OST 0.001 # skip large 0.007
self.optimizer_Dec_b = torch.optim.Adam(
self.netDec_b.parameters(), lr=0.000007,
betas=(opt.beta1,
0.999)) # OST 0.000007 skip 0.00002 skip large 0.000007
self.optimizer_D_b = torch.optim.Adam(self.netD_b.parameters(),
lr=0.0002,
betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_Net)
self.optimizers.append(self.optimizer_Dec_b)
self.optimizers.append(self.optimizer_D_b)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
def initialize(self, opt):
BaseModel.initialize(self, opt)
self.opt = opt
self.isTrain = opt.isTrain
# specify the training losses you want to print out. The program will call base_model.get_current_losses
self.loss_names = ['G_GAN', 'G_L1', 'D']
# specify the images you want to save/display. The program will call base_model.get_current_visuals
if self.opt.show_flow:
self.visual_names = ['real_A', 'fake_B', 'real_B', 'flow_srcs']
else:
self.visual_names = ['real_A', 'fake_B', 'real_B']
# specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
if self.isTrain:
self.model_names = ['G', 'D']
else: # during test time, only load Gs
self.model_names = ['G']
# batchsize should be 1 for mask_global
self.mask_global = torch.ByteTensor(1, 1, \
opt.fineSize, opt.fineSize)
# Here we need to set an artificial mask_global(not to make it broken, so center hole is ok.)
self.mask_global.zero_()
self.mask_global[:, :, int(self.opt.fineSize/4) + self.opt.overlap : int(self.opt.fineSize/2) + int(self.opt.fineSize/4) - self.opt.overlap,\
int(self.opt.fineSize/4) + self.opt.overlap: int(self.opt.fineSize/2) + int(self.opt.fineSize/4) - self.opt.overlap] = 1
self.mask_type = opt.mask_type
self.gMask_opts = {}
self.wgan_gp = False
# added for wgan-gp
if opt.gan_type == 'wgan_gp':
self.gp_lambda = opt.gp_lambda
self.ncritic = opt.ncritic
self.wgan_gp = True
if len(opt.gpu_ids) > 0:
self.use_gpu = True
self.mask_global = self.mask_global.to(self.device)
# load/define networks
# self.ng_innerCos_list is the constraint list in netG inner layers.
# self.ng_mask_list is the mask list constructing shift operation.
if opt.add_mask2input:
input_nc = opt.input_nc + 1
else:
input_nc = opt.input_nc
self.netG, self.ng_innerCos_list, self.ng_shift_list = networks.define_G(
input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt,
self.mask_global, opt.norm, opt.use_dropout,
opt.use_spectral_norm_G, opt.init_type, self.gpu_ids,
opt.init_gain) # add opt, we need opt.shift_sz and other stuffs
if self.isTrain:
use_sigmoid = False
if opt.gan_type == 'vanilla':
use_sigmoid = True # only vanilla GAN using BCECriterion
# don't use cGAN
self.netD = networks.define_D(opt.input_nc, opt.ndf,
opt.which_model_netD, opt.n_layers_D,
opt.norm, use_sigmoid,
opt.use_spectral_norm_D,
opt.init_type, self.gpu_ids,
opt.init_gain)
if self.isTrain:
self.old_lr = opt.lr
# define loss functions
self.criterionGAN = networks.GANLoss(gan_type=opt.gan_type).to(
self.device)
self.criterionL1 = torch.nn.L1Loss()
self.criterionL1_mask = util.Discounted_L1(opt).to(
self.device
) # make weights/buffers transfer to the correct device
# initialize optimizers
self.schedulers = []
self.optimizers = []
if self.wgan_gp:
opt.beta1 = 0
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))
else:
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)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
if not self.isTrain or opt.continue_train:
self.load_networks(opt.which_epoch)
self.print_networks(opt.verbose)
def initialize(self, opt):
BaseModel.initialize(self, opt)
self.opt = opt
self.isTrain = opt.isTrain
# specify the training losses you want to print out. The program will call base_model.get_current_losses
self.loss_names = ['G_GAN', 'G_L1', 'D', 'style', 'content', 'tv']
# specify the images you want to save/display. The program will call base_model.get_current_visuals
if self.opt.show_flow:
self.visual_names = ['real_A', 'fake_B', 'real_B', 'flow_srcs']
else:
self.visual_names = ['real_A', 'fake_B', 'real_B']
# specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
if self.isTrain:
self.model_names = ['G', 'D']
else: # during test time, only load Gs
self.model_names = ['G']
# batchsize should be 1 for mask_global
self.mask_global = torch.zeros((self.opt.batchSize, 1, \
opt.fineSize, opt.fineSize), dtype=torch.bool)
# Here we need to set an artificial mask_global(center hole is ok.)
self.mask_global.zero_()
# self.mask_global[:, :, int(self.opt.fineSize/4) + self.opt.overlap : int(self.opt.fineSize/2) + int(self.opt.fineSize/4) - self.opt.overlap,\
# int(self.opt.fineSize/4) + self.opt.overlap: int(self.opt.fineSize/2) + int(self.opt.fineSize/4) - self.opt.overlap] = 1
self.mask_global[:, :, int(self.opt.fineSize * 3 / 8) + self.opt.overlap: int(self.opt.fineSize / 2) + int(self.opt.fineSize / 8) - self.opt.overlap, \
int(self.opt.fineSize * 3 / 8) + self.opt.overlap: int(self.opt.fineSize / 2) + int(self.opt.fineSize / 8) - self.opt.overlap] = 1
if len(opt.gpu_ids) > 0:
self.mask_global = self.mask_global.to(self.device)
# load/define networks
# self.ng_innerCos_list is the guidance loss list in netG inner layers.
# self.ng_shift_list is the mask list constructing shift operation.
if opt.add_mask2input:
input_nc = opt.input_nc + 1
else:
input_nc = opt.input_nc
self.netG, self.ng_innerCos_list, self.ng_shift_list = networks.define_G(
input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt,
self.mask_global, opt.norm, opt.use_spectral_norm_G, opt.init_type,
self.gpu_ids, opt.init_gain)
if self.isTrain:
use_sigmoid = False
if opt.gan_type == 'vanilla':
use_sigmoid = True # only vanilla GAN using BCECriterion
# don't use cGAN
self.netD = networks.define_D(1, opt.ndf, opt.which_model_netD,
opt.n_layers_D, opt.norm,
use_sigmoid, opt.use_spectral_norm_D,
opt.init_type, self.gpu_ids,
opt.init_gain)
# add style extractor
self.vgg16_extractor = util.VGG16FeatureExtractor().to(self.gpu_ids[0])
self.vgg16_extractor = torch.nn.DataParallel(self.vgg16_extractor,
self.gpu_ids)
if self.isTrain:
self.old_lr = opt.lr
# define loss functions
self.criterionGAN = networks.GANLoss(gan_type=opt.gan_type).to(
self.device)
self.criterionL1 = torch.nn.L1Loss()
self.criterionL1_mask = networks.Discounted_L1(opt).to(
self.device
) # make weights/buffers transfer to the correct device
# VGG loss
self.criterionL2_style_loss = torch.nn.MSELoss()
self.criterionL2_content_loss = torch.nn.MSELoss()
# TV loss
self.tv_criterion = networks.TVLoss(self.opt.tv_weight)
# initialize optimizers
self.schedulers = []
self.optimizers = []
if self.opt.gan_type == 'wgan_gp':
opt.beta1 = 0
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.9))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.9))
else:
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)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
if not self.isTrain or opt.continue_train:
self.load_networks(opt.which_epoch)
self.print_networks(opt.verbose)
def __init__(self, args):
self.gpu_ids=[0]
self.isTrain = True
self.checkpoints_dir = './checkpoints'
self.which_epoch = 'latest' # which epoch to load? set to latest to use latest cached model
self.args = args
# self.name = 'G_GAN_%s_lambdar_%s_lambdas_%s_alpha_%s' % (self.args.lambda_d, self.args.lambda_r, self.args.lambda_s, self.args.alpha)
self.name = 'Res_convolution_Gram'
expr_dir = os.path.join(self.checkpoints_dir, self.name)
if not os.path.exists(expr_dir):
os.makedirs(expr_dir)
self.save_dir = os.path.join(self.checkpoints_dir, self.name)
self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
# specify the training losses you want to print out. The program will call base_model.get_current_losses
self.loss_names = ['G_style', 'G_content']
# 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 = ['C', 'E', 'G']
self.visual_names = ['A', 'B', 'C', 'R']
self.input_nc = 3
self.output_nc = 3
self.ndf = 64 #number of filters in the first layer of discriminator
self.ngf = 64
use_sigmoid = False
# define networks
self.netCA = networks.define_channel_attention(self.gpu_ids)
self.netKA = networks.define_kernel_attention(self.gpu_ids)
self.netSA = networks.define_spatial_attention(self.gpu_ids)
self.netC = networks.define_Convolution(self.gpu_ids)
self.netKC = networks.define_K_Convolution(self.gpu_ids)
self.netVGG = networks.define_VGG()
self.netE = networks.define_E(self.input_nc, self.ngf, self.gpu_ids)
self.netG = networks.define_G(self.input_nc, self.output_nc, self.ngf, self.gpu_ids)
self.criterionMSE = torch.nn.MSELoss()
self.criterionL1 = torch.nn.L1Loss()
# initialize optimizers
self.schedulers = []
self.optimizers = []
self.optimizer_CA = torch.optim.Adam(self.netCA.parameters(),
lr=0.0002, betas=(0.5, 0.999))
self.optimizer_KA = torch.optim.Adam(self.netKA.parameters(),
lr=0.0002, betas=(0.5, 0.999))
self.optimizer_SA = torch.optim.Adam(self.netSA.parameters(),
lr=0.0002, betas=(0.5, 0.999))
self.optimizer_C = torch.optim.Adam(self.netC.parameters(),
lr=0.0002, betas=(0.5, 0.999))
self.optimizer_KC = torch.optim.Adam(self.netKC.parameters(),
lr=0.0002, betas=(0.5, 0.999))
self.optimizer_E = torch.optim.Adam(self.netE.parameters(),
lr=0.0002, betas=(0.5, 0.999))
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=0.0002, betas=(0.5, 0.999))
self.optimizers.append(self.optimizer_CA)
self.optimizers.append(self.optimizer_KA)
self.optimizers.append(self.optimizer_SA)
self.optimizers.append(self.optimizer_C)
self.optimizers.append(self.optimizer_KC)
self.optimizers.append(self.optimizer_E)
self.optimizers.append(self.optimizer_G)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, lr_policy='lambda', epoch_count=1, niter=100, niter_decay=100, lr_decay_iters=50))
if not self.isTrain or args.continue_train:
self.load_networks(self.which_epoch)
self.print_networks()
