def TrainDiscriminator(real_A: tp.Numpy.Placeholder( (1, 3, args.crop_size, args.crop_size), dtype=flow.float32), fake_A: tp.Numpy.Placeholder( (1, 3, args.crop_size, args.crop_size), dtype=flow.float32), real_B: tp.Numpy.Placeholder( (1, 3, args.crop_size, args.crop_size), dtype=flow.float32), fake_B: tp.Numpy.Placeholder( (1, 3, args.crop_size, args.crop_size), dtype=flow.float32)): with flow.scope.placement("gpu", "0:0-0"): # Calculate GAN loss for discriminator D_A # Real pred_real_B = networks.define_D(real_B, "netD_A", ndf=args.ndf, n_layers_D=3, trainable=True, reuse=True) loss_D_A_real = networks.GANLoss(pred_real_B, True) # Fake pred_fake_B = networks.define_D(fake_B, "netD_A", ndf=args.ndf, n_layers_D=3, trainable=True, reuse=True) loss_D_A_fake = networks.GANLoss(pred_fake_B, False) # Combined loss and calculate gradients loss_D_A = (loss_D_A_real + loss_D_A_fake) * 0.5 # Calculate GAN loss for discriminator D_B # Real pred_real_A = networks.define_D(real_A, "netD_B", ndf=args.ndf, n_layers_D=3, trainable=True, reuse=True) loss_D_B_real = networks.GANLoss(pred_real_A, True) # Fake pred_fake_A = networks.define_D(fake_A, "netD_B", ndf=args.ndf, n_layers_D=3, trainable=True, reuse=True) loss_D_B_fake = networks.GANLoss(pred_fake_A, False) # Combined loss and calculate gradients loss_D_B = (loss_D_B_real + loss_D_B_fake) * 0.5 loss_D = loss_D_A + loss_D_B flow.optimizer.Adam(flow.optimizer.PiecewiseConstantScheduler( [], [args.learning_rate]), beta1=0.5).minimize(loss_D) return loss_D
def __init__(self, opt): super().__init__() self.is_train = opt.is_train self.gpu_ids = opt.gpu_ids self.save_dir = opt.checkpoints_dir self.device = torch.device('cuda:{}'.format( self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') # load/define networks #self.netG = networks.define_X(1, 3, opt.init_type, opt.init_gain, self.gpu_ids) self.netG = networks.define_G(3, 2, opt.ngf_S12, opt.netG, opt.norm_G_D, not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids) if self.is_train: # define loss functions self.criterionGAN = networks.GANLoss( use_lsgan=not opt.no_lsgan).to(self.device) self.criterionL1 = torch.nn.L1Loss() # initialize optimizers self.optimizers = [] self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers.append(self.optimizer_G)
def __init__(self, input_nc=3, output_nc=3, gpu_id=None): self.device = torch.device( f"cuda:{gpu_id}" if gpu_id is not None else 'cpu') print(f"Using device {self.device}") # Hyperparameters self.lambda_idt = 0.5 self.lambda_A = 10.0 self.lambda_B = 10.0 # Define generator networks self.netG_A = networks.define_netG(input_nc, output_nc, ngf=64, n_blocks=9, device=self.device) self.netG_B = networks.define_netG(output_nc, input_nc, ngf=64, n_blocks=9, device=self.device) # Define discriminator networks self.netD_A = networks.define_netD(output_nc, ndf=64, n_layers=3, device=self.device) self.netD_B = networks.define_netD(input_nc, ndf=64, n_layers=3, device=self.device) # Define image pools self.fake_A_pool = utils.ImagePool(pool_size=50) self.fake_B_pool = utils.ImagePool(pool_size=50) # Define loss functions self.criterionGAN = networks.GANLoss().to(self.device) self.criterionCycle = torch.nn.L1Loss() self.criterionIdt = torch.nn.L1Loss() # Define optimizers netG_params = itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()) netD_params = itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()) self.optimizer_G = torch.optim.Adam(netG_params, lr=0.0002, betas=(0.5, 0.999)) self.optimizer_D = torch.optim.Adam(netD_params, lr=0.0002, betas=(0.5, 0.999)) # Learning rate schedulers self.scheduler_G = utils.get_lr_scheduler(self.optimizer_G) self.scheduler_D = utils.get_lr_scheduler(self.optimizer_D)
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) # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X) # input file shape (batch, length, dims-256) 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 = AudioPool(opt.pool_size) # create image buffer to store previously generated images self.fake_B_pool = AudioPool(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 = nn.L1Loss() self.criterionIdt = 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, opt): """ Parameters: """ super(CycleGANModel, self).__init__() # Stuff ensure_existance_paths(opt) # Create necessary directories for saving stuff self.opt = opt self.gpu_ids = opt.gpu_ids self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids and torch.cuda.is_available() else torch.device('cpu') self.lambda_idt = 0.5 self.loss_lambda = 10.0 self.visual_names = ['real_A', 'fake_B', 'rec_A', 'idt_B', 'real_B', 'fake_A', 'rec_B', 'idt_A'] self.name = self.opt.name # Define models self.model_names = ['G_A', 'G_B', 'D_A', 'D_B'] self.netG_A = networks.to_device(networks.ResNetGenerator(opt.n_input, opt.n_output, opt.forward_mask),self.gpu_ids) self.netG_B = networks.to_device(networks.ResNetGenerator(opt.n_input, opt.n_output, opt.forward_mask),self.gpu_ids) self.netD_A = networks.to_device(networks.NLayerDiscriminator(opt.n_output),self.gpu_ids) self.netD_B = networks.to_device(networks.NLayerDiscriminator(opt.n_output),self.gpu_ids) # Define losses self.criterionGAN = networks.GANLoss().to(self.device) self.criterionCycle = nn.L1Loss().to(self.device) self.criterionIdt = nn.L1Loss().to(self.device) self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B'] # Define optimizers self.optimizers = [] 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.optimizers.append(self.optimizer_G) 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_D) # Define schedulers self.schedulers = [networks.get_optimizer_scheduler(optimizer,opt) for optimizer in self.optimizers] self.set_train() self.step = 0 wandb.init( entity="star-witchers", project="cycleGAN", config=self.opt, name=self.name, )
def initialize(self, opt, writer=None): BaseModel.initialize(self, opt) self.writer = writer self.num_step = 0 self.opt = opt if self.opt.use_lbp_network: self.model_names = ['G', 'LBP', 'D', 'D2'] else: self.model_names = ['G', 'D'] self.netG = networks.define_G(self.opt) if self.opt.use_lbp_network: self.netLBP = networks.define_LBP(self.opt) self.netD = networks.define_D( opt.input_nc, opt.ndf, self.opt.device) # Discriminator for netG if self.opt.use_lbp_network: self.netD2 = networks.define_D( opt.input_nc - 2, opt.ndf, self.opt.device) # Discriminator for netLBP self.vgg16_extractor = util.VGG16FeatureExtractor().to(self.opt.device) self.criterionGAN = networks.GANLoss(gan_type=opt.gan_type).to( self.opt.device) self.criterionL1 = torch.nn.L1Loss() self.criterionL2 = torch.nn.MSELoss() self.criterionL1_mask = networks.Discounted_L1(opt).to(self.opt.device) self.criterionL2_style_loss = torch.nn.MSELoss() self.criterionL2_perceptual_loss = torch.nn.MSELoss() self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(0.5, 0.999)) if self.opt.use_lbp_network: self.optimizer_LBP = torch.optim.Adam(self.netLBP.parameters(), lr=opt.lr, betas=(0.5, 0.999)) self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(0.5, 0.999)) if self.opt.use_lbp_network: self.optimizer_D2 = torch.optim.Adam(self.netD2.parameters(), lr=opt.lr, betas=(0.5, 0.999)) _, self.rand_t, self.rand_l = util.create_rand_mask(self.opt)
def initialize(self, opt): self.opt = opt self.init_architecture(opt) if opt.use_gpu: self.Tensor = torch.cuda.FloatTensor else: self.Tensor = torch.FloatTensor self.critGAN = networks.GANLoss(mse_loss=True, tensor=self.Tensor) self.critL1 = torch.nn.L1Loss() if self.opt.isTrain: self.schedulers = [] for optimizer in self.optimizers: self.schedulers.append(get_scheduler(optimizer, opt))
def initialize(self, opt, train_mode=True): # Model transforms from A --> B and uses Adv as the # adversarial example. BaseModel.initialize(self, opt) self.train_mode = train_mode # define tensors self.input_B = self.Tensor(opt['batchSize'], opt['input_nc'], opt['B_height'], opt['B_width']) self.input_A = self.Tensor(opt['batchSize'], opt['output_nc'], opt['A_height'], opt['A_width']) # load/define networks self.netG = networks.define_G(opt['input_nc'], opt['output_nc'], opt['ngf'], opt['norm'], self.gpu_ids) if self.train_mode: 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'], use_sigmoid, self.gpu_ids) if self.train_mode: # self.fake_AB_pool = ImagePool(opt['pool_size']) self.old_lr = opt['lr'] # define loss functions self.criterionGAN = networks.GANLoss(use_lsgan=not opt['no_lsgan'], tensor=self.Tensor) self.content_loss = torch.nn.MSELoss() # initialize 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)) print('---------- Networks initialized -------------') networks.print_network(self.netG) networks.print_network(self.netD) print('-----------------------------------------------')
def __init__(self, opt): self.opt = opt self.device = "cuda" self.netG_A = networks.ResnetGenerator(n_blocks=opt.n_blocks).to( self.device) self.netG_B = networks.ResnetGenerator(n_blocks=opt.n_blocks).to( self.device) self.netD_A = networks.NLayerDiscriminator().to(self.device) self.netD_B = networks.NLayerDiscriminator().to(self.device) self.fake_A_pool = ImagePool(opt.pool_size) self.fake_B_pool = ImagePool(opt.pool_size) self.criterionGAN = networks.GANLoss("lsgan").to(self.device) self.criterionCycle = flow.nn.L1Loss() self.criterionIdt = flow.nn.L1Loss() self.optimizer_G = flow.optim.Adam( itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, opt.beta2), ) self.optimizer_D = flow.optim.Adam( itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, opt.beta2), ) self.optimizers = [self.optimizer_G, self.optimizer_D] self.schedulers = [ flow.optim.lr_scheduler.CosineDecayLR(optimizer, decay_steps=100, alpha=0.0) for optimizer in self.optimizers ]
def initialize(self, opt): super(DesignerGAN, self).initialize(opt) ################################### # define data tensors ################################### # self.input['img'] = self.Tensor() # self.input['img_attr'] = self.Tensor() # self.input['lm_map'] = self.Tensor() # self.input['seg_mask'] = self.Tensor() # self.input['attr_label'] = self.Tensor() # self.input['id'] = [] ################################### # load/define networks ################################### # Todo modify networks.define_G # 1. add specified generator networks self.netG = networks.define_G(opt) self.netAE, self.opt_AE = network_loader.load_attribute_encoder_net( id=opt.which_model_AE, gpu_ids=opt.gpu_ids) if opt.which_model_FeatST != 'none': self.netFeatST, self.opt_FeatST = network_loader.load_feature_spatial_transformer_net( id=opt.which_model_FeatST, gpu_ids=opt.gpu_ids) self.use_FeatST = True # assert self.opt_FeatST.shape_encode == self.opt.shape_encode, 'GAN model and FeatST model has different shape encode mode' # assert self.opt_FeatST.input_mask_mode == self.opt.input_mask_mode, 'GAN model and FeatST model has different segmentation input mode' else: self.use_FeatST = False if self.is_train: self.netD = networks.define_D(opt) if opt.which_model_init_netG != 'none' and not opt.continue_train: self.load_network(self.netG, 'G', 'latest', opt.which_model_init_netG) if not self.is_train or opt.continue_train: self.load_network(self.netG, 'G', opt.which_epoch) if self.is_train: self.load_network(self.netD, 'D', opt.which_epoch) if self.is_train: self.fake_pool = ImagePool(opt.pool_size) ################################### # define loss functions and loss buffers ################################### if opt.which_gan in {'dcgan', 'lsgan'}: self.crit_GAN = networks.GANLoss( use_lsgan=opt.which_gan == 'lsgan', tensor=self.Tensor) else: # WGAN loss will be calculated in self.backward_D_wgangp and self.backward_G self.crit_GAN = None self.crit_L1 = nn.L1Loss() self.crit_attr = nn.BCELoss() self.loss_functions = [] self.loss_functions.append(self.crit_GAN) self.loss_functions.append(self.crit_L1) self.loss_functions.append(self.crit_attr) if self.opt.loss_weight_vgg > 0: self.crit_vgg = networks.VGGLoss(self.gpu_ids) self.loss_functions.append(self.crit_vgg) ################################### # create optimizers ################################### self.schedulers = [] self.optimizers = [] self.optim_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, opt.beta2)) self.optim_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim_G) self.optimizers.append(self.optim_D) for optim in self.optimizers: self.schedulers.append(networks.get_scheduler(optim, opt)) # color transformation from std to imagenet # img_imagenet = img_std * a + b self.trans_std_to_imagenet = { 'a': Variable(self.Tensor([0.5 / 0.229, 0.5 / 0.224, 0.5 / 0.225]), requires_grad=False).view(3, 1, 1), 'b': Variable(self.Tensor([(0.5 - 0.485) / 0.229, (0.5 - 0.456) / 0.224, (0.5 - 0.406) / 0.225]), requires_grad=False).view(3, 1, 1) }
def initialize(self, opt): super(VUnetPoseTransferModel, self).initialize(opt) ################################### # define transformer ################################### self.netT = networks.VariationalUnet( input_nc_dec = self.get_pose_dim(opt.pose_type), input_nc_enc = self.get_appearance_dim(opt.appearance_type), output_nc = self.get_output_dim(opt.output_type), nf = opt.vunet_nf, max_nf = opt.vunet_max_nf, input_size = opt.fine_size, n_latent_scales = opt.vunet_n_latent_scales, bottleneck_factor = opt.vunet_bottleneck_factor, box_factor = opt.vunet_box_factor, n_residual_blocks = 2, norm_layer = networks.get_norm_layer(opt.norm), activation = nn.ReLU(False), use_dropout = False, gpu_ids = opt.gpu_ids, output_tanh = False, ) if opt.gpu_ids: self.netT.cuda() networks.init_weights(self.netT, init_type=opt.init_type) ################################### # define discriminator ################################### self.use_GAN = self.is_train and opt.loss_weight_gan > 0 if self.use_GAN: self.netD = networks.define_D_from_params( input_nc=3+self.get_pose_dim(opt.pose_type) if opt.D_cond else 3, ndf=opt.D_nf, which_model_netD='n_layers', n_layers_D=opt.D_n_layer, norm=opt.norm, which_gan=opt.which_gan, init_type=opt.init_type, gpu_ids=opt.gpu_ids) else: self.netD = None ################################### # loss functions ################################### self.crit_psnr = networks.PSNR() self.crit_ssim = networks.SSIM() if self.is_train: self.optimizers =[] self.crit_vgg = networks.VGGLoss_v2(self.gpu_ids, opt.content_layer_weight, opt.style_layer_weight, opt.shifted_style) # self.crit_vgg_old = networks.VGGLoss(self.gpu_ids) self.optim = torch.optim.Adam(self.netT.parameters(), lr=opt.lr, betas=(opt.beta1, opt.beta2), weight_decay=opt.weight_decay) self.optimizers += [self.optim] if self.use_GAN: self.crit_GAN = networks.GANLoss(use_lsgan=opt.which_gan=='lsgan', tensor=self.Tensor) self.optim_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr_D, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim_D) # todo: add pose loss self.fake_pool = ImagePool(opt.pool_size) ################################### # load trained model ################################### if not self.is_train: self.load_network(self.netT, 'netT', opt.which_epoch) elif opt.continue_train: self.load_network(self.netT, 'netT', opt.which_epoch) self.load_optim(self.optim, 'optim', opt.which_epoch) if self.use_GAN: self.load_network(self.netD, 'netD', opt.which_epoch) self.load_optim(self.optim_D, 'optim_D', opt.which_epoch) ################################### # schedulers ################################### if self.is_train: self.schedulers = [] for optim in self.optimizers: self.schedulers.append(networks.get_scheduler(optim, opt))
def initialize(self, cfg): self.cfg = cfg ## set devices if cfg['GPU_IDS']: assert (torch.cuda.is_available()) self.device = torch.device('cuda:{}'.format(cfg['GPU_IDS'][0])) torch.backends.cudnn.benchmark = True print('Using %d GPUs' % len(cfg['GPU_IDS'])) else: self.device = torch.device('cpu') # define network if cfg['ARCHI'] == 'alexnet': self.netB = networks.netB_alexnet() self.netH = networks.netH_alexnet() if self.cfg['USE_DA'] and self.cfg['TRAIN']: self.netD = networks.netD_alexnet(self.cfg['DA_LAYER']) elif cfg['ARCHI'] == 'vgg16': raise NotImplementedError self.netB = networks.netB_vgg16() self.netH = networks.netH_vgg16() if self.cfg['USE_DA'] and self.cfg['TRAIN']: self.netD = netD_vgg16(self.cfg['DA_LAYER']) elif 'resnet' in cfg['ARCHI']: self.netB = networks.netB_resnet34() self.netH = networks.netH_resnet34() if self.cfg['USE_DA'] and self.cfg['TRAIN']: self.netD = networks.netD_resnet(self.cfg['DA_LAYER']) else: raise ValueError('Un-supported network') ## initialize network param. self.netB = networks.init_net(self.netB, cfg['GPU_IDS'], 'xavier') self.netH = networks.init_net(self.netH, cfg['GPU_IDS'], 'xavier') if self.cfg['USE_DA'] and self.cfg['TRAIN']: self.netD = networks.init_net(self.netD, cfg['GPU_IDS'], 'xavier') # loss, optimizer, and scherduler if cfg['TRAIN']: self.total_steps = 0 ## Output path self.save_dir = os.path.join( cfg['OUTPUT_PATH'], cfg['ARCHI'], datetime.now().strftime("%Y-%m-%d_%H-%M-%S")) if not os.path.isdir(self.save_dir): os.makedirs(self.save_dir) # self.logger = Logger(self.save_dir) ## model names self.model_names = ['netB', 'netH'] ## loss self.criterionGAN = networks.GANLoss().to(self.device) self.criterionDepth1 = torch.nn.MSELoss().to(self.device) self.criterionNorm = torch.nn.CosineEmbeddingLoss().to(self.device) self.criterionEdge = torch.nn.BCELoss().to(self.device) ## optimizers self.lr = cfg['LR'] self.optimizers = [] self.optimizer_B = torch.optim.Adam(self.netB.parameters(), lr=cfg['LR'], betas=(cfg['BETA1'], cfg['BETA2'])) self.optimizer_H = torch.optim.Adam(self.netH.parameters(), lr=cfg['LR'], betas=(cfg['BETA1'], cfg['BETA2'])) self.optimizers.append(self.optimizer_B) self.optimizers.append(self.optimizer_H) if cfg['USE_DA']: self.real_pool = ImagePool(cfg['POOL_SIZE']) self.syn_pool = ImagePool(cfg['POOL_SIZE']) self.model_names.append('netD') ## use SGD for discriminator self.optimizer_D = torch.optim.SGD( self.netD.parameters(), lr=cfg['LR'], momentum=cfg['MOMENTUM'], weight_decay=cfg['WEIGHT_DECAY']) self.optimizers.append(self.optimizer_D) ## LR scheduler self.schedulers = [ networks.get_scheduler(optimizer, cfg) for optimizer in self.optimizers ] else: ## testing self.model_names = ['netB', 'netH'] self.criterionDepth1 = torch.nn.MSELoss().to(self.device) self.criterionNorm = torch.nn.CosineEmbeddingLoss( reduction='none').to(self.device) self.load_dir = os.path.join(cfg['CKPT_PATH']) self.criterionNorm_eval = torch.nn.CosineEmbeddingLoss( reduction='none').to(self.device) if cfg['TEST'] or cfg['RESUME']: self.load_networks(cfg['EPOCH_LOAD'])
def __init__(self, hyperparameters): super(myMUNIT_Trainer, self).__init__() lr = hyperparameters['lr'] # Initiate the networks self.style_dim = hyperparameters['gen']['style_dim'] self.enc_a = networks.define_E(input_nc=3, output_nc=self.style_dim, ndf=64) # encoder for domain a self.enc_b = networks.define_E(input_nc=3, output_nc=self.style_dim, ndf=64) # encoder for domain b self.gen_a = networks.define_G(input_nc=3, output_nc=3, nz=self.style_dim, ngf=64) # generator for domain a self.gen_b = networks.define_G(input_nc=3, output_nc=3, nz=self.style_dim, ngf=64) # generator for domain b self.dis_a = networks.define_D(input_nc=3, ndf=64, norm='instance', num_Ds=2) # discriminator for domain a self.dis_b = networks.define_D(input_nc=3, ndf=64, norm='instance', num_Ds=2) # discriminator for domain b self.netVGGF = networks.define_VGGF() self.instancenorm = nn.InstanceNorm2d(512, affine=False) # Initiate the criterions or loss functions self.criterionGAN = networks.GANLoss( mse_loss=True, tensor=torch.cuda.FloatTensor) # criterion GAN adversarial loss self.wGANloss = networks.wGANLoss( tensor=torch.cuda.FloatTensor) # wGAN adversarial loss self.criterionL1 = torch.nn.L1Loss() # L1 loss self.criterionL2 = networks.L2Loss() # L2 loss self.criterionZ = torch.nn.L1Loss() # L1 loss between code self.criterionC = networks.ContentLoss( vgg_features=self.netVGGF) # content loss self.criterionS = networks.StyleLoss( vgg_features=self.netVGGF) # style loss self.criterionC_l = networks.ContentLoss( vgg_features=self.netVGGF) # local content loss self.criterionS_l = networks.StyleLoss( vgg_features=self.netVGGF) # local style loss self.criterionHisogram = networks.HistogramLoss( vgg_features=self.netVGGF) # histogram loss self.Feature_map_im = networks.Feature_map_im( vgg_features=self.netVGGF) # show feature map # fix the noise used in sampling self.s_a = torch.randn(8, self.style_dim, 1, 1).cuda() self.s_b = torch.randn(8, self.style_dim, 1, 1).cuda() # Setup the optimizers beta1 = hyperparameters['beta1'] beta2 = hyperparameters['beta2'] dis_params = list(self.dis_a.parameters()) + list( self.dis_b.parameters()) gen_params = list(self.gen_a.parameters()) + list( self.gen_b.parameters()) self.dis_opt = torch.optim.Adam( [p for p in dis_params if p.requires_grad], lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay']) self.gen_opt = torch.optim.Adam( [p for p in gen_params if p.requires_grad], lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay']) self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters) self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters) # Load VGG model if needed if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0: self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models') self.vgg.eval() for param in self.vgg.parameters(): param.requires_grad = False
def __init__(self, opt): """Initialize the CycleGAN class. Parameters: opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions """ self.opt = opt self.isTrain = opt.isTrain self.device = opt.device self.model_save_dir = opt.model_dir self.loss_names = [] self.model_names = [] self.optimizers = [] self.image_paths = [] self.epoch = 0 self.num_epochs = opt.nr_epochs # 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', 'D_B', 'G_B', 'cycle_B'] # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals> 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.Generator(opt.input_nc, opt.output_nc, opt.ngf, opt.n_layers_G).to(self.device) self.netG_B = networks.Generator(opt.output_nc, opt.input_nc, opt.ngf, opt.n_layers_G).to(self.device) if self.isTrain: # define discriminators self.netD_A = networks.NLayerDiscriminator( opt.output_nc, opt.ndf, opt.n_layers_D).to(self.device) self.netD_B = networks.NLayerDiscriminator( opt.input_nc, opt.ndf, opt.n_layers_D).to(self.device) # 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: self.fake_A_pool = networks.MotionPool( opt.pool_size ) # create image buffer to store previously generated images self.fake_B_pool = networks.MotionPool( 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) self.schedulers = [ networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers ]
device = torch.device("cuda:3") # Shouldn't be patchGAN, add fc layer at end #model = networks.define_D(6, 64, 'n_layers', n_layers_D=3, use_sigmoid=False, out_channels=256, glob=True) model = networks.define_D(6, 64, 'n_layers', n_layers_D=3, use_sigmoid=False) #model = networks.Siamese() #model = networks.GlobalLocal() #model = torch.nn.DataParallel(model, device_ids=(0,1,3,4)) #model = networks.SiameseResnet() chkpt = torch.load('checkpoints/patch/179.pth') #chkpt = torch.load('checkpoints/localGlobal/190.pth') #chkpt = torch.load('checkpoints/SiameseResnet/17.pth') model.load_state_dict(chkpt['state_dict']) model.to(device) patch_loss = networks.GANLoss(use_lsgan=False).to(device) total_steps = 0 dataset = FrankensteinDataset() #dataset.initialize('../datasets/street_view/sides/') #dataset.initialize('../../../data/semanticLandscapes512/train_img') dataset.initialize('../../../data/MITCVCL/imgs') train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True) total = 0 correct = 0 real = [] fake = [] fake_pred = []
def __init__(self, conf): # Acquire configuration self.conf = conf self.cur_iter = 0 self.max_iters = conf.max_iters # Define input tensor self.input_tensor = torch.FloatTensor(1, 3, conf.input_crop_size, conf.input_crop_size).cuda() self.real_example = torch.FloatTensor(1, 3, conf.output_crop_size, conf.output_crop_size).cuda() # Define networks # self.G = networks.Generator(conf.G_base_channels, conf.G_num_resblocks, conf.G_num_downscales, conf.G_use_bias, # conf.G_skip) self.G = Generator(32, 32) self.D = networks.MultiScaleDiscriminator(conf.output_crop_size, self.conf.D_max_num_scales, self.conf.D_scale_factor, self.conf.D_base_channels) self.GAN_loss_layer = networks.GANLoss() self.Reconstruct_loss = networks.WeightedMSELoss(use_L1=conf.use_L1) self.RandCrop = networks.RandomCrop( [conf.input_crop_size, conf.input_crop_size], must_divide=conf.must_divide) self.SwapCrops = networks.SwapCrops(conf.crop_swap_min_size, conf.crop_swap_max_size) # Make all networks run on GPU self.G.cuda() self.D.cuda() self.GAN_loss_layer.cuda() self.Reconstruct_loss.cuda() self.RandCrop.cuda() self.SwapCrops.cuda() # Define loss function self.criterionGAN = self.GAN_loss_layer.forward self.criterionReconstruction = self.Reconstruct_loss.forward # Keeping track of losses- prepare tensors self.losses_G_gan = torch.FloatTensor(conf.print_freq).cuda() self.losses_D_real = torch.FloatTensor(conf.print_freq).cuda() self.losses_D_fake = torch.FloatTensor(conf.print_freq).cuda() self.losses_G_reconstruct = torch.FloatTensor(conf.print_freq).cuda() if self.conf.reconstruct_loss_stop_iter > 0: self.losses_D_reconstruct = torch.FloatTensor( conf.print_freq).cuda() # Initialize networks self.G.apply(networks.weights_init) self.D.apply(networks.weights_init) # Initialize optimizers if self.conf.gopt_sgd: g_opt = torch.optim.SGD(self.G.parameters(), lr=conf.g_lr) else: g_opt = torch.optim.Adam(self.G.parameters(), lr=conf.g_lr, betas=(conf.beta1, 0.999)) self.optimizer_G = g_opt self.optimizer_D = torch.optim.Adam(self.D.parameters(), lr=conf.d_lr, betas=(conf.beta1, 0.999)) # Learning rate scheduler # First define linearly decaying functions (decay starts at a special iter) start_decay = conf.lr_start_decay_iter end_decay = conf.max_iters # def lr_function(n_iter): # return 1 - max(0, 1.0 * (n_iter - start_decay) / (conf.max_iters - start_decay)) lr_function = LRPolicy(start_decay, end_decay) # Define learning rate schedulers self.lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR( self.optimizer_G, lr_function) self.lr_scheduler_D = torch.optim.lr_scheduler.LambdaLR( self.optimizer_D, lr_function)
def initialize(self, opt): super(TwoStagePoseTransferModel, self).initialize(opt) ################################### # load pretrained stage-1 (coarse) network ################################### self._create_stage_1_net(opt) ################################### # define stage-2 (refine) network ################################### # local patch encoder if opt.which_model_s2e == 'patch_embed': self.netT_s2e = networks.LocalPatchEncoder( n_patch=len(opt.patch_indices), input_nc=3, output_nc=opt.s2e_nof, nf=opt.s2e_nf, max_nf=opt.s2e_max_nf, input_size=opt.patch_size, bottleneck_factor=opt.s2e_bottleneck_factor, n_residual_blocks=2, norm_layer=networks.get_norm_layer(opt.norm), activation=nn.ReLU(False), use_dropout=False, gpu_ids=opt.gpu_ids, ) s2e_nof = opt.s2e_nof elif opt.which_model_s2e == 'patch': self.netT_s2e = networks.LocalPatchRearranger( n_patch=len(opt.patch_indices), image_size=opt.fine_size, ) s2e_nof = 3 elif opt.which_model_s2e == 'seg_embed': self.netT_s2e = networks.SegmentRegionEncoder( seg_nc=self.opt.seg_nc, input_nc=3, output_nc=opt.s2e_nof, nf=opt.s2d_nf, input_size=opt.fine_size, n_blocks=3, norm_layer=networks.get_norm_layer(opt.norm), activation=nn.ReLU, use_dropout=False, grid_level=opt.s2e_grid_level, gpu_ids=opt.gpu_ids, ) s2e_nof = opt.s2e_nof + opt.s2e_grid_level else: raise NotImplementedError() if opt.gpu_ids: self.netT_s2e.cuda() # decoder if self.opt.which_model_s2d == 'resnet': self.netT_s2d = networks.ResnetGenerator( input_nc=3 + s2e_nof, output_nc=3, ngf=opt.s2d_nf, norm_layer=networks.get_norm_layer(opt.norm), activation=nn.ReLU, use_dropout=False, n_blocks=opt.s2d_nblocks, gpu_ids=opt.gpu_ids, output_tanh=False, ) elif self.opt.which_model_s2d == 'unet': self.netT_s2d = networks.UnetGenerator_v2( input_nc=3 + s2e_nof, output_nc=3, num_downs=8, ngf=opt.s2d_nf, max_nf=opt.s2d_nf * 2**3, norm_layer=networks.get_norm_layer(opt.norm), use_dropout=False, gpu_ids=opt.gpu_ids, output_tanh=False, ) elif self.opt.which_model_s2d == 'rpresnet': self.netT_s2d = networks.RegionPropagationResnetGenerator( input_nc=3 + s2e_nof, output_nc=3, ngf=opt.s2d_nf, norm_layer=networks.get_norm_layer(opt.norm), activation=nn.ReLU, use_dropout=False, nblocks=opt.s2d_nblocks, gpu_ids=opt.gpu_ids, output_tanh=False) else: raise NotImplementedError() if opt.gpu_ids: self.netT_s2d.cuda() ################################### # define discriminator ################################### self.use_GAN = self.is_train and opt.loss_weight_gan > 0 if self.use_GAN: self.netD = networks.define_D_from_params( input_nc=3 + self.get_pose_dim(opt.pose_type) if opt.D_cond else 3, ndf=opt.D_nf, which_model_netD='n_layers', n_layers_D=opt.D_n_layer, norm=opt.norm, which_gan=opt.which_gan, init_type=opt.init_type, gpu_ids=opt.gpu_ids) else: self.netD = None ################################### # loss functions ################################### self.crit_psnr = networks.PSNR() self.crit_ssim = networks.SSIM() if self.is_train: self.optimizers = [] self.crit_vgg = networks.VGGLoss_v2(self.gpu_ids, opt.content_layer_weight, opt.style_layer_weight, opt.shifted_style) self.optim = torch.optim.Adam([{ 'params': self.netT_s2e.parameters() }, { 'params': self.netT_s2d.parameters() }], lr=opt.lr, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim) if opt.train_s1: self.optim_s1 = torch.optim.Adam(self.netT_s1.parameters(), lr=opt.lr_s1, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim_s1) if self.use_GAN: self.crit_GAN = networks.GANLoss( use_lsgan=opt.which_gan == 'lsgan', tensor=self.Tensor) self.optim_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr_D, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim_D) self.fake_pool = ImagePool(opt.pool_size) ################################### # init/load model ################################### if self.is_train: if not opt.continue_train: self.load_network(self.netT_s1, 'netT', 'latest', self.opt_s1.id) networks.init_weights(self.netT_s2e, init_type=opt.init_type) networks.init_weights(self.netT_s2d, init_type=opt.init_type) if self.use_GAN: networks.init_weights(self.netD, init_type=opt.init_type) else: self.load_network(self.netT_s1, 'netT_s1', opt.which_epoch) self.load_network(self.netT_s2e, 'netT_s2e', opt.which_epoch) self.load_network(self.netT_s2d, 'netT_s2d', opt.which_epoch) self.load_optim(self.optim, 'optim', opt.which_epoch) if self.use_GAN: self.load_network(self.netD, 'netD', opt.which_epoch) self.load_optim(self.optim_D, 'optim_D', opt.which_epoch) else: self.load_network(self.netT_s1, 'netT_s1', opt.which_epoch) self.load_network(self.netT_s2e, 'netT_s2e', opt.which_epoch) self.load_network(self.netT_s2d, 'netT_s2d', opt.which_epoch) ################################### # schedulers ################################### if self.is_train: self.schedulers = [] for optim in self.optimizers: self.schedulers.append(networks.get_scheduler(optim, opt))
def initialize(self, opt): super(SupervisedPoseTransferModel, self).initialize(opt) ################################### # define transformer ################################### if opt.which_model_T == 'resnet': self.netT = networks.ResnetGenerator( input_nc=3 + self.get_pose_dim(opt.pose_type), output_nc=3, ngf=opt.T_nf, norm_layer=networks.get_norm_layer(opt.norm), use_dropout=not opt.no_dropout, n_blocks=9, gpu_ids=opt.gpu_ids) elif opt.which_model_T == 'unet': self.netT = networks.UnetGenerator_v2( input_nc=3 + self.get_pose_dim(opt.pose_type), output_nc=3, num_downs=8, ngf=opt.T_nf, norm_layer=networks.get_norm_layer(opt.norm), use_dropout=not opt.no_dropout, gpu_ids=opt.gpu_ids) else: raise NotImplementedError() if opt.gpu_ids: self.netT.cuda() networks.init_weights(self.netT, init_type=opt.init_type) ################################### # define discriminator ################################### self.use_GAN = self.is_train and opt.loss_weight_gan > 0 if self.use_GAN > 0: self.netD = networks.define_D_from_params( input_nc=3 + self.get_pose_dim(opt.pose_type) if opt.D_cond else 3, ndf=opt.D_nf, which_model_netD='n_layers', n_layers_D=3, norm=opt.norm, which_gan=opt.which_gan, init_type=opt.init_type, gpu_ids=opt.gpu_ids) else: self.netD = None ################################### # loss functions ################################### if self.is_train: self.loss_functions = [] self.schedulers = [] self.optimizers = [] self.crit_L1 = nn.L1Loss() self.crit_vgg = networks.VGGLoss_v2(self.gpu_ids) # self.crit_vgg_old = networks.VGGLoss(self.gpu_ids) self.crit_psnr = networks.PSNR() self.crit_ssim = networks.SSIM() self.loss_functions += [self.crit_L1, self.crit_vgg] self.optim = torch.optim.Adam(self.netT.parameters(), lr=opt.lr, betas=(opt.beta1, opt.beta2)) self.optimizers += [self.optim] if self.use_GAN: self.crit_GAN = networks.GANLoss( use_lsgan=opt.which_gan == 'lsgan', tensor=self.Tensor) self.optim_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr_D, betas=(opt.beta1, opt.beta2)) self.loss_functions.append(self.use_GAN) self.optimizers.append(self.optim_D) # todo: add pose loss for optim in self.optimizers: self.schedulers.append(networks.get_scheduler(optim, opt)) self.fake_pool = ImagePool(opt.pool_size) ################################### # load trained model ################################### if not self.is_train: self.load_network(self.netT, 'netT', opt.which_model)
def initialize(self, opt): BaseModel.initialize(self, opt) if opt.resize_or_crop != 'none' or not opt.isTrain: # when training at full res this causes OOM torch.backends.cudnn.benchmark = True self.isTrain = opt.isTrain input_nc = opt.label_nc if opt.label_nc != 0 else opt.input_nc ##### define networks # Generator network netG_input_nc = input_nc # Main Generator self.netG = networks.define_G(11, opt.output_nc, opt.ngf, opt.netG, opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers, opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids) self.netP = networks.define_P(44, 20, opt.ngf, opt.netG, opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers, opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids) self.netP.load_state_dict( torch.load( os.path.dirname(os.path.realpath(__file__)) + "/checkpoints/generate/parse.pth")) # Discriminator network if self.isTrain: use_sigmoid = opt.no_lsgan netD_input_nc = input_nc + opt.output_nc netB_input_nc = opt.output_nc * 2 self.netD = networks.define_D(netD_input_nc, opt.ndf, opt.n_layers_D, opt.norm, use_sigmoid, opt.num_D, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids) #self.netB = networks.define_B(netB_input_nc, opt.output_nc, 32, 3, 3, opt.norm, gpu_ids=self.gpu_ids) if self.opt.verbose: print('---------- Networks initialized -------------') # load networks if not self.isTrain or opt.continue_train or opt.load_pretrain: pretrained_path = '' if not self.isTrain else opt.load_pretrain self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path) if self.isTrain: self.load_network(self.netD, 'D', opt.which_epoch, pretrained_path) # set loss functions and optimizers if self.isTrain: if opt.pool_size > 0 and (len(self.gpu_ids)) > 1: raise NotImplementedError( "Fake Pool Not Implemented for MultiGPU") self.fake_pool = ImagePool(opt.pool_size) self.old_lr = opt.lr # define loss functions self.loss_filter = self.init_loss_filter(not opt.no_ganFeat_loss, not opt.no_vgg_loss) self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor) self.criterionFeat = torch.nn.L1Loss() if not opt.no_vgg_loss: self.criterionVGG = networks.VGGLoss(self.gpu_ids) self.criterionStyle = networks.StyleLoss(self.gpu_ids) # Names so we can breakout loss self.loss_names = self.loss_filter('G_GAN', 'G_GAN_Feat', 'G_VGG', 'D_real', 'D_fake') # initialize optimizers # optimizer G if opt.niter_fix_global > 0: import sys if sys.version_info >= (3, 0): finetune_list = set() else: from sets import Set finetune_list = Set() params_dict = dict(self.netG.named_parameters()) params = [] for key, value in params_dict.items(): if key.startswith('model' + str(opt.n_local_enhancers)): params += [value] finetune_list.add(key.split('.')[0]) print( '------------- Only training the local enhancer network (for %d epochs) ------------' % opt.niter_fix_global) print('The layers that are finetuned are ', sorted(finetune_list)) else: params = list(self.netG.parameters())+list(self.netimage.parameters())+list(self.netcolor.parameters())\ +list(self.netlabel.parameters())+list(self.netsketch.parameters())+list(self.classfier.parameters()) # params.extend(list(self.netimage.parameters())) self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999)) # optimizer D params = list(self.netD.parameters()) self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
def initialize(self, opt): super(PoseTransferModel, self).initialize(opt) ################################### # define generator ################################### if opt.which_model_G == 'unet': self.netG = networks.UnetGenerator( input_nc=self.get_tensor_dim('+'.join( [opt.G_appearance_type, opt.G_pose_type])), output_nc=3, nf=opt.G_nf, max_nf=opt.G_max_nf, num_scales=opt.G_n_scale, n_residual_blocks=2, norm=opt.G_norm, activation=nn.LeakyReLU(0.1) if opt.G_activation == 'leaky_relu' else nn.ReLU(), use_dropout=opt.use_dropout, gpu_ids=opt.gpu_ids) elif opt.which_model_G == 'dual_unet': self.netG = networks.DualUnetGenerator( pose_nc=self.get_tensor_dim(opt.G_pose_type), appearance_nc=self.get_tensor_dim(opt.G_appearance_type), output_nc=3, aux_output_nc=[], nf=opt.G_nf, max_nf=opt.G_max_nf, num_scales=opt.G_n_scale, num_warp_scales=opt.G_n_warp_scale, n_residual_blocks=2, norm=opt.G_norm, vis_mode=opt.G_vis_mode, activation=nn.LeakyReLU(0.1) if opt.G_activation == 'leaky_relu' else nn.ReLU(), use_dropout=opt.use_dropout, no_end_norm=opt.G_no_end_norm, gpu_ids=opt.gpu_ids, ) if opt.gpu_ids: self.netG.cuda() networks.init_weights(self.netG, init_type=opt.init_type) ################################### # define external pixel warper ################################### if opt.G_pix_warp: pix_warp_n_scale = opt.G_n_scale self.netPW = networks.UnetGenerator_MultiOutput( input_nc=self.get_tensor_dim(opt.G_pix_warp_input_type), output_nc=[1], # only use one output branch (weight mask) nf=32, max_nf=128, num_scales=pix_warp_n_scale, n_residual_blocks=2, norm=opt.G_norm, activation=nn.ReLU(False), use_dropout=False, gpu_ids=opt.gpu_ids) if opt.gpu_ids: self.netPW.cuda() networks.init_weights(self.netPW, init_type=opt.init_type) ################################### # define discriminator ################################### self.use_gan = self.is_train and self.opt.loss_weight_gan > 0 if self.use_gan: self.netD = networks.NLayerDiscriminator( input_nc=self.get_tensor_dim(opt.D_input_type_real), ndf=opt.D_nf, n_layers=opt.D_n_layers, use_sigmoid=(opt.gan_type == 'dcgan'), output_bias=True, gpu_ids=opt.gpu_ids, ) if opt.gpu_ids: self.netD.cuda() networks.init_weights(self.netD, init_type=opt.init_type) ################################### # load optical flow model ################################### if opt.flow_on_the_fly: self.netF = load_flow_network(opt.pretrained_flow_id, opt.pretrained_flow_epoch, opt.gpu_ids) self.netF.eval() if opt.gpu_ids: self.netF.cuda() ################################### # loss and optimizers ################################### self.crit_psnr = networks.PSNR() self.crit_ssim = networks.SSIM() if self.is_train: self.crit_vgg = networks.VGGLoss( opt.gpu_ids, shifted_style=opt.shifted_style_loss, content_weights=opt.vgg_content_weights) if opt.G_pix_warp: # only optimze netPW self.optim = torch.optim.Adam(self.netPW.parameters(), lr=opt.lr, betas=(opt.beta1, opt.beta2), weight_decay=opt.weight_decay) else: self.optim = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, opt.beta2), weight_decay=opt.weight_decay) self.optimizers = [self.optim] if self.use_gan: self.crit_gan = networks.GANLoss( use_lsgan=(opt.gan_type == 'lsgan')) if self.gpu_ids: self.crit_gan.cuda() self.optim_D = torch.optim.Adam( self.netD.parameters(), lr=opt.lr_D, betas=(opt.beta1, opt.beta2), weight_decay=opt.weight_decay_D) self.optimizers += [self.optim_D] ################################### # load trained model ################################### if not self.is_train: # load trained model for testing self.load_network(self.netG, 'netG', opt.which_epoch) if opt.G_pix_warp: self.load_network(self.netPW, 'netPW', opt.which_epoch) elif opt.pretrained_G_id is not None: # load pretrained network self.load_network(self.netG, 'netG', opt.pretrained_G_epoch, opt.pretrained_G_id) elif opt.resume_train: # resume training self.load_network(self.netG, 'netG', opt.which_epoch) self.load_optim(self.optim, 'optim', opt.which_epoch) if self.use_gan: self.load_network(self.netD, 'netD', opt.which_epoch) self.load_optim(self.optim_D, 'optim_D', opt.which_epoch) if opt.G_pix_warp: self.load_network(self.netPW, 'netPW', opt.which_epoch) ################################### # schedulers ################################### if self.is_train: self.schedulers = [] for optim in self.optimizers: self.schedulers.append(networks.get_scheduler(optim, opt))
def initialize(self, opt): super(MultimodalDesignerGAN_V2, self).initialize(opt) ################################### # define networks ################################### self.modules = {} # shape branch if opt.which_model_netG != 'unet': self.shape_encoder = networks.define_image_encoder(opt, 'shape') self.modules['shape_encoder'] = self.shape_encoder else: self.shape_encoder = None # edge branch if opt.use_edge: self.edge_encoder = networks.define_image_encoder(opt, 'edge') self.modules['edge_encoder'] = self.edge_encoder else: self.encoder_edge = None # color branch if opt.use_color: self.color_encoder = networks.define_image_encoder(opt, 'color') self.modules['color_encoder'] = self.color_encoder else: self.color_encoder = None # fusion model if opt.ftn_model == 'none': # shape_feat, edge_feat and color_feat will be simply upmpled to same size (size of shape_feat) and concatenated pass elif opt.ftn_model == 'concat': assert opt.use_edge or opt.use_color if opt.use_edge: self.edge_trans_net = networks.define_feature_fusion_network( name='FeatureConcatNetwork', feat_nc=opt.edge_nof, guide_nc=opt.shape_nof, nblocks=opt.ftn_nblocks, norm=opt.norm, gpu_ids=self.gpu_ids, init_type=opt.init_type) self.modules['edge_trans_net'] = self.edge_trans_net if opt.use_color: self.color_trans_net = networks.define_feature_fusion_network( name='FeatureConcatNetwork', feat_nc=opt.color_nof, guide_nc=opt.shape_nof, nblocks=opt.ftn_nblocks, norm=opt.norm, gpu_ids=self.gpu_ids, init_type=opt.init_type) self.modules['color_trans_net'] = self.color_trans_net elif opt.ftn_model == 'reduce': assert opt.use_edge or opt.use_color if opt.use_edge: self.edge_trans_net = networks.define_feature_fusion_network( name='FeatureReduceNetwork', feat_nc=opt.edge_nof, guide_nc=opt.shape_nof, nblocks=opt.ftn_nblocks, ndowns=opt.ftn_ndowns, norm=opt.norm, gpu_ids=self.gpu_ids, init_type=opt.init_type) self.modules['edge_trans_net'] = self.edge_trans_net if opt.use_color: self.color_trans_net = networks.define_feature_fusion_network( name='FeatureReduceNetwork', feat_nc=opt.color_nof, guide_nc=opt.shape_nof, nblocks=opt.ftn_nblocks, ndowns=opt.ftn_ndowns, norm=opt.norm, gpu_ids=self.gpu_ids, init_type=opt.init_type) self.modules['color_trans_net'] = self.color_trans_net elif opt.ftn_model == 'trans': assert opt.use_edge or opt.use_color if opt.use_edge: self.edge_trans_net = networks.define_feature_fusion_network( name='FeatureTransformNetwork', feat_nc=opt.edge_nof, guide_nc=opt.shape_nof, nblocks=opt.ftn_nblocks, feat_size=opt.feat_size_lr, norm=opt.norm, gpu_ids=self.gpu_ids, init_type=opt.init_type) self.modules['edge_trans_net'] = self.edge_trans_net if opt.use_color: self.color_trans_net = networks.define_feature_fusion_network( name='FeatureTransformNetwork', feat_nc=opt.color_nof, guide_nc=opt.shape_nof, nblocks=opt.ftn_nblocks, feat_size=opt.feat_size_lr, norm=opt.norm, gpu_ids=self.gpu_ids, init_type=opt.init_type) self.modules['color_trans_net'] = self.color_trans_net # netG self.netG = networks.define_generator(opt) self.modules['netG'] = self.netG # netD if self.is_train: self.netD = networks.define_D(opt) self.modules['netD'] = self.netD ################################### # load weights ################################### if self.is_train: if opt.continue_train: for label, net in self.modules.iteritems(): self.load_network(net, label, opt.which_epoch) else: if opt.which_model_init != 'none': # load pretrained entire model for label, net in self.modules.iteritems(): self.load_network(net, label, 'latest', opt.which_model_init, forced=False) else: # load pretrained encoder if opt.which_model_netG != 'unet' and opt.pretrain_shape: self.load_network(self.shape_encoder, 'shape_encoder', 'latest', opt.which_model_init_shape_encoder) if opt.use_edge and opt.pretrain_edge: self.load_network(self.edge_encoder, 'edge_encoder', 'latest', opt.which_model_init_edge_encoder) if opt.use_color and opt.pretrain_color: self.load_network(self.color_encoder, 'color_encoder', 'latest', opt.which_model_init_color_encoder) else: for label, net in self.modules.iteritems(): if label != 'netD': self.load_network(net, label, opt.which_epoch) ################################### # prepare for training ################################### if self.is_train: self.fake_pool = ImagePool(opt.pool_size) ################################### # define loss functions ################################### self.loss_functions = [] if opt.which_gan in {'dcgan', 'lsgan'}: self.crit_GAN = networks.GANLoss( use_lsgan=opt.which_gan == 'lsgan', tensor=self.Tensor) self.loss_functions.append(self.crit_GAN) else: # WGAN loss will be calculated in self.backward_D_wgangp and self.backward_G self.crit_GAN = None self.crit_L1 = nn.L1Loss() self.loss_functions.append(self.crit_L1) if self.opt.loss_weight_vgg > 0: self.crit_vgg = networks.VGGLoss(self.gpu_ids) self.loss_functions.append(self.crit_vgg) if self.opt.G_output_seg: self.crit_CE = nn.CrossEntropyLoss() self.loss_functions.append(self.crit_CE) self.crit_psnr = networks.SmoothLoss(networks.PSNR()) self.loss_functions.append(self.crit_psnr) ################################### # create optimizers ################################### self.schedulers = [] self.optimizers = [] # G optimizer G_module_list = [ 'shape_encoder', 'edge_encoder', 'color_encoder', 'netG' ] G_param_groups = [{ 'params': self.modules[m].parameters() } for m in G_module_list if m in self.modules] self.optim_G = torch.optim.Adam(G_param_groups, lr=opt.lr, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim_G) # D optimizer self.optim_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr_D, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim_D) # feature transfer network optimizer FTN_module_list = ['edge_trans_net', 'color_trans_net'] FTN_param_groups = [{ 'params': self.modules[m].parameters() } for m in FTN_module_list if m in self.modules] if len(FTN_param_groups) > 0: self.optim_FTN = torch.optim.Adam(FTN_param_groups, lr=opt.lr_FTN, betas=(0.9, 0.999)) self.optimizers.append(self.optim_FTN) else: self.optim_FTN = None # schedulers for optim in self.optimizers: self.schedulers.append(networks.get_scheduler(optim, opt))
def __init__(self, p): super(CycleGAN, self).__init__(p) nb = p.batchSize size = p.cropSize # load/define models # 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) self.netG_A = networks.define_G(p.input_nc, p.output_nc, p.ngf, p.which_model_netG, p.norm, not p.no_dropout, p.init_type, self.gpu_ids) self.netG_B = networks.define_G(p.output_nc, p.input_nc, p.ngf, p.which_model_netG, p.norm, not p.no_dropout, p.init_type, self.gpu_ids) if self.isTrain: use_sigmoid = p.no_lsgan self.netD_A = networks.define_D(p.output_nc, p.ndf, p.which_model_netD, p.n_layers_D, p.norm, use_sigmoid, p.init_type, self.gpu_ids) self.netD_B = networks.define_D(p.input_nc, p.ndf, p.which_model_netD, p.n_layers_D, p.norm, use_sigmoid, p.init_type, self.gpu_ids) if not self.isTrain or p.continue_train: which_epoch = p.which_epoch self.load_model(self.netG_A, 'G_A', which_epoch) self.load_model(self.netG_B, 'G_B', which_epoch) if self.isTrain: self.load_model(self.netD_A, 'D_A', which_epoch) self.load_model(self.netD_B, 'D_B', which_epoch) if self.isTrain: self.old_lr = p.lr self.fake_A_pool = ImagePool(p.pool_size) self.fake_B_pool = ImagePool(p.pool_size) # define loss functions self.criterionGAN = networks.GANLoss(use_lsgan=not p.no_lsgan, tensor=self.Tensor) self.criterionCycle = torch.nn.L1Loss() self.criterionIdt = torch.nn.L1Loss() # initialize optimizers self.optimizer_G = torch.optim.Adam(itertools.chain( self.netG_A.parameters(), self.netG_B.parameters()), lr=p.lr, betas=(p.beta1, 0.999)) self.optimizer_D_A = torch.optim.Adam(self.netD_A.parameters(), lr=p.lr, betas=(p.beta1, 0.999)) self.optimizer_D_B = torch.optim.Adam(self.netD_B.parameters(), lr=p.lr, betas=(p.beta1, 0.999)) self.optimizers = [] self.schedulers = [] self.optimizers.append(self.optimizer_G) 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, p))
def initialize(self, opt): super(EncoderDecoderFramework_DFN, self).initialize(opt) ################################### # define encoder ################################### self.encoder = networks.define_encoder_v2(opt) ################################### # define decoder ################################### self.decoder = networks.define_decoder_v2(opt) ################################### # guide encoder ################################### self.guide_encoder, self.opt_guide = networks.load_encoder_v2(opt, opt.which_model_guide) self.guide_encoder.eval() for p in self.guide_encoder.parameters(): p.requires_grad = False ################################### # DFN Modules ################################### if self.opt.use_dfn: self.dfn = networks.define_DFN_from_params(nf=opt.nof, ng=self.opt_guide.nof, nmid=opt.dfn_nmid, feat_size=opt.feat_size, local_size=opt.dfn_local_size, nblocks=opt.dfn_nblocks, norm=opt.norm, gpu_ids=opt.gpu_ids, init_type=opt.init_type) else: self.dfn = None ################################### # Discriminator ################################### self.use_GAN = self.is_train and opt.loss_weight_gan > 0 if self.is_train: if self.use_GAN: # if not self.opt.D_cond: # input_nc = self.decoder.output_nc # else: # input_nc = self.decoder.output_nc + self.encoder.input_nc if self.opt.gan_level == 'image': input_nc = self.decoder.output_nc elif self.opt.gan_level == 'feature': input_nc = self.opt.nof if self.opt.D_cond: if self.opt.D_cond_type == 'cond': input_nc += self.encoder.input_nc elif self.opt.D_cond_type == 'pair': input_nc += input_nc self.netD = networks.define_D_from_params(input_nc=input_nc, ndf=64, which_model_netD='n_layers', n_layers_D=3, norm=opt.norm, which_gan=opt.which_gan, init_type=opt.init_type, gpu_ids=opt.gpu_ids) else: self.netD = None ################################### # loss functions ################################### if self.is_train: self.loss_functions = [] self.schedulers = [] self.crit_image = nn.L1Loss() self.crit_seg = nn.CrossEntropyLoss() self.crit_edge = nn.BCELoss() self.loss_functions += [self.crit_image, self.crit_seg, self.crit_edge] if self.opt.use_dfn: self.optim = torch.optim.Adam([{'params': self.encoder.parameters()}, {'params': self.decoder.parameters()}, {'params': self.dfn.parameters()}], lr=opt.lr, betas=(opt.beta1, opt.beta2)) else: self.optim = torch.optim.Adam([{'params': self.encoder.parameters()}, {'params': self.decoder.parameters()}], lr=opt.lr, betas=(opt.beta1, opt.beta2)) self.optimizers = [self.optim] # GAN loss and optimizers if self.use_GAN > 0: self.crit_GAN = networks.GANLoss(use_lsgan=opt.which_gan=='lsgan', tensor=self.Tensor) self.optim_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr_D, betas=(0.5, 0.999)) self.loss_functions += [self.crit_GAN] self.optimizers += [self.optim_D] for optim in self.optimizers: self.schedulers.append(networks.get_scheduler(optim, opt)) ################################### # load trained model ################################### if not self.is_train: self.load_network(self.encoder, 'encoder', opt.which_epoch) self.load_network(self.decoder, 'decoder', opt.which_epoch) if opt.use_dfn: self.load_network(self.dfn, 'dfn', opt.which_epoch)
netD_B = networks.define_D(opt.input_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, False) device = 'cpu' if opt.cuda: netG.cuda() netD_B.cuda() device = 'cuda' # netG.apply(weights_init_normal) # netD_A.apply(weights_init_normal) netD_B.apply(weights_init_normal) # Lossess criterionGAN = networks.GANLoss().to(device) criterion_MSE = torch.nn.MSELoss(reduce=False) criterion_cycle = torch.nn.L1Loss() criterion_identity = torch.nn.L1Loss() ############### only use B to A ########################### optimizer_G = torch.optim.Adam(netG.parameters(), lr=opt.lr, betas=(0.5, 0.999)) optimizer_D_B = torch.optim.Adam(netD_B.parameters(), lr=opt.lr, betas=(0.5, 0.999)) lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(optimizer_G, lr_lambda=LambdaLR( opt.n_epochs, opt.epoch,
def initialize(self, opt): super(MultimodalDesignerGAN, self).initialize(opt) ################################### # load/define networks ################################### # basic G self.netG = networks.define_G(opt) # encoders self.encoders = {} if opt.use_edge: self.edge_encoder = networks.define_image_encoder(opt, 'edge') self.encoders['edge_encoder'] = self.edge_encoder if opt.use_color: self.color_encoder = networks.define_image_encoder(opt, 'color') self.encoders['color_encoder'] = self.color_encoder if opt.use_attr: self.attr_encoder, self.opt_AE = network_loader.load_attribute_encoder_net( id=opt.which_model_AE, gpu_ids=opt.gpu_ids) # basic D and auxiliary Ds if self.is_train: # basic D self.netD = networks.define_D(opt) # auxiliary Ds self.auxiliaryDs = {} if opt.use_edge_D: assert opt.use_edge self.netD_edge = networks.define_D_from_params( input_nc=opt.edge_nof + 3, ndf=opt.ndf, which_model_netD=opt.which_model_netD, n_layers_D=opt.n_layers_D, norm=opt.norm, which_gan='dcgan', init_type=opt.init_type, gpu_ids=opt.gpu_ids) self.auxiliaryDs['D_edge'] = self.netD_edge if opt.use_color_D: assert opt.use_color self.netD_color = networks.define_D_from_params( input_nc=opt.color_nof + 3, ndf=opt.ndf, which_model_netD=opt.which_model_netD, n_layers_D=opt.n_layers_D, norm=opt.norm, which_gan='dcgan', init_type=opt.init_type, gpu_ids=opt.gpu_ids) self.auxiliaryDs['D_color'] = self.netD_color if opt.use_attr_D: assert opt.use_attr attr_nof = opt.n_attr_feat if opt.attr_cond_type in { 'feat', 'feat_map' } else opt.n_attr self.netD_attr = networks.define_D_from_params( input_nc=attr_nof + 3, ndf=opt.ndf, which_model_netD=opt.which_model_netD, n_layers_D=opt.n_layers_D, norm=opt.norm, which_gan='dcgan', init_type=opt.init_type, gpu_ids=opt.gpu_ids) self.auxiliaryDs['D_attr'] = self.netD_attr # load weights if not opt.continue_train: if opt.which_model_init != 'none': self.load_network(self.netG, 'G', 'latest', opt.which_model_init) self.load_network(self.netD, 'D', 'latest', opt.which_model_init) for l, net in self.encoders.iteritems(): self.load_network(net, l, 'latest', opt.which_model_init) for l, net in self.auxiliaryDs.iteritems(): self.load_network(net, l, 'latest', opt.which_model_init) else: self.load_network(self.netG, 'G', opt.which_epoch) self.load_network(self.netD, 'D', opt.which_epoch) for l, net in self.encoders.iteritems(): self.load_network(net, l, opt.which_epoch) for l, net in self.auxiliaryDs.iteritems(): self.load_network(net, l, opt.which_epoch) else: self.load_network(self.netG, 'G', opt.which_epoch) for l, net in self.encoders.iteritems(): self.load_network(net, l, opt.which_epoch) if self.is_train: self.fake_pool = ImagePool(opt.pool_size) ################################### # define loss functions and loss buffers ################################### self.loss_functions = [] if opt.which_gan in {'dcgan', 'lsgan'}: self.crit_GAN = networks.GANLoss( use_lsgan=opt.which_gan == 'lsgan', tensor=self.Tensor) else: # WGAN loss will be calculated in self.backward_D_wgangp and self.backward_G self.crit_GAN = None self.loss_functions.append(self.crit_GAN) self.crit_L1 = nn.L1Loss() self.loss_functions.append(self.crit_L1) if self.opt.loss_weight_vgg > 0: self.crit_vgg = networks.VGGLoss(self.gpu_ids) self.loss_functions.append(self.crit_vgg) self.crit_psnr = networks.SmoothLoss(networks.PSNR()) self.loss_functions.append(self.crit_psnr) ################################### # create optimizers ################################### self.schedulers = [] self.optimizers = [] # optim_G will optimize parameters of netG and all image encoders (except attr_encoder) G_param_groups = [{'params': self.netG.parameters()}] for l, net in self.encoders.iteritems(): G_param_groups.append({'params': net.parameters()}) self.optim_G = torch.optim.Adam(G_param_groups, lr=opt.lr, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim_G) # optim_D will optimize parameters of netD self.optim_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr_D, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim_D) # optim_D_aux will optimize parameters of auxiliaryDs if len(self.auxiliaryDs) > 0: aux_D_param_groups = [{ 'params': net.parameters() } for net in self.auxiliaryDs.values()] self.optim_D_aux = torch.optim.Adam(aux_D_param_groups, lr=opt.lr_D, betas=(opt.beta1, opt.beta2)) self.optimizers.append(self.optim_D_aux) for optim in self.optimizers: self.schedulers.append(networks.get_scheduler(optim, opt)) # color transformation from std to imagenet # img_imagenet = img_std * a + b self.trans_std_to_imagenet = { 'a': Variable(self.Tensor([0.5 / 0.229, 0.5 / 0.224, 0.5 / 0.225]), requires_grad=False).view(3, 1, 1), 'b': Variable(self.Tensor([(0.5 - 0.485) / 0.229, (0.5 - 0.456) / 0.224, (0.5 - 0.406) / 0.225]), requires_grad=False).view(3, 1, 1) }
def TrainGenerator(real_A: tp.Numpy.Placeholder( (1, 3, args.crop_size, args.crop_size), dtype=flow.float32), real_B: tp.Numpy.Placeholder( (1, 3, args.crop_size, args.crop_size), dtype=flow.float32)): with flow.scope.placement("gpu", "0:0-0"): # G_A(A) fake_B = networks.define_G(real_A, "netG_A", ngf=args.ngf, n_blocks=9, trainable=True, reuse=True) # G_B(G_A(A)) rec_A = networks.define_G(fake_B, "netG_B", ngf=args.ngf, n_blocks=9, trainable=True, reuse=True) # G_B(B) fake_A = networks.define_G(real_B, "netG_B", ngf=args.ngf, n_blocks=9, trainable=True, reuse=True) # G_A(G_B(B)) rec_B = networks.define_G(fake_A, "netG_A", ngf=args.ngf, n_blocks=9, trainable=True, reuse=True) # Identity loss # G_A should be identity if real_B is fed: ||G_A(B) - B|| idt_A = networks.define_G(real_B, "netG_A", ngf=args.ngf, n_blocks=9, trainable=True, reuse=True) loss_idt_A = networks.L1Loss( idt_A - real_B) * args.lambda_B * args.lambda_identity # G_B should be identity if real_A is fed: ||G_B(A) - A|| idt_B = networks.define_G(real_A, "netG_B", ngf=args.ngf, n_blocks=9, trainable=True, reuse=True) loss_idt_B = networks.L1Loss( idt_B - real_A) * args.lambda_A * args.lambda_identity # GAN loss D_A(G_A(A)) netD_A_out = networks.define_D(fake_B, "netD_A", ndf=args.ndf, n_layers_D=3, trainable=False, reuse=True) loss_G_A = networks.GANLoss(netD_A_out, True) # GAN loss D_B(G_B(B)) netD_B_out = networks.define_D(fake_A, "netD_B", ndf=args.ndf, n_layers_D=3, trainable=False, reuse=True) loss_G_B = networks.GANLoss(netD_B_out, True) # Forward cycle loss || G_B(G_A(A)) - A|| loss_cycle_A = networks.L1Loss(rec_A - real_A) * args.lambda_A # Backward cycle loss || G_A(G_B(B)) - B|| loss_cycle_B = networks.L1Loss(rec_B - real_B) * args.lambda_B # combined loss and calculate gradients loss_G = loss_G_A + loss_G_B + loss_cycle_A + loss_cycle_B + loss_idt_A + loss_idt_B flow.optimizer.Adam(flow.optimizer.PiecewiseConstantScheduler( [], [args.learning_rate]), beta1=0.5).minimize(loss_G) return fake_B, rec_A, fake_A, rec_B, loss_G
import torch.nn.functional as F import torch import numpy as np # Panorama parameters batch_size = 1 NUM_SLICES = 3 SAMPLES = 5 # Number of horizon offsets to sample before finetuning # Set up deep learning stuff device = torch.device("cuda") model = networks.define_D(6, 64, 'n_layers', n_layers_D=3, use_sigmoid=False) chkpt = torch.load('checkpoints/patch_horizon/49.pth') model.load_state_dict(chkpt['state_dict']) model.to(device) patch_loss = networks.GANLoss() # Create dataset dataset = EvalDataset() #dataset.initialize('../../../data/semanticLandscapes512/train_img', allrandom=True, return_idx=True) #dataset.initialize('../../../data/MITCVCL/coast', allrandom=True) dataset.initialize('../../../data/instagram_landscapes/anne_karin_69', allrandom=True) def convertImage(im): ''' Takes output of dataset and returns a np image ''' # undo right image flip when cat-ing im = np.concatenate((im[:3, :, :], im[3:, :, ::-1]), 2)
def __init__( self, name="experiment", phase="train", which_epoch="latest", batch_size=1, image_size=128, map_nc=1, input_nc=3, output_nc=3, num_downs=7, ngf=64, ndf=64, norm_layer="batch", pool_size=50, lr=0.0002, beta1=0.5, lambda_D=0.5, lambda_MSE=10, lambda_P=5.0, use_dropout=True, gpu_ids=[], n_layers=3, use_sigmoid=False, use_lsgan=True, upsampling="nearest", continue_train=False, checkpoints_dir="checkpoints/" ): # Define input data that will be consumed by networks self.input_A = torch.FloatTensor( batch_size, 3, image_size, image_size ) self.input_map = torch.FloatTensor( batch_size, map_nc, image_size, image_size ) norm_layer = nn.BatchNorm2d \ if norm_layer == "batch" else nn.InstanceNorm2d # Define netD and netG self.netG = networks.UnetGenerator( input_nc=input_nc, output_nc=map_nc, num_downs=num_downs, ngf=ngf, use_dropout=use_dropout, gpu_ids=gpu_ids, norm_layer=norm_layer, upsampling_layer=upsampling ) self.netD = networks.NLayerDiscriminator( input_nc=input_nc + map_nc, ndf=ndf, n_layers=n_layers, use_sigmoid=use_sigmoid, gpu_ids=gpu_ids ) # Transfer data to GPU if len(gpu_ids) > 0: self.input_A = self.input_A.cuda() self.input_map = self.input_map.cuda() self.netD.cuda() self.netG.cuda() # Initialize parameters of netD and netG self.netG.apply(networks.weights_init) self.netD.apply(networks.weights_init) # Load trained netD and netG if phase == "test" or continue_train: netG_checkpoint_file = os.path.join( checkpoints_dir, name, "netG_{}.pth".format(which_epoch) ) self.netG.load_state_dict( torch.load(netG_checkpoint_file) ) print("Restoring netG from {}".format(netG_checkpoint_file)) if continue_train: netD_checkpoint_file = os.path.join( checkpoints_dir, name, "netD_{}.pth".format(which_epoch) ) self.netD.load_state_dict( torch.load(netD_checkpoint_file) ) print("Restoring netD from {}".format(netD_checkpoint_file)) self.name = name self.gpu_ids = gpu_ids self.checkpoints_dir = checkpoints_dir # Criterions if phase == "train": self.count = 0 self.lr = lr self.lambda_D = lambda_D self.lambda_MSE = lambda_MSE self.image_pool = ImagePool(pool_size) self.criterionGAN = networks.GANLoss(use_lsgan=use_lsgan) self.criterionL1 = torch.nn.L1Loss() self.criterionMSE = torch.nn.MSELoss() # Landmark loss self.optimizer_G = torch.optim.Adam( self.netG.parameters(), lr=self.lr, betas=(beta1, 0.999) ) self.optimizer_D = torch.optim.Adam( self.netD.parameters(), lr=self.lr, betas=(beta1, 0.999) ) print('---------- Networks initialized -------------') networks.print_network(self.netG) networks.print_network(self.netD) print('-----------------------------------------------')