def print_network(self): print('--------------------- Model ---------------------') print('##################### NetG #####################') networks.print_network(self.net_i) if self.isTrain and self.opt.lambda_gan > 0: print('##################### NetD #####################') networks.print_network(self.netD)
def print_network(self): if self.opt.debug: print('--------------------- Model ---------------------') if self.vgg is not None: print('##################### VGG #####################') from torchsummary import summary summary(self.vgg, input_size=( 3, 224, 224)) #torch vgg model summary for 224*224 i/p size print('##################### NetG #####################') networks.print_network(self.net_i) if self.isTrain and self.opt.lambda_gan > 0: print('##################### NetD #####################') networks.print_network(self.netD)
def __init__(self, opt): super(Pose_GAN, self).__init__() # load generator and discriminator models # adding extra layers for larger image size if(opt.checkMode == 0): nfilters_decoder = (512, 512, 512, 256, 128, 3) if max(opt.image_size) < 256 else (512, 512, 512, 512, 256, 128, 3) nfilters_encoder = (64, 128, 256, 512, 512, 512) if max(opt.image_size) < 256 else (64, 128, 256, 512, 512, 512, 512) else: nfilters_decoder = (128, 3) if max(opt.image_size) < 256 else (256, 128, 3) nfilters_encoder = (64, 128) if max(opt.image_size) < 256 else (64, 128, 256) if (opt.use_input_pose): input_nc = 3 + 2*opt.pose_dim else: input_nc = 3 + opt.pose_dim self.num_stacks = opt.num_stacks self.batch_size = opt.batch_size self.pose_dim = opt.pose_dim if(opt.gen_type=='stacked'): self.gen = Stacked_Generator(input_nc, opt.num_stacks, opt.pose_dim, nfilters_encoder, nfilters_decoder, use_input_pose=opt.use_input_pose) elif(opt.gen_type=='baseline'): self.gen = Generator(input_nc, nfilters_encoder, nfilters_decoder, use_input_pose=opt.use_input_pose) else: raise Exception('Invalid gen_type') # discriminator also sees the output image for the target pose self.disc = Discriminator(input_nc + 3, use_input_pose=opt.use_input_pose, checkMode=opt.checkMode) print('---------- Networks initialized -------------') print_network(self.gen) print_network(self.disc) print('-----------------------------------------------') # Setup the optimizers lr = opt.learning_rate self.disc_opt = torch.optim.Adam(self.disc.parameters(), lr=lr, betas=(0.5, 0.999)) self.gen_opt = torch.optim.Adam(self.gen.parameters(), lr=lr, betas=(0.5, 0.999)) # Network weight initialization self.gen.cuda() self.disc.cuda() self.disc.apply(xavier_weights_init) self.gen.apply(xavier_weights_init) # Setup the loss function for training self.ll_loss_criterion = torch.nn.L1Loss()
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, _isTrain=False): self.initialize(opt) self.mode = opt.mode if opt.input == 'single_view': self.num_input = 3 elif opt.input == 'two_view': self.num_input = 6 elif opt.input == 'two_view_k': self.num_input = 7 else: raise ValueError("Unknown input type %s" % opt.input) if self.mode == 'Ours_Bilinear': print( '====================================== DIW NETWORK TRAIN FROM %s=======================' % self.mode) new_model = hourglass.HourglassModel(self.num_input) print( '===================Loading Pretrained Model OURS ===================================' ) if not _isTrain: if self.num_input == 7: model_parameters = self.load_network( new_model, 'G', 'best_depth_Ours_Bilinear_inc_7') elif self.num_input == 3: model_parameters = self.load_network( new_model, 'G', 'best_depth_Ours_Bilinear_inc_3') elif self.num_input == 6: model_parameters = self.load_network( new_model, 'G', 'best_depth_Ours_Bilinear_inc_6') else: print('Something Wrong') sys.exit() new_model.load_state_dict(model_parameters) new_model = torch.nn.parallel.DataParallel( new_model.cuda(), device_ids=range(torch.cuda.device_count())) self.netG = new_model else: print('ONLY SUPPORT Ours_Bilinear') sys.exit() self.old_lr = opt.lr self.netG.train() if True: self.criterion_joint = networks.JointLoss(opt) # initialize optimizers self.optimizer_G = torch.optim.Adam( self.netG.parameters(), lr=opt.lr, betas=(0.9, 0.999)) self.scheduler = networks.get_scheduler(self.optimizer_G, opt) print('---------- Networks initialized -------------') networks.print_network(self.netG) print('-----------------------------------------------')
def initialize(self, opt): BaseModel.initialize(self, opt) nb = opt.batchSize size = opt.fineSize self.target_weight = [] self.input_A = self.Tensor(nb, opt.input_nc, size, size) self.input_B = self.Tensor(nb, opt.output_nc, size, size) self.input_C = self.Tensor(nb, opt.output_nc, size, size) self.input_C_sr = self.Tensor(nb, opt.output_nc, size, size) if opt.aux: self.A_aux = self.Tensor(nb, opt.input_nc, size, size) self.B_aux = self.Tensor(nb, opt.output_nc, size, size) self.C_aux = self.Tensor(nb, opt.output_nc, size, size) self.netE_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'ResnetEncoder_my', opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt, n_downsampling=2) mult = self.netE_A.get_mult() self.netE_C = networks.define_G(opt.input_nc, opt.output_nc, 64, 'ResnetEncoder_my', opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt, n_downsampling=3) self.net_D = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'ResnetDecoder_my', opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt, mult=mult) mult = self.net_D.get_mult() self.net_Dc = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'ResnetDecoder_my', opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt, mult=mult, n_upsampling=1) self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'GeneratorLL', opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt, mult=mult) mult = self.net_Dc.get_mult() self.netG_C = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, 'GeneratorLL', opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt, mult=mult) # self.netG_A_running = networks.define_G(opt.input_nc, opt.output_nc, # opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt) # set_eval(self.netG_A_running) # accumulate(self.netG_A_running, self.netG_A, 0) # self.netG_B = networks.define_G(opt.output_nc, opt.input_nc, # opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt) # self.netG_B_running = networks.define_G(opt.output_nc, opt.input_nc, # opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, self.gpu_ids, opt=opt) # set_eval(self.netG_B_running) # accumulate(self.netG_B_running, self.netG_B, 0) if self.isTrain: use_sigmoid = opt.no_lsgan self.netD_A = networks.define_D(opt.output_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, self.gpu_ids, opt=opt) # self.netD_B = networks.define_D(opt.input_nc, opt.ndf, # opt.which_model_netD, # opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, self.gpu_ids, opt=opt) print('---------- Networks initialized -------------') # networks.print_network(self.netG_B, opt, (opt.input_nc, opt.fineSize, opt.fineSize)) networks.print_network(self.netE_C, opt, (opt.input_nc, opt.fineSize, opt.fineSize)) networks.print_network( self.net_D, opt, (opt.ngf * 4, opt.fineSize / 4, opt.fineSize / 4)) networks.print_network(self.net_Dc, opt, (opt.ngf, opt.CfineSize / 2, opt.CfineSize / 2)) # networks.print_network(self.netG_B, opt) if self.isTrain: networks.print_network(self.netD_A, opt) # networks.print_network(self.netD_B, opt) print('-----------------------------------------------') if not self.isTrain or opt.continue_train: print('Loaded model') which_epoch = opt.which_epoch self.load_network(self.netG_A, 'G_A', which_epoch) self.load_network(self.netG_B, 'G_B', which_epoch) if self.isTrain: self.load_network(self.netG_A_running, 'G_A', which_epoch) self.load_network(self.netG_B_running, 'G_B', which_epoch) self.load_network(self.netD_A, 'D_A', which_epoch) self.load_network(self.netD_B, 'D_B', which_epoch) if self.isTrain and opt.load_path != '': print('Loaded model from load_path') which_epoch = opt.which_epoch load_network_with_path(self.netG_A, 'G_A', opt.load_path, epoch_label=which_epoch) load_network_with_path(self.netG_B, 'G_B', opt.load_path, epoch_label=which_epoch) load_network_with_path(self.netD_A, 'D_A', opt.load_path, epoch_label=which_epoch) load_network_with_path(self.netD_B, 'D_B', opt.load_path, epoch_label=which_epoch) if self.isTrain: self.old_lr = opt.lr self.fake_A_pool = ImagePool(opt.pool_size) self.fake_B_pool = ImagePool(opt.pool_size) self.fake_C_pool = ImagePool(opt.pool_size) # define loss functions if len(self.target_weight) == opt.num_D: print(self.target_weight) self.criterionGAN = networks.GANLoss( use_lsgan=not opt.no_lsgan, tensor=self.Tensor, target_weight=self.target_weight, gan=opt.gan) else: self.criterionGAN = networks.GANLoss( use_lsgan=not opt.no_lsgan, tensor=self.Tensor, gan=opt.gan) self.criterionCycle = torch.nn.L1Loss() self.criterionIdt = torch.nn.L1Loss() self.criterionColor = networks.ColorLoss() # initialize optimizers self.optimizer_G = torch.optim.Adam(itertools.chain( self.netE_A.parameters(), self.net_D.parameters(), self.netG_A.parameters(), self.net_Dc.parameters(), self.netG_C.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_AE = torch.optim.Adam(itertools.chain( self.netE_C.parameters(), self.net_D.parameters(), self.net_Dc.parameters(), self.netG_C.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_G_A_sr = torch.optim.Adam(itertools.chain( self.netE_A.parameters(), self.net_D.parameters(), self.net_Dc.parameters(), self.netG_C.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_AE_sr = torch.optim.Adam(itertools.chain( self.netE_C.parameters(), self.net_D.parameters(), self.netG_A.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_D_A = torch.optim.Adam(self.netD_A.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) # self.optimizer_D_B = torch.optim.Adam(self.netD_B.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers = [] self.schedulers = [] self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_AE) # self.optimizers.append(self.optimizer_G_A_sr) self.optimizers.append(self.optimizer_AE_sr) self.optimizers.append(self.optimizer_D_A) # self.optimizers.append(self.optimizer_D_B) for optimizer in self.optimizers: self.schedulers.append(networks.get_scheduler(optimizer, opt))
save_path = os.path.join(opt.checkpoints_dir, opt.name, save_filename) torch.save(network.cpu().state_dict(), save_path) if torch.cuda.is_available(): network.cuda() netG_deblur = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.gpu_ids, False, opt.learn_residual) netG_blur = multi_in_networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.gpu_ids, False, opt.learn_residual) load_network(netG_deblur, 'deblur_G', opt.which_epoch) load_network(netG_blur, 'blur_G', opt.which_epoch) print('------- Networks deblur_G initialized ---------') networks.print_network(netG_deblur) print('-----------------------------------------------') # ### Freeze layers # In[6]: def freeze_single_input(model,num_layers_frozen=19): ct=0 for child in list(model.children())[0]: ct+=1 if ct<num_layers_frozen: for param in child.parameters():
def print_network(self): print('--------------------- NetworkWrapper ---------------------') networks.print_network(self.net)
# visualizer = Visualizer(opt) # if opt.distributed: # torch.cuda.set_device(args.local_rank) # torch.distributed.init_process_group(backend="nccl", init_method="env://") # synchronize() optimizer_G, optimizer_D = model.module.optimizer_G, model.module.optimizer_D total_steps = (start_epoch - 1) * dataset_size + epoch_iter display_delta = total_steps % opt.display_freq print_delta = total_steps % opt.print_freq save_delta = total_steps % opt.save_latest_freq netVAE = networks.define_VAE(opt.input_nc) networks.print_network(netVAE) for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() if epoch != start_epoch: epoch_iter = epoch_iter % dataset_size for i, data in enumerate(dataset, start=epoch_iter): if total_steps % opt.print_freq == print_delta: iter_start_time = time.time() total_steps += opt.batchSize epoch_iter += opt.batchSize # whether to collect output images save_fake = total_steps % opt.display_freq == display_delta # print("DATA.shape: ", data)
def print_network(self): print('--------------------- Model ---------------------') networks.print_network(self.netG)
def initialize(self): # , opt BaseModel.initialize(self) # , opt batchSize = 32 fineSize = 256 input_nc = 3 output_nc = 3 vgg = 0 skip = 0.8 ngf = 64 pool_size = 50 norm = 'instance' lr = 0.0001 no_dropout = True no_lsgan = True continue_train = True use_wgan = 0.0 use_mse = True beta1 = 0.5 global which_direction new_lr = True niter_decay = 100 l1 = 10.0 # batch size nb = batchSize # 图像size size = fineSize #self.opt = opt self.input_A = self.Tensor(nb, input_nc, size, size) self.input_B = self.Tensor(nb, output_nc, size, size) self.input_img = self.Tensor(nb, input_nc, size, size) self.input_A_gray = self.Tensor(nb, 1, size, size) # Default 0, use perceptrual loss if vgg > 0: self.vgg_loss = networks.PerceptualLoss() self.vgg_loss.cuda() self.vgg = networks.load_vgg16("./model") self.vgg.eval() for param in self.vgg.parameters(): param.requires_grad = False # load/define networks # The naming conversion is different from those used in the paper # Code (paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X) # default=0.8, help='B = net.forward(A) + skip*A' skip = True if skip > 0 else False # which_model_netG, default = 'unet-256', selects model to use for netG # ngf, default = 64, of gen filters in first conv layer' # norm, default = 'instance', instance normalization or batch normalization # no_dropout, default = 'True', no dropout for the generator self.netG_A = networks.define_G(input_nc, output_nc, ngf, which_model_netG, norm, not no_dropout, self.gpu_ids, skip=skip) if not self.isTrain or continue_train: #which epoch to load which_epoch = 'lastest' self.load_network(self.netG_A, 'G_A', which_epoch) # --pool_size', default=50, help='the size of image buffer that stores previously generated images' # lr, default=0.0001 if self.isTrain: self.old_lr = lr self.fake_A_pool = ImagePool(pool_size) self.fake_B_pool = ImagePool(pool_size) # define loss functions if use_wgan: self.criterionGAN = networks.DiscLossWGANGP() else: # no_lsgan = True self.criterionGAN = networks.GANLoss(use_lsgan=not no_lsgan, tensor=self.Tensor) if use_mse: self.criterionCycle = torch.nn.MSELoss() else: self.criterionCycle = torch.nn.L1Loss() self.criterionL1 = torch.nn.L1Loss() self.criterionIdt = torch.nn.L1Loss() # initialize optimizers self.optimizer_G = torch.optim.Adam(self.netG_A.parameters(), lr=lr, betas=(beta1, 0.999)) print('---------- Networks initialized -------------') networks.print_network(self.netG_A) if isTrain: self.netG_A.train() else: self.netG_A.eval() print('-----------------------------------------------')