def initialize(self, opt):
BaseModel.initialize(self, opt)
nb = opt.batchSize
size = opt.fineSize
self.input_P1_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_BP1_set = self.Tensor(nb, opt.BP_input_nc, size, size)
self.input_P2_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_BP2_set = self.Tensor(nb, opt.BP_input_nc, size, size)
self.mask_1 = None
self.mask_2 = None
self.fg_img = None
self.bg_img = None
self.bg_fill = None
self.out_patn_img = None
self.target_img_bg = None
self.target_img_fg = None
self.source_fg_mask = None
input_nc = [opt.P_input_nc, opt.BP_input_nc + opt.BP_input_nc]
self.netG = networks.define_G(input_nc,
opt.P_input_nc,
opt.ngf,
opt.which_model_netG,
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
n_downsampling=opt.G_n_downsampling)
if self.isTrain:
use_sigmoid = opt.no_lsgan
if opt.with_D_PB:
self.netD_PB = networks.define_D(
opt.P_input_nc + opt.BP_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if opt.with_D_PP:
self.netD_PP = networks.define_D(
opt.P_input_nc + opt.P_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if not self.isTrain or opt.continue_train:
which_epoch = opt.which_epoch
self.load_network(self.netG, 'netG', which_epoch)
if self.isTrain:
if opt.with_D_PB:
self.load_network(self.netD_PB, 'netD_PB', which_epoch)
if opt.with_D_PP:
self.load_network(self.netD_PP, 'netD_PP', which_epoch)
if self.isTrain:
self.old_lr = opt.lr
self.fake_PP_pool = ImagePool(opt.pool_size)
self.fake_PB_pool = ImagePool(opt.pool_size)
# define loss functions
# if opt.which_model_netD == 'patchgan':
# self.criterionGAN = networks.GANLoss(gan_mode='lsgan', target_real_label=1.0, target_fake_label=0.0)
# else:
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
self.loss_l1 = torch.nn.L1Loss()
if opt.L1_type == 'origin':
self.criterionL1 = torch.nn.L1Loss()
elif opt.L1_type == 'l1_plus_perL1':
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
else:
raise Exception('Unsurportted type of L1!')
# initialize optimizers
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PB:
self.optimizer_D_PB = torch.optim.Adam(
self.netD_PB.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PP:
self.optimizer_D_PP = torch.optim.Adam(
self.netD_PP.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
if opt.with_D_PB:
self.optimizers.append(self.optimizer_D_PB)
if opt.with_D_PP:
self.optimizers.append(self.optimizer_D_PP)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
if self.isTrain:
if opt.with_D_PB:
networks.print_network(self.netD_PB)
if opt.with_D_PP:
networks.print_network(self.netD_PP)
print('-----------------------------------------------')
def initialize(self, opt):
BaseModel.initialize(self, opt)
nb = opt.batchSize #6
size = opt.fineSize #256
self.input_P1_set = self.Tensor(nb, opt.P_input_nc, size,
size) #(6,3,256,256),图像
self.input_BP1_set = self.Tensor(nb, opt.BP_input_nc, size,
size) #(6,18,256,256)，骨骼点
self.input_P2_set = self.Tensor(nb, opt.P_input_nc, size,
size) ##(6,3,256,256)
self.input_BP2_set = self.Tensor(nb, opt.BP_input_nc, size,
size) #(6,18,256,256)
self.input_SP1_set = self.Tensor(nb, opt.SP_input_nc, size,
size) ##(6,8,256,256)#语义图
input_nc = [opt.P_input_nc,
opt.BP_input_nc + opt.BP_input_nc] #[3,18+18]
self.netG = networks.define_G(input_nc,
opt.P_input_nc,
opt.ngf,
opt.which_model_netG,
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
n_downsampling=opt.G_n_downsampling)
# define_G([3,36], 3,
# 64, opt.which_model_netG, instance, not opt.no_dropout, opt.init_type, self.gpu_ids,
# n_downsampling=2)
if self.isTrain:
use_sigmoid = opt.no_lsgan
if opt.with_D_PB:
self.netD_PB = networks.define_D(
opt.P_input_nc + opt.BP_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if opt.with_D_PP:
self.netD_PP = networks.define_D(
opt.P_input_nc + opt.P_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if len(opt.gpu_ids) > 1:
self.load_VGG(self.netG.module.enc_style.vgg)
else:
self.load_VGG(self.netG.enc_style.vgg)
if not self.isTrain or opt.continue_train:
which_epoch = opt.which_epoch
self.load_network(self.netG, 'netG', which_epoch)
if self.isTrain:
if opt.with_D_PB:
self.load_network(self.netD_PB, 'netD_PB', which_epoch)
if opt.with_D_PP:
self.load_network(self.netD_PP, 'netD_PP', which_epoch)
if self.isTrain:
self.old_lr = opt.lr
self.fake_PP_pool = ImagePool(opt.pool_size)
self.fake_PB_pool = ImagePool(opt.pool_size)
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
if opt.L1_type == 'origin':
self.criterionL1 = torch.nn.L1Loss()
elif opt.L1_type == 'l1_plus_perL1':
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
else:
raise Excption('Unsurportted type of L1!')
if opt.use_cxloss:
self.CX_loss = CXLoss(sigma=0.5)
if torch.cuda.is_available():
self.CX_loss.cuda()
self.vgg = VGG()
self.vgg.load_state_dict(
torch.load(
os.path.abspath(os.path.dirname(opt.dataroot)) +
'/vgg_conv.pth'))
for param in self.vgg.parameters():
param.requires_grad = False
if torch.cuda.is_available():
self.vgg.cuda()
# initialize optimizers
self.optimizer_G = torch.optim.Adam(filter(
lambda p: p.requires_grad, self.netG.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PB:
self.optimizer_D_PB = torch.optim.Adam(
self.netD_PB.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PP:
self.optimizer_D_PP = torch.optim.Adam(
self.netD_PP.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
if opt.with_D_PB:
self.optimizers.append(self.optimizer_D_PB)
if opt.with_D_PP:
self.optimizers.append(self.optimizer_D_PP)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
if self.isTrain:
if opt.with_D_PB:
networks.print_network(self.netD_PB)
if opt.with_D_PP:
networks.print_network(self.netD_PP)
print('-----------------------------------------------')
def initialize(self, opt):
BaseModel.initialize(self, opt)
input_nc = 18 + 1 + 18
self.netG = networks.define_G(opt,
input_nc,
1,
opt.ngf,
opt.which_model_netG,
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
n_downsampling=opt.G_n_downsampling)
if self.isTrain:
use_sigmoid = opt.no_lsgan
if opt.with_D_PB:
input_nc = 18 + 1
self.netD_PB = networks.define_D(
input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if opt.with_D_PP:
input_nc = 1 + 1
self.netD_PP = networks.define_D(
input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if not self.isTrain or opt.continue_train:
which_epoch = opt.which_epoch
self.load_network(self.netG, 'netG', which_epoch)
if self.isTrain:
if opt.with_D_PB:
self.load_network(self.netD_PB, 'netD_PB', which_epoch)
if opt.with_D_PP:
self.load_network(self.netD_PP, 'netD_PP', which_epoch)
if self.isTrain:
self.old_lr = opt.lr
self.fake_PP_pool = ImagePool(opt.pool_size)
self.fake_PB_pool = ImagePool(opt.pool_size)
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
if opt.L1_type == 'origin':
self.criterionL1 = torch.nn.L1Loss()
elif opt.L1_type == 'l1_plus_perL1':
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
else:
raise Exception('Unsurportted type of L1!')
if opt.use_cxloss:
self.CX_loss = CXLoss(sigma=0.5)
if torch.cuda.is_available():
self.CX_loss.cuda()
self.vgg = VGG()
self.vgg.load_state_dict(
torch.load('./fashion_data/vgg_conv.pth'))
for param in self.vgg.parameters():
param.requires_grad = False
if torch.cuda.is_available():
self.vgg.cuda()
# initialize optimizers
self.optimizer_G = torch.optim.Adam(filter(
lambda p: p.requires_grad, self.netG.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PB:
self.optimizer_D_PB = torch.optim.Adam(
self.netD_PB.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PP:
self.optimizer_D_PP = torch.optim.Adam(
self.netD_PP.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
if opt.with_D_PB:
self.optimizers.append(self.optimizer_D_PB)
if opt.with_D_PP:
self.optimizers.append(self.optimizer_D_PP)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
if self.isTrain:
if opt.with_D_PB:
networks.print_network(self.netD_PB)
if opt.with_D_PP:
networks.print_network(self.netD_PP)
print('-----------------------------------------------')
def initialize(self, opt):
BaseModel.initialize(self, opt)
nb = opt.batchSize
size = opt.fineSize
input_nc = [opt.P_input_nc, opt.BP_input_nc + opt.BP_input_nc]
self.netG1 = networks.define_G(input_nc,
opt.P_input_nc,
opt.ngf,
opt.which_model_netG,
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
n_downsampling=opt.G_n_downsampling)
self.netG2 = networks.define_G(input_nc,
opt.P_input_nc,
opt.ngf,
'SelectionGAN',
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
n_downsampling=opt.G_n_downsampling)
if self.isTrain:
use_sigmoid = opt.no_lsgan
if opt.with_D_PB:
self.netD_PB = networks.define_D(
opt.P_input_nc + opt.BP_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if opt.with_D_PP:
self.netD_PP = networks.define_D(
opt.P_input_nc + opt.P_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if not self.isTrain or opt.continue_train:
which_epoch = opt.which_epoch
self.load_network(self.netG1, 'netG1', which_epoch)
self.load_network(self.netG2, 'netG2', which_epoch)
if self.isTrain:
if opt.with_D_PB:
self.load_network(self.netD_PB, 'netD_PB', which_epoch)
if opt.with_D_PP:
self.load_network(self.netD_PP, 'netD_PP', which_epoch)
if self.isTrain:
self.old_lr = opt.lr
self.fake_PP_pool = ImagePool(opt.pool_size)
self.fake_PB_pool = ImagePool(opt.pool_size)
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
self.criterionSSIM = SSIM()
if opt.L1_type == 'origin':
self.criterionL1 = torch.nn.L1Loss()
elif opt.L1_type == 'l1_plus_perL1':
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
else:
raise Excption('Unsurportted type of L1!')
# initialize optimizers
self.optimizer_G = torch.optim.Adam(itertools.chain(
self.netG1.parameters(), self.netG2.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PB:
self.optimizer_D_PB = torch.optim.Adam(
self.netD_PB.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PP:
self.optimizer_D_PP = torch.optim.Adam(
self.netD_PP.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
if opt.with_D_PB:
self.optimizers.append(self.optimizer_D_PB)
if opt.with_D_PP:
self.optimizers.append(self.optimizer_D_PP)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
print('---------- Networks initialized -------------')
networks.print_network(self.netG1)
networks.print_network(self.netG2)
if self.isTrain:
if opt.with_D_PB:
networks.print_network(self.netD_PB)
if opt.with_D_PP:
networks.print_network(self.netD_PP)
print('-----------------------------------------------')
def __init__(self, opt):
super(MMHandModel, self).__init__(opt)
'''Added for Depth Map'''
nb = opt.batchSize
size = opt.fineSize
self.overflow = False
input_nc = [
opt.H_input_nc, opt.P_input_nc + opt.P_input_nc,
opt.D_input_nc + opt.D_input_nc
]
self.netG = self.define_G(input_nc,
opt.output_nc,
opt.ngf,
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
n_downsampling=opt.G_n_downsampling)
if self.isTrain:
self.netD_PB = self.define_D(opt.H_input_nc + opt.P_input_nc,
opt.ndf,
opt.n_layers_D,
opt.norm,
opt.no_lsgan,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
self.netD_PP = self.define_D(opt.H_input_nc + opt.H_input_nc,
opt.ndf,
opt.n_layers_D,
opt.norm,
opt.no_lsgan,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if not self.isTrain or opt.continue_train:
self.load_network()
if self.isTrain:
self.old_lr = opt.lr
self.fake_PP_pool = ImagePool(opt.pool_size)
self.fake_PB_pool = ImagePool(opt.pool_size)
# define loss functions
self.criterionGAN = GANLoss(use_lsgan=not opt.no_lsgan,
gpu=self.opt.local_rank)
if opt.L1_type == 'origin':
self.criterionL1 = torch.nn.L1Loss()
elif opt.L1_type == 'l1_plus_perL1':
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
else:
raise Exception('Unsurportted type of L1!')
# initialize optimizers
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizer_D_PB = torch.optim.Adam(self.netD_PB.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizer_D_PP = torch.optim.Adam(self.netD_PP.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if self.opt.distributed:
[self.netG, self.netD_PB, self.netD_PP], \
[self.optimizer_G, self.optimizer_D_PB, self.optimizer_D_PP] = \
amp.initialize(models=[self.netG,
self.netD_PB,
self.netD_PP],
optimizers=[self.optimizer_G,
self.optimizer_D_PB,
self.optimizer_D_PP],
num_losses=3, opt_level=opt.opt_level)
self.netG = apex.parallel.convert_syncbn_model(
DDP(self.netG, delay_allreduce=True))
self.netD_PB = apex.parallel.convert_syncbn_model(
DDP(self.netD_PB, delay_allreduce=True))
self.netD_PP = apex.parallel.convert_syncbn_model(
DDP(self.netD_PP, delay_allreduce=True))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D_PB)
self.optimizers.append(self.optimizer_D_PP)
for optimizer in self.optimizers:
self.schedulers.append(get_scheduler(optimizer, opt))
if self.master:
print('---------- Networks initialized -------------')
print_network(self.netG)
if self.isTrain:
print_network(self.netD_PB)
print_network(self.netD_PP)
print(opt.local_rank)
print('-----------------------------------------------')
def initialize(self, opt):
BaseModel.initialize(self, opt)
nb = opt.batchSize
size = opt.fineSize
self.input_P1_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_BP1_set = self.Tensor(nb, opt.BP_input_nc, size, size)
self.input_P2_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_BP2_set = self.Tensor(nb, opt.BP_input_nc, size, size)
self.input_BP2_mask_set = self.Tensor(nb, opt.BP_limbs_nc, size, size)
input_nc = [opt.P_input_nc, opt.BP_input_nc + opt.BP_input_nc]
self.netG = networks.define_G(input_nc,
opt.P_input_nc,
opt.ngf,
opt.which_model_netG,
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
n_blocks=opt.n_blocks,
n_downsampling=opt.G_n_downsampling)
if self.isTrain:
use_sigmoid = opt.no_lsgan
if opt.with_D_PB:
self.netD_PB = networks.define_D(
opt.P_input_nc + opt.BP_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if opt.with_D_PP:
self.netD_PP = networks.define_D(
opt.P_input_nc + opt.P_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if not self.isTrain or opt.continue_train:
which_epoch = opt.which_epoch
self.load_network(self.netG, 'netG', which_epoch)
if self.isTrain:
if opt.with_D_PB:
self.load_network(self.netD_PB, 'netD_PB', which_epoch)
if opt.with_D_PP:
self.load_network(self.netD_PP, 'netD_PP', which_epoch)
if self.isTrain:
self.old_lr = opt.lr
self.fake_PP_pool = ImagePool(opt.pool_size)
self.fake_PB_pool = ImagePool(opt.pool_size)
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
if opt.L1_type == 'origin': # L1 loss
self.criterionL1 = torch.nn.L1Loss()
elif opt.L1_type == 'perL1': # PerL1 loss
self.criterionL1 = PerceptualLoss(opt.lambda_B,
opt.perceptual_layers,
self.gpu_ids,
opt.percep_is_l1)
elif opt.L1_type == 'SSIM': # SSIM loss
self.criterionSSIM = SSIM(win_size=opt.win_size,
win_sigma=opt.win_sigma,
data_range=1.0,
size_average=True)
elif opt.L1_type == 'Style': # Style loss
self.criterionStyle = StyleLoss(opt.lambda_style,
opt.perceptual_layers,
self.gpu_ids)
elif opt.L1_type == 'PerSSIM': # PerceptualSSIM loss
self.criterionPerSSIM = PerceptualSSIMLoss(
opt.lambda_perssim,
opt.perceptual_layers,
self.gpu_ids,
win_size=opt.win_size,
win_sigma=opt.win_sigma)
elif opt.L1_type == 'l1_plus_perL1': # L1 + PerL1 loss
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
elif opt.L1_type == 'SSIM_plus_perL1_l1': # L1 + PerL1 + SSIM loss
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
self.criterionSSIM = SSIM(win_size=opt.win_size,
win_sigma=opt.win_sigma,
data_range=1.0,
size_average=True)
elif opt.L1_type == 'FPart_BSSIM_plus_perL1_L1': # FPart_BSSIM + PerL1 + L1 loss
self.criterionSSIM = FPart_BSSIM(data_range=1.0,
size_average=True,
win_size=opt.win_size,
win_sigma=opt.win_sigma)
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
elif opt.L1_type == 'FPart_BSSIM_plus_perL1_style': #FPart_BSSIM + PerL1 + style loss
self.criterionSSIM = FPart_BSSIM(data_range=1.0,
size_average=True,
win_size=opt.win_size,
win_sigma=opt.win_sigma)
self.criterionL1 = L1_plus_perceptual_styleLoss(
opt.lambda_style, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
else:
self.criterionL1 = None
# initialize optimizers
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PB:
self.optimizer_D_PB = torch.optim.Adam(
self.netD_PB.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PP:
self.optimizer_D_PP = torch.optim.Adam(
self.netD_PP.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
if opt.with_D_PB:
self.optimizers.append(self.optimizer_D_PB)
if opt.with_D_PP:
self.optimizers.append(self.optimizer_D_PP)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
def initialize(self, opt):
self.opt = opt
self.gpu_ids = opt.gpu_ids
self.isTrain = opt.isTrain
self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
nb = opt.batchSize
size = opt.fineSize
self.input_P1_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_KP1_set = self.Tensor(nb, opt.BP_input_nc, size, size)
self.input_P2_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_KP2_set = self.Tensor(nb, opt.BP_input_nc, size, size)
self.input_SPL1_set = self.Tensor(nb, 1, size, size)
self.input_SPL2_set = self.Tensor(nb, 1, size, size)
self.input_SPL1_onehot_set = self.Tensor(nb, 12, size, size)
self.input_SPL2_onehot_set = self.Tensor(nb, 12, size, size)
self.input_syn_set = self.Tensor(nb, opt.P_input_nc, size, size)
input_nc = [
opt.P_input_nc, opt.BP_input_nc + opt.BP_input_nc, opt.P_input_nc
]
self.netG = networks.define_G(input_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
if opt.with_D_PB:
self.netD_PB = networks.define_D(
3 + 18,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
'instance',
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if opt.with_D_PP:
self.netD_PP = networks.define_D(
opt.P_input_nc + opt.P_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
'instance',
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if not self.isTrain or opt.continue_train:
which_epoch = opt.which_epoch
self.load_network(self.netG, 'netG', which_epoch)
if self.isTrain:
if opt.with_D_PB:
self.load_network(self.netD_PB, 'netD_PB', which_epoch)
if opt.with_D_PP:
self.load_network(self.netD_PP, 'netD_PP', which_epoch)
if self.isTrain:
self.old_lr = opt.lr
self.fake_PP_pool = ImagePool(opt.pool_size)
self.fake_PB_pool = ImagePool(opt.pool_size)
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
#define shape loss
if False: #self._opt.mask_bce:
self.parseLoss = torch.nn.BCELoss()
else:
self.parseLoss = CrossEntropyLoss2d()
if opt.L1_type == 'origin':
self.criterionL1 = torch.nn.L1Loss()
elif opt.L1_type == 'l1_plus_perL1':
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
else:
raise Excption('Unsurportted type of L1!')
# initialize optimizers
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PB:
self.optimizer_D_PB = torch.optim.Adam(
self.netD_PB.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PP:
self.optimizer_D_PP = torch.optim.Adam(
self.netD_PP.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
if opt.with_D_PB:
self.optimizers.append(self.optimizer_D_PB)
if opt.with_D_PP:
self.optimizers.append(self.optimizer_D_PP)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
if self.isTrain:
if opt.with_D_PB:
networks.print_network(self.netD_PB)
if opt.with_D_PP:
networks.print_network(self.netD_PP)
print('-----------------------------------------------')
def initialize(self, opt):
BaseModel.initialize(self, opt)
nb = opt.batchSize
size = opt.fineSize
self.opt = opt
self.mask = torch.ones([nb, 1, 46, 32]).cuda()
cpm_model_path = 'openpose_coco_best.pth.tar'
self.cpm_model = heatmap_pose.construct_model(cpm_model_path)
# self.cpm_model = cpm_test.construct_model(cpm_model_path)
input_nc = [opt.P_input_nc, opt.BP_input_nc + opt.BP_input_nc]
self.netG = networks.define_G(input_nc,
opt.P_input_nc,
opt.ngf,
opt.which_model_netG,
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
n_downsampling=opt.G_n_downsampling)
if self.isTrain:
use_sigmoid = opt.no_lsgan
if opt.with_D_PB:
self.netD_PB = networks.define_D(
opt.P_input_nc + opt.BP_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if opt.with_D_PP:
self.netD_PP = networks.define_D(
opt.P_input_nc + opt.P_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if not self.isTrain or opt.continue_train:
which_epoch = opt.which_epoch
self.load_network(self.netG, 'netG', which_epoch)
if self.isTrain:
if opt.with_D_PB:
self.load_network(self.netD_PB, 'netD_PB', which_epoch)
if opt.with_D_PP:
self.load_network(self.netD_PP, 'netD_PP', which_epoch)
if self.isTrain:
self.old_lr = opt.lr
self.fake_PP_pool = ImagePool(opt.pool_size)
self.fake_PB_pool = ImagePool(opt.pool_size)
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
if opt.pose_loss:
self.pose_loss = torch.nn.MSELoss()
if opt.L1_type == 'origin':
self.criterionL1 = torch.nn.L1Loss()
elif opt.L1_type == 'l1_plus_perL1':
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
else:
raise Excption('Unsurportted type of L1!')
# initialize optimizers
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PB:
self.optimizer_D_PB = torch.optim.Adam(
self.netD_PB.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PP:
self.optimizer_D_PP = torch.optim.Adam(
self.netD_PP.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
if opt.with_D_PB:
self.optimizers.append(self.optimizer_D_PB)
if opt.with_D_PP:
self.optimizers.append(self.optimizer_D_PP)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
if self.isTrain:
if opt.with_D_PB:
networks.print_network(self.netD_PB)
if opt.with_D_PP:
networks.print_network(self.netD_PP)
print('-----------------------------------------------')
self.reset_acc()
def initialize(self, opt):
BaseModel.initialize(self, opt)
nb = opt.batchSize
size = opt.fineSize
self.input_P1_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_BP1_set = self.Tensor(nb, opt.BP_input_nc, size, size)
self.input_P2_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_BP2_set = self.Tensor(nb, opt.BP_input_nc, size, size)
self.input_P2_64_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_P2_32_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_P2_16_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.posenum_set = self.Tensor(nb, size, 1)
self.posenum_para_set = self.Tensor(nb, size, 1)
self.idnum_set = self.Tensor(nb, size, 1)
self.input_P1_64_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.input_P1_32_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.left_eye_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.right_eye_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.nose_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.mouth_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.target_left_eye_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.target_right_eye_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.target_nose_set = self.Tensor(nb, opt.P_input_nc, size, size)
self.target_mouth_set = self.Tensor(nb, opt.P_input_nc, size, size)
input_nc = [opt.P_input_nc, opt.BP_input_nc + opt.BP_input_nc]
self.netG = networks.define_G(input_nc,
opt.P_input_nc,
opt.ngf,
opt.which_model_netG,
opt.norm,
not opt.no_dropout,
opt.init_type,
self.gpu_ids,
n_downsampling=opt.G_n_downsampling)
if self.isTrain:
use_sigmoid = opt.no_lsgan
if opt.with_D_PB:
self.netD_PB = networks.define_D(
opt.P_input_nc + opt.BP_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if opt.with_D_PP:
self.netD_PP = networks.define_D(
opt.P_input_nc + opt.P_input_nc,
opt.ndf,
opt.which_model_netD,
opt.n_layers_D,
opt.norm,
use_sigmoid,
opt.init_type,
self.gpu_ids,
not opt.no_dropout_D,
n_downsampling=opt.D_n_downsampling)
if not self.isTrain or opt.continue_train == 1:
which_epoch = opt.which_epoch
self.load_network(self.netG, 'netG', which_epoch)
if self.isTrain:
if opt.with_D_PB:
self.load_network(self.netD_PB, 'netD_PB', which_epoch)
if opt.with_D_PP:
self.load_network(self.netD_PP, 'netD_PP', which_epoch)
if self.isTrain:
self.old_lr = opt.lr
self.fake_PP_pool = ImagePool(opt.pool_size)
self.fake_PB_pool = ImagePool(opt.pool_size)
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
if opt.L1_type == 'origin':
self.criterionL1 = torch.nn.L1Loss()
elif opt.L1_type == 'l1_plus_perL1':
self.criterionL1 = L1_plus_perceptualLoss(
opt.lambda_A, opt.lambda_B, opt.perceptual_layers,
self.gpu_ids, opt.percep_is_l1)
else:
raise Excption('Unsurportted type of L1!')
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=(opt.lr / 2),
betas=(opt.beta1, 0.999))
if opt.with_D_PB:
self.optimizer_D_PB = torch.optim.Adam(
self.netD_PB.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
if opt.with_D_PP:
self.optimizer_D_PP = torch.optim.Adam(
self.netD_PP.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers = []
self.schedulers = []
self.optimizers.append(self.optimizer_G)
if opt.with_D_PB:
self.optimizers.append(self.optimizer_D_PB)
if opt.with_D_PP:
self.optimizers.append(self.optimizer_D_PP)
for optimizer in self.optimizers:
self.schedulers.append(networks.get_scheduler(optimizer, opt))