def define_losses(self):
opt = self.opt
# set loss functions
if self.isTrain or opt.finetune:
self.fake_pool = ImagePool(0)
self.old_lr = opt.lr
# define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.Tensor,
opt=opt)
self.criterionFeat = torch.nn.L1Loss()
self.criterionFlow = networks.MaskedL1Loss()
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(opt, self.gpu_ids)
# Names so we can breakout loss
self.loss_names_G = [
'G_GAN', 'G_GAN_Feat', 'G_VGG', 'Gf_GAN', 'Gf_GAN_feat',
'GT_GAN', 'GT_GAN_Feat', 'F_Flow', 'F_Warp', 'W'
]
self.loss_names_D = [
'D_real', 'D_fake', 'Df_real', 'Df_fake', 'DT_real', 'DT_fake'
]
self.loss_names = self.loss_names_G + self.loss_names_D
def __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.alpha = 1
self.net = torch.nn.ModuleDict(self.initialize_networks(opt))
# set loss functions
if opt.isTrain:
self.vggnet_fix = networks.correspondence.VGG19_feature_color_torchversion(
vgg_normal_correct=opt.vgg_normal_correct)
self.vggnet_fix.load_state_dict(
torch.load('models/vgg19_conv.pth'))
self.vggnet_fix.eval()
for param in self.vggnet_fix.parameters():
param.requires_grad = False
self.vggnet_fix.to(self.opt.gpu_ids[0])
self.contextual_forward_loss = networks.ContextualLoss_forward(opt)
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
self.MSE_loss = torch.nn.MSELoss()
if opt.which_perceptual == '5_2':
self.perceptual_layer = -1
elif opt.which_perceptual == '4_2':
self.perceptual_layer = -2
def __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() else torch.ByteTensor
self.netG, self.netD, self.netE = self.initialize_networks(opt)
# set loss functions
if opt.isTrain:
if self.opt.dataset_mode == 'cityscapes':
self.background_list = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 14, 15, 16, 21, 23]
elif self.opt.dataset_mode == 'coco':
self.background_list = [114, 115, 116, 117, 118, 124, 126, 127, 128, 135, 136,
144, 145, 147, 148, 149, 151, 154, 155, 157, 158, 159, 160, 161, 162, 164, 169,
170, 171, 172, 173, 174, 175, 176, 177, 178, 182]
elif self.opt.dataset_mode == 'ade20k':
self.background_list = [0, 1, 2, 3, 4, 5, 6, 9, 11, 13, 16, 17, 21, 25, 26, 29, 32,
46, 48, 51, 52, 54, 60, 61, 68, 69, 81, 91, 94, 101, 128]
self.criterionGAN = networks.GANLoss(opt.gan_mode, tensor=self.FloatTensor, opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
if opt.use_vae:
self.KLDLoss = networks.KLDLoss()
def __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
torch_version = torch.__version__.split('.')
if int(torch_version[1]) >= 2:
self.ByteTensor = torch.cuda.BoolTensor if self.use_gpu() \
else torch.BoolTensor
else:
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.netG, self.netD, self.netE = self.initialize_networks(opt)
self.amp = True if AMP and opt.use_amp and opt.isTrain else False
# set loss functions
if opt.isTrain:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
if opt.use_vae:
self.KLDLoss = networks.KLDLoss()
def __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.netG, self.netD, self.netE, self.netIG, self.netFE, self.netB, self.netD2, self.netSIG = self.initialize_networks(
opt)
# set loss functions
if opt.isTrain:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
self.criterionGANFeat = networks.GANFeatLoss(self.opt)
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt)
if opt.use_vae:
self.KLDLoss = networks.KLDLoss()
if not opt.no_orient_loss:
self.criterionOrient = networks.L1OLoss(self.opt)
self.criterionStyleContent = networks.StyleContentLoss(opt)
# the loss of RGB background
self.criterionBackground = networks.RGBBackgroundL1Loss()
self.criterionRGBL1 = nn.L1Loss()
self.criterionRGBL2 = nn.MSELoss()
self.criterionLabL1 = networks.LabColorLoss(opt)
if opt.unpairTrain:
self.criterionHairAvgLab = networks.HairAvgLabLoss(opt)
def __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.netG, self.netD, self.netE = self.initialize_networks(opt)
# set loss functions
if opt.isTrain:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
self.criterionFeat = nn.L1Loss()
self.criterionL1 = nn.L1Loss()
self.criterionL2 = nn.MSELoss()
self.criterionOpenEDS = MSECalculator.calculate_mse_for_tensors
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
if opt.lambda_style_feat > 0:
# loss on style feature maps
self.criterion_style_feat = nn.MSELoss()
if opt.lambda_style_w > 0:
# Loss on latent style code
self.criterion_style_w = nn.MSELoss()
if opt.lambda_gram > 0:
self.criterion_gram = networks.StyleLoss()
self.reset_loss_log()
def __init__(self, opt):
super(RotateModel, self).__init__()
self.opt = opt
self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.real_image = torch.zeros(opt.batchSize, 3, opt.crop_size,
opt.crop_size)
self.input_semantics = torch.zeros(opt.batchSize, 3, opt.crop_size,
opt.crop_size)
self.netG, self.netD, self.netE, self.netD_rotate = self.initialize_networks(
opt)
# set loss functions
if opt.isTrain:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt)
if opt.use_vae:
self.KLDLoss = networks.KLDLoss()
def __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.netGbg, self.netG, self.netD = self.initialize_networks(opt)
# set loss functions
if opt.isTrain:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
if opt.two_step_model:
self.criterionMask = torch.nn.BCEWithLogitsLoss()
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
if opt.embed_captions and not opt.use_precomputed_captions:
print('Loading pretrained BERT model...')
from transformers import BertModel
self.bert = BertModel.from_pretrained(
'bert-base-uncased', output_hidden_states=True).eval()
if self.use_gpu():
self.bert = self.bert.cuda()
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)
# 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 initialize_loss(self, opt):
self.criterionReg = torch.nn.L1Loss()
self.criterionCycle = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
self.criterionGAN = networks.GANLoss(opt['gan_mode'],
tensor=self.FloatTensor,
opt=opt)
def __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.netG, self.netD, self.netE, self.netF = self.initialize_networks(
opt)
# set loss functions
if opt.isTrain:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
if opt.booru_loss and opt.no_vgg_loss:
self.criterionBooru = networks.BooruLoss(self.opt.gpu_ids)
if opt.use_vae:
self.KLDLoss = networks.KLDLoss()
if opt.L2_loss:
self.L2Loss = torch.nn.MSELoss()
if opt.L1_loss:
self.L1Loss = torch.nn.L1Loss()
if self.opt.hsv_tv:
self.hsvTVLoss = networks.HSVTVLoss()
if self.opt.high_sv:
self.HighSVLoss = networks.HighSVLoss()
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 initialize_loss(self, opt):
self.criterionGAN = networks.GANLoss(opt['gan_mode'],
tensor=self.FloatTensor,
opt=opt)
# if self.use_gpu:
# self.criterionGAN = DataParallelWithCallback(
# self.criterionGAN, device_ids=opt['gpu_ids'])
self.criterionReg = torch.nn.L1Loss()
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 set_losses(self):
# set loss functions
opt = self.opt
if opt.isTrain:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
if opt.use_vae:
self.KLDLoss = networks.KLDLoss()
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 __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__(self, opts, model, vgg, output_dim):
super(AdaINOptimize, self).__init__()
# parameter
self.refine_iter = opts.n_optimize
self.isTrain = (opts.reg_phase == 'train')
self.gpu_ids = opts.gpu_ids
self.device = torch.device('cuda:{}'.format(
self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
self.refine_lr = 0.1
if self.isTrain:
self.lambda_reconstruction = opts.lambda_reconstruction
# model and vgg model (dont require gradient calculation)
self.num_params = model.netG.module.num_adain_params
self.learner = model
self.learner.eval()
self.vgg = vgg
self.vgg.eval()
self.set_requires_grad(self.learner, False)
self.set_requires_grad(self.vgg, False)
# regularizer
self.net = L2Regularizer(self.num_params, opts.batch_size)
# discriminator
if self.isTrain:
self.netD = networks.define_D(output_dim,
64,
netD='basic_256_multi',
norm='instance',
num_Ds=2,
gpu_ids=self.gpu_ids)
ckt = torch.load(opts.load, map_location=str(self.device))
self.cktD = ckt['bicycle_{}'.format(opts.reg_stage)]['netD']
self.netD.load_state_dict(self.cktD)
self.criterionGAN = networks.GANLoss(gan_mode=opts.gan_mode).to(
self.device)
# optimizer
if self.isTrain:
self.opt = torch.optim.Adam(self.net.parameters(),
lr=opts.lr,
weight_decay=opts.weight_decay)
self.optD = torch.optim.Adam(self.netD.parameters(),
lr=2e-4,
betas=(0.5, 0.999))
return
def __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() else torch.ByteTensor
self.netG, self.netD, self.netE = self.initialize_networks(opt)
# set loss functions
if opt.isTrain:
self.criterionGAN = networks.GANLoss(opt.gan_mode, tensor=self.FloatTensor, opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
if opt.use_vae:
self.KLDLoss = networks.KLDLoss()
def __init__(self, opt):
"""Initialize the vox2vox class. Modified from models.pix2pix_model .
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 = ['G_GAN', 'G_L1L2', '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_3d.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_3d.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)
if opt.loss_norm == "L1":
self.criterionL1L2 = torch.nn.L1Loss()
elif opt.loss_norm == "L2":
self.criterionL1L2 = torch.nn.MSELoss()
else:
raise NotImplementedError(
"specify either an L1 or L2 as the loss_norm")
# 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 __init__(self, opt):
BaseModel.__init__(self, opt)
self.opt = opt
self.netrec = networks.define_rec(opt)
if self.opt.isTrain:
self.netD = networks.define_D(3, opt.ndf, 'basic', 3, opt.norm,
opt.init_type, opt.init_gain,
self.gpu_ids, opt.rec)
self.model_names = ['rec', 'D']
self.optimizer = torch.optim.Adam(self.netrec.parameters(),
lr=opt.lr)
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(0.5, 0.999))
self.criterionL1 = torch.nn.L1Loss()
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
else:
self.model_names = ['rec']
def __init__(self, opt):
"""
Initialize the Branch class
:param opt: 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 = ['G_GAN', 'G_L1','D_real','D_fake','D_origin']
self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake']
# specity the models you want to save/dispaly. The training/ test scripts will call <BaseModel.get_current_visuals>
self.visual_names = [
'real_A', 'real_B', 'fake_B', 'mask_B', 'mask_fake_B'
]
# specity the models you want to save to the disk. The training/ test scripts will call <BaseModel.save_network> and <BaseModel.load_network>
if self.isTrain:
self.model_names = ['G', 'D']
else: # during test time, only load G
self.model_names = ['G']
self.gpu_ids = opt.gpu_ids
# 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,
opt.gpu_ids)
if self.isTrain:
self.netD = networks.define_D(1536, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm,
opt.init_type, opt.init_gain,
self.gpu_ids)
# define loss function
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
self.criterionL1 = torch.nn.L1Loss()
# initialize optimizersl 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 __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor
self.perturbation = opt.perturbation
self.netG, self.netD, self.netE = self.initialize_networks(opt)
self.generator = opt.netG
self.noise_range = opt.noise_range
if self.perturbation:
self.netP = PerturbationNet(opt)
self.netP.cuda()
if self.opt.manipulation:
self.vocab = pickle.load(open(opt.vocab_path, 'rb'))
self.txt_enc = EncoderText(len(self.vocab), 300, 1024, 1)
self.txt_enc.load_state_dict(
torch.load(opt.vse_enc_path, map_location='cpu')['model'][1])
self.txt_enc.eval().cuda()
# set loss functions
if self.perturbation:
self.criterionPix = torch.nn.L1Loss()
self.criterionVGG = networks.VGGLoss(self.opt.gpu)
elif opt.isTrain:
self.loss_l1pix = opt.l1pix_loss
self.loss_gan = not opt.no_disc
self.loss_ganfeat = not opt.no_ganFeat_loss
self.loss_vgg = not opt.no_vgg_loss
if self.loss_l1pix:
self.criterionPix = torch.nn.L1Loss()
if self.loss_gan:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
if self.loss_ganfeat:
self.criterionFeat = torch.nn.L1Loss()
if self.loss_vgg:
self.criterionVGG = networks.VGGLoss(self.opt.gpu)
def __init__(self, opts, input_dim, output_dim):
super(Pix2Pix, self).__init__()
self.isTrain = (opts.phase == 'train')
self.gpu_ids = opts.gpu_ids
self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
# generator
self.netG = networks.define_G(input_dim, output_dim, 0, 64, netG='unet_256', norm='instance', upsample='bilinear', padding_type='reflect', gpu_ids=self.gpu_ids)
# discriminator
if self.isTrain:
self.netD = networks.define_D(input_dim + output_dim, 64, netD='basic_256_multi', norm='instance', num_Ds=2, gpu_ids=self.gpu_ids)
if self.isTrain:
# define loss functions, optimizers
self.criterionGAN = networks.GANLoss(gan_mode=opts.gan_mode).to(self.device)
self.criterionL1 = torch.nn.L1Loss()
self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opts.lr, betas=(0.5, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opts.lr, betas=(0.5, 0.999))
self.optimizers = [self.optimizer_G, self.optimizer_D]
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 __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.netG, self.netD, self.netE = self.initialize_networks(opt)
# set loss functions
if opt.isTrain:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
self.criterionFeat = torch.nn.L1Loss()
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
if opt.use_vae:
self.KLDLoss = networks.KLDLoss()
self.segmentation_model = \
torch.load('/home/qasima/venv_spade/SPADE/checkpoints/isic_fold_1/'
'model_epochs100_percent100_isic_256')
def __init__(self, opt):
super(AvModel, self).__init__()
self.opt = opt
self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.netG, self.netD, self.netA, self.netA_sync, self.netV, self.netE = \
self.initialize_networks(opt)
# set loss functions
if opt.isTrain:
self.loss_cls = CrossEntropyLoss()
self.criterionFeat = torch.nn.L1Loss()
if opt.softmax_contrastive:
self.criterionSoftmaxContrastive = networks.SoftmaxContrastiveLoss(
)
if opt.train_recognition or opt.train_sync:
pass
else:
self.criterionGAN = networks.GANLoss(opt.gan_mode,
tensor=self.FloatTensor,
opt=self.opt)
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt)
if opt.vgg_face:
self.VGGFace = VGGFace19(self.opt)
self.criterionVGGFace = networks.VGGLoss(
self.opt, self.VGGFace)
if opt.disentangle:
self.criterionLogSoftmax = networks.L2SoftmaxLoss()
def __init__(self, opt):
super().__init__()
self.opt = opt
self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \
else torch.FloatTensor
self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \
else torch.ByteTensor
self.netG_for_CT, self.netD_aligned, self.netG_for_MR, self.netD_unaligned = self.initialize_networks(opt)
# set loss functions
if opt.isTrain:
self.criterionGAN = networks.GANLoss(
opt.gan_mode, tensor=self.FloatTensor, opt=self.opt)
if not opt.no_vgg_loss:
self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
if not opt.no_L1_loss:
self.criterionL1 = torch.nn.L1Loss()
if not opt.no_SSIMLoss:
self.criterionSSIM = networks.SSIMLoss()
if not opt.no_GradientDifferenceLoss:
self.criterionGDL = networks.GradientDifferenceLoss()
if not opt.no_cycle_loss:
self.criterionCYC = torch.nn.L1Loss()
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)