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))