def __init__( self, p, ): super(CycleGAN_inference, self).__init__(p) size = p.cropSize 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) def load_model(model, checkpoint_path): model.load_state_dict(torch.load(checkpoint_path)) which_epoch = p.which_epoch G_A_checkpoint_path = os.path.join(p.checkpoints_dir, p.name, f'{which_epoch}_net_G_A.pth') G_B_checkpoint_path = os.path.join(p.checkpoints_dir, p.name, f'{which_epoch}_net_G_B.pth') load_model(self.netG_A, G_A_checkpoint_path) load_model(self.netG_B, G_B_checkpoint_path)
def init_models(self): """ Models: G_UM, G_MU, D_M, D_U """ # Networks self.G_UM = networks.define_G(input_nc=1, output_nc=1, ngf=self.config.g_conv_dim, which_model_netG=self.config.which_model_netG, norm='batch', init_type='normal', gpu_ids=self.gpu_ids) self.G_MU = networks.define_G(input_nc=1, output_nc=1, ngf=self.config.g_conv_dim, which_model_netG=self.config.which_model_netG, norm='batch', init_type='normal', gpu_ids=self.gpu_ids) self.D_M = networks.define_D(input_nc=1, ndf=self.config.d_conv_dim, which_model_netD=self.config.which_model_netD, n_layers_D=3, norm='instance', use_sigmoid=True, init_type='normal', gpu_ids=self.gpu_ids) self.D_U = networks.define_D(input_nc=1, ndf=self.config.d_conv_dim, which_model_netD=self.config.which_model_netD, n_layers_D=3, norm='instance', use_sigmoid=True, init_type='normal', gpu_ids=self.gpu_ids) # Optimisers # single optimiser for both generators self.G_optim = optim.Adam(itertools.chain(self.G_UM.parameters(), self.G_MU.parameters()), self.config.lr, betas=(self.config.beta1, self.config.beta2)) self.D_M_optim = optim.Adam(self.D_M.parameters(), lr=self.config.lr, betas=(self.config.beta1, self.config.beta2)) self.D_U_optim = optim.Adam(self.D_U.parameters(), lr=self.config.lr, betas=(self.config.beta1, self.config.beta2)) self.optimizers = [self.G_optim, self.D_M_optim, self.D_U_optim] # Schedulers self.schedulers = [] for optimizer in self.optimizers: self.schedulers.append(networks.get_scheduler(optimizer, self.config))
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_models(self): # Networks self.G_AB = networks.define_G( input_nc=self.config.input_nc, output_nc=self.config.output_nc, ngf=self.config.g_conv_dim, which_model_netG=self.config.which_model_netG, norm='batch', init_type='normal', gpu_ids=self.gpu_ids) self.D_B = networks.define_D( input_nc=self.config.input_nc, ndf=self.config.d_conv_dim, which_model_netD=self.config.which_model_netD, n_layers_D=3, norm='instance', use_sigmoid=True, init_type='normal', gpu_ids=self.gpu_ids, image_size=self.config.image_size) # Optimisers self.G_optim = optim.Adam(self.G_AB.parameters(), lr=self.config.lr, betas=(self.config.beta1, self.config.beta2)) self.D_optim = optim.Adam(self.D_B.parameters(), lr=self.config.lr, betas=(self.config.beta1, self.config.beta2)) self.optimizers = [self.G_optim, self.D_optim] # Schedulers self.schedulers = [] for optimizer in self.optimizers: self.schedulers.append( networks.get_scheduler(optimizer, self.config))
def __init__(self, args): self.ckptdir = args.ckptdir self.datadir = args.datadir self.input_mode = args.input_mode self.in_size_w, self.in_size_h = args.in_size_w, args.in_size_h self.out_size_w, self.out_size_h = args.out_size_w, args.out_size_h self.skyboxengine = SkyBox(args) self.net_G = define_G(input_nc=3, output_nc=1, ngf=64, netG=args.net_G).to(device) self.load_model() self.video_writer = cv2.VideoWriter('demo.mp4', cv2.VideoWriter_fourcc(*'MP4V'), 20.0, (args.out_size_w, args.out_size_h)) self.video_writer_cat = cv2.VideoWriter( 'demo-cat.mp4', cv2.VideoWriter_fourcc(*'MP4V'), 20.0, (2 * args.out_size_w, args.out_size_h)) if os.path.exists(args.output_dir) is False: os.mkdir(args.output_dir) self.save_jpgs = args.save_jpgs
def __init__(self, args): self.args = args self._name = args.model_name self.resume = args.resume self.device = args.device self.output_directory = args.output_dir self.model_dir = args.model_dir self.model_store_path = os.path.join(args.model_dir, args.model_name) if not os.path.exists(self.model_store_path) and args.mode == 'train': os.mkdir(self.model_store_path) # Also store all args in a text_file. self.write_args_string_to_file() if not os.path.exists(self.model_store_path) and args.mode == 'test': raise FileNotFoundError( 'Model does not exist. Please check if the model has yet been run.' ) self.model_names = [] self.loss_names = [] self.optimizer_names = [] self.losses = {} self.start_epoch = 0 self.G = define_G(args) self.G = self.G.to(self.device) if args.use_multiple_gpu: self.G = torch.nn.DataParallel(self.G) print("G [{}] initiated with {} trainable parameters".format( args.generator_model, self.num_parameters)) self.loader = None
def __init__(self, args, loader): self.args = args self._name = args.model_name self.device = args.device self.mode = args.mode self.model_dir = args.model_dir self.model_store_path = os.path.join(args.model_dir, args.model_name) if not os.path.exists(self.model_store_path) and args.mode == 'train': os.mkdir(self.model_store_path) # Also store all args in a text_file. self.write_args_string_to_file() if not os.path.exists(self.model_store_path) and args.mode == 'test': raise FileNotFoundError( 'Model does not exist. Please check if the model has yet been run.' ) self.losses = {} self.loader = loader self.epochs = args.epochs # Define the generator network. self.G = define_G(args).to(self.device) print("[{}] initiated with {} trainable parameters".format( args.backbone, self.num_parameters)) # Set the optimizer and scheduler, but wait for method-specific parameters. self.criterion = None self.optimizer = None self.learning_rate_scheduler = None if args.mode == 'train': self.learning_rate_scheduler = LearningRateScheduler( args.adjust_lr, args.lr_mode, args.learning_rate, args.epochs, args.num_iterations)
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 TestGenerator(real: tp.Numpy.Placeholder( (1, 3, args.network_input_size, args.network_input_size), dtype=flow.float32)) -> tp.Numpy: with flow.scope.placement("gpu", "0:0-0"): fake = networks.define_G(real, netG_name, ngf=64, n_blocks=9, trainable=False, reuse=True) return fake
def __init__(self, p): super(TestModel, self).__init__(p) assert (not p.isTrain) 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) which_epoch = p.which_epoch self.load_model(self.netG, 'G', which_epoch) print('---------- Networks initialized -------------') networks.print_network(self.netG) print('-----------------------------------------------')
def __init__(self, opt, ignore_noise=False, testing=False): self.ignore_noise = ignore_noise ##### model options self.old_lr = opt.lr opt.use_sigmoid = opt.no_lsgan self.opt = opt ##### define all networks we need here self.netG_A_B = networks.define_stochastic_G(nlatent=opt.nlatent, input_nc=opt.input_nc, output_nc=opt.output_nc, ngf=opt.ngf, which_model_netG=opt.which_model_netG, norm=opt.norm, use_dropout=opt.use_dropout, gpu_ids=opt.gpu_ids) self.netG_B_A = networks.define_G(input_nc=opt.output_nc, output_nc=opt.input_nc, ngf=opt.ngf, which_model_netG=opt.which_model_netG, norm=opt.norm, use_dropout=opt.use_dropout, gpu_ids=opt.gpu_ids) self.netD_A = networks.define_D_A(input_nc=opt.input_nc, ndf=32, which_model_netD=opt.which_model_netD, norm=opt.norm, use_sigmoid=opt.use_sigmoid, gpu_ids=opt.gpu_ids) self.netD_B = networks.define_D_B(input_nc=opt.output_nc, ndf=opt.ndf, which_model_netD=opt.which_model_netD, norm=opt.norm, use_sigmoid=opt.use_sigmoid, gpu_ids=opt.gpu_ids) ##### define all optimizers here self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A_B.parameters(), self.netG_B_A.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/5., betas=(opt.beta1, 0.999)) self.criterionGAN = functools.partial(criterion_GAN, use_sigmoid=opt.use_sigmoid) self.criterionCycle = F.l1_loss if not testing: with open("%s/nets.txt" % opt.expr_dir, 'w') as nets_f: networks.print_network(self.netG_A_B, nets_f) networks.print_network(self.netG_B_A, nets_f) networks.print_network(self.netD_A, nets_f) networks.print_network(self.netD_B, nets_f)
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 __init__(self, base_model='unet_256', gpu_ids=[0]): self.gpu_ids = gpu_ids self.device = torch.device('cuda:{}'.format( self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') print("Using device", self.device) self.netG = networks.define_G(3, 3, 64, netG=base_model, gpu_ids=self.gpu_ids) print("Network Initialized") self.opt = { 'crop_size': 256, 'no_flip': True, 'load_size': 256, 'preprocess': 'resize_and_crop' }
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('-----------------------------------------------')
torch.manual_seed(opt.seed) if opt.cuda: torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') train_set = get_dataset(path=opt.dataset) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batch_size, shuffle=True) device = torch.device("cuda:0" if opt.cuda else "cpu") print('===> Building models') net_g = define_G(opt.input_nc, opt.output_nc, opt.ngf, 'batch', False, 'normal', 0.02, gpu_id=device) net_d = define_D(opt.input_nc + opt.output_nc, opt.ndf, 'basic', gpu_id=device) criterionGAN = GANLoss().to(device) criterionL1 = nn.L1Loss().to(device) criterionMSE = nn.MSELoss().to(device) # setup optimizer optimizer_g = optim.Adam(net_g.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizer_d = optim.Adam(net_d.parameters(), lr=opt.lr,
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 __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
ngpu = int(opt.ngpu) nz = int(opt.nz) ngf = int(opt.ngf) ndf = int(opt.ndf) nc = 1 print('===> Loading datasets') dataset = DatasetFromFolder(opt.dataroot, opt.imageSize) assert dataset dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize, shuffle=True, num_workers=int(opt.workers)) print('===> Building model') netG = define_G(nc, nz, ngf, ngpu, device, opt.netG) netD = define_D(nc, ndf, ngpu, device, opt.netD) criterion = nn.BCELoss() fixed_noise = torch.randn(opt.batchSize, nz, 1, 1, device=device) real_label = 1 fake_label = 0 # setup optimizer optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) print('---------- Networks initialized -------------') print_network(netG) print_network(netD)
torch.cuda.manual_seed(opt.seed) print('===> Loading datasets') root_path = "dataset/" train_set = get_training_set(root_path + opt.dataset) test_set = get_test_set(root_path + opt.dataset) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=False) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) max_dataset_num = 1500 #trainsetの要素削減 train_set.image_filenames = train_set.image_filenames[:max_dataset_num] test_set.image_filenames = test_set.image_filenames[:max_dataset_num] print('===> Building model') netG = define_G(4, 3, opt.ngf, 'batch', False, [0]) #NetDを3つ構築するのがよい #netD = define_D(opt.input_nc + opt.output_nc, opt.ndf, 'batch', False, [0]) #そもそもいくつが入力なのか disc_input_nc = 6 disc_outpuc_nc = 1024 netD_Global = define_D_Global(disc_input_nc , disc_outpuc_nc, opt.ndf, [0]) netD_Local = define_D_Local(disc_input_nc , disc_outpuc_nc, opt.ndf, [0]) netD_Edge = define_D_Edge(disc_input_nc , disc_outpuc_nc, opt.ndf, [0]) #netD_Global = define_D_Global(hogehoge) #netD_Local = define_D_Local(hogehoge) #netD_Edge = define_D_Edge(hogehoge) criterionGAN = GANLoss() criterionL1 = nn.L1Loss()
print('===> Loading datasets') root_path = "dataset/" train_set = get_training_set(root_path + opt.dataset) test_set = get_test_set(root_path + opt.dataset) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model') netG = define_G(opt.input_nc, opt.output_nc, opt.ngf, 'batch', False, [0]) netD = define_D(opt.input_nc + opt.output_nc, opt.ndf, 'batch', False, [0]) print('loading done') criterionGAN = GANLoss() criterionL1 = nn.L1Loss() criterionMSE = nn.MSELoss() # setup optimizer optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) print('---------- Networks initialized -------------') print_network(netG)
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))
from networks import define_G import torch import numpy as np from torch.autograd import Variable if __name__ == '__main__': net = define_G(input_nc=3, output_nc=3, ngf=64, which_model_netG='unet_80', norm='batch') xfake = np.random.normal(0, 1, size=(2, 3, 80, 80)) xfake = Variable(torch.from_numpy(xfake).float()) dat = net(xfake) import pdb pdb.set_trace()
def main(name_exp, segloss=False, cuda=True, finetune=False): # Training settings parser = argparse.ArgumentParser(description='pix2pix-PyTorch-implementation') parser.add_argument('--batchSize', type=int, default=8, help='training batch size') parser.add_argument('--testBatchSize', type=int, default=8, help='testing batch size') parser.add_argument('--nEpochs', type=int, default=100, help='number of epochs to train for') parser.add_argument('--input_nc', type=int, default=3, help='input image channels') parser.add_argument('--output_nc', type=int, default=3, help='output image channels') parser.add_argument('--ngf', type=int, default=64, help='generator filt+ers in first conv layer') parser.add_argument('--ndf', type=int, default=64, help='discriminator filters in first conv layer') parser.add_argument('--lr', type=float, default=0.0002, help='Learning Rate. Default=0.0002') parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5') parser.add_argument('--threads', type=int, default=8, help='number of threads for data loader to use') parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123') parser.add_argument('--lamb', type=int, default=10, help='weight on L1 term in objective') opt = parser.parse_args() cudnn.benchmark = True def val(): net_current = "path_exp/checkpoint/DFS/{}/netG_model_current.pth".format(name_exp) netVal = torch.load(net_current) netVal.eval() SEG_NET.eval() features.eval() with torch.no_grad(): total_mse = 0 total_mse2 = 0 avg_psnr_depth = 0 avg_psnr_dehaze = 0 avg_ssim_depth = 0 avg_ssim_dehaze = 0 for batch in validation_data_loader: input, target, depth = Variable(batch[0]), Variable(batch[1]), Variable(batch[2]) if cuda == True: input = input.cuda() target = target.cuda() depth = depth.cuda() dehaze = netVal(input) prediction = SEG_NET(dehaze) avg_ssim_dehaze += pytorch_ssim.ssim(dehaze, target).item() mse = criterionMSE(prediction, depth) total_mse += mse.item() avg_psnr_depth += 10 * log10(1 / mse.item()) mse2 = criterionMSE(dehaze, target) total_mse2 += mse2.item() avg_psnr_dehaze += 10 * log10(1 / mse2.item()) avg_ssim_depth += pytorch_ssim.ssim(prediction, depth).item() visual_ret_val = OrderedDict() visual_ret_val['Haze'] = input visual_ret_val['Seg estimate'] = prediction visual_ret_val['Dehaze '] = dehaze visual_ret_val['GT dehaze'] = target visual_ret_val['GT Seg '] = depth visualizer.display_current_results(visual_ret_val, epoch, True) print("===> Validation") #f.write("===> Testing: \r\n") print("===> PSNR seg: {:.4f} ".format(avg_psnr_depth / len(validation_data_loader))) #f.write("===> PSNR depth: {:.4f} \r\n".format(avg_psnr_depth / len(validation_data_loader))) print("===> Mse seg: {:.4f} ".format(total_mse / len(validation_data_loader))) #f.write("===> Mse depth: {:.4f} \r\n".format(total_mse / len(validation_data_loader))) print("===> SSIM seg: {:.4f} ".format(avg_ssim_depth / len(validation_data_loader))) #f.write("===> SSIM depth: {:.4f} \r\n".format(avg_ssim_depth / len(validation_data_loader))) return total_mse / len(validation_data_loader) def testing(): path = "path_exp/checkpoint/DFS/{}/netG_model_best.pth".format(name_exp) net = torch.load(path) net.eval() SEG_NET.eval() features.eval() with torch.no_grad(): total_mse = 0 total_mse2 = 0 avg_psnr_depth = 0 avg_psnr_dehaze = 0 avg_ssim_depth = 0 avg_ssim_dehaze = 0 for batch in testing_data_loader: input, target, depth = Variable(batch[0]), Variable(batch[1]), Variable(batch[2]) if cuda == True: input = input.cuda() target = target.cuda() depth = depth.cuda() dehaze = net(input) prediction = SEG_NET(dehaze) avg_ssim_dehaze += pytorch_ssim.ssim(dehaze, target).item() mse = criterionMSE(prediction, depth) total_mse += mse.item() avg_psnr_depth += 10 * log10(1 / mse.item()) mse2 = criterionMSE(dehaze, target) total_mse2 += mse2.item() avg_psnr_dehaze += 10 * log10(1 / mse2.item()) avg_ssim_depth += pytorch_ssim.ssim(prediction, depth).item() print("===> Testing") print("===> PSNR seg: {:.4f} ".format(avg_psnr_depth / len(testing_data_loader))) print("===> Mse seg: {:.4f} ".format(total_mse / len(testing_data_loader))) print("===> SSIM seg: {:.4f} ".format(avg_ssim_depth / len(testing_data_loader))) print("===> PSNR dehaze: {:.4f} ".format(avg_psnr_dehaze / len(testing_data_loader))) print("===> SSIM dehaze: {:.4f} ".format(avg_ssim_dehaze / len(testing_data_loader))) def checkpoint(): if not os.path.exists("checkpoint"): os.mkdir("checkpoint") if not os.path.exists(os.path.join("path_exp/checkpoint/DFS", name_exp)): os.mkdir(os.path.join("path_exp/checkpoint/DFS", name_exp)) net_g_model_out_path = "path_exp/checkpoint/DFS/{}/netG_model_best.pth".format(name_exp) net_d_model_out_path = "path_exp/checkpoint/DFS/{}/netD_model_best.pth".format(name_exp) torch.save(netG, net_g_model_out_path) torch.save(netD, net_d_model_out_path) def checkpoint_current(): if not os.path.exists(os.path.join("path_exp/checkpoint/DFS", name_exp)): os.mkdir(os.path.join("path_exp/checkpoint/DFS", name_exp)) net_g_model_out_path = "path_exp/checkpoint/DFS/{}/netG_model_current.pth".format(name_exp) torch.save(netG, net_g_model_out_path) def checkpoint_seg(): if not os.path.exists(os.path.join("path_exp/checkpoint/DFS", name_exp)): os.mkdir(os.path.join("path_exp/checkpoint/DFS", name_exp)) net_g_model_out_path = "path_exp/checkpoint/DFS/{}/seg_net.pth".format(name_exp) torch.save(SEG_NET, net_g_model_out_path) torch.manual_seed(opt.seed) if cuda==True: torch.cuda.manual_seed(opt.seed) print(" ") print(name_exp) print(" ") print('===> Loading datasets') train_set = get_training_set('path_exp/cityscape/HAZE') val_set = get_val_set('path_exp/cityscape/HAZE') test_set = get_test_set('path_exp/cityscape/HAZE') training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) validation_data_loader = DataLoader(dataset=val_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader= DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) print('===> Building model') netG = define_G(opt.input_nc, opt.output_nc, opt.ngf, 'batch', False, [0]) netD = define_D(opt.input_nc + opt.output_nc, opt.ndf, 'batch', False, [0]) criterionGAN = GANLoss() criterionL1 = nn.L1Loss() criterionMSE = nn.MSELoss() # setup optimizer optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) print('---------- Networks initialized -------------') print_network(netG) print_network(netD) print('-----------------------------------------------') real_a = torch.FloatTensor(opt.batchSize, opt.input_nc, 256, 256) real_b = torch.FloatTensor(opt.batchSize, opt.output_nc, 256, 256) real_c = torch.FloatTensor(opt.batchSize, opt.output_nc, 256, 256) if cuda==True: netD = netD.cuda() netG = netG.cuda() criterionGAN = criterionGAN.cuda() criterionL1 = criterionL1.cuda() criterionMSE = criterionMSE.cuda() real_a = real_a.cuda() real_b = real_b.cuda() real_c=real_c.cuda() real_a = Variable(real_a) real_b = Variable(real_b) real_c = Variable(real_c) SEG_NET = torch.load("path_exp/SEG_NET.pth") optimizerSeg = optim.Adam(SEG_NET.parameters(), lr=opt.lr/10, betas=(opt.beta1, 0.999)) features = Vgg16() if cuda==True: SEG_NET.cuda() features.cuda() bon =100000000 for epoch in range(opt.nEpochs): features.eval() if finetune== True and epoch>50: SEG_NET.train() else: SEG_NET.eval() loss_epoch_gen=0 loss_epoch_dis=0 total_segloss=0 loss_seg=0 i=0 for iteration, batch in enumerate(training_data_loader, 1): netG.train() i=i+1 # forward real_a_cpu, real_b_cpu, real_c_cpu = batch[0], batch[1], batch[2] with torch.no_grad(): real_a = real_a.resize_(real_a_cpu.size()).copy_(real_a_cpu) with torch.no_grad(): real_b = real_b.resize_(real_b_cpu.size()).copy_(real_b_cpu) with torch.no_grad(): real_c = real_c.resize_(real_c_cpu.size()).copy_(real_c_cpu) fake_b = netG(real_a) ############################ # (1) Update D network: maximize log(D(x,y)) + log(1 - D(x,G(x))) ########################### optimizerD.zero_grad() # train with fake fake_ab = torch.cat((real_a, fake_b), 1) pred_fake = netD.forward(fake_ab.detach()) loss_d_fake = criterionGAN(pred_fake, False) # train with real real_ab = torch.cat((real_a, real_b), 1) pred_real = netD.forward(real_ab) loss_d_real = criterionGAN(pred_real, True) # Combined loss loss_d = (loss_d_fake + loss_d_real) * 0.5 loss_d.backward() optimizerD.step() ############################ # (2) Update G network: maximize log(D(x,G(x))) + L1(y,G(x)) ########################## optimizerG.zero_grad() # First, G(A) should fake the discriminator fake_ab = torch.cat((real_a, fake_b), 1) pred_fake = netD.forward(fake_ab) loss_g_gan = criterionGAN(pred_fake, True) # Second, G(A) = B loss_g_l1 = criterionL1(fake_b, real_b) * opt.lamb features_y = features(fake_b) features_x = features(real_b) loss_content = criterionMSE(features_y[1], features_x[1])*10 if segloss == True: fake_seg = SEG_NET(fake_b) loss_seg = criterionMSE(fake_seg, real_c) * 10 total_segloss += loss_seg.item() features_y = features(fake_seg) features_x = features(real_c) ssim_seg = criterionMSE(features_y[1], features_x[1]) * 10 loss_g = loss_g_gan + loss_g_l1 + loss_content + loss_seg else: loss_g = loss_g_gan + loss_g_l1+loss_content loss_epoch_gen+=loss_g.item() loss_epoch_dis+=loss_d.item() if finetune== True and epoch>50: loss_g.backward(retain_graph=True) optimizerG.step() loss_seg=loss_seg loss_seg.backward() optimizerSeg.zero_grad() optimizerSeg.step() else: loss_g.backward() optimizerG.step() errors_ret = OrderedDict() errors_ret['Total_G'] = float(loss_g) errors_ret['Content'] = float(loss_content) errors_ret['GAN'] = float(loss_g_gan) errors_ret['L1'] = float(loss_g_l1) errors_ret['D'] = float(loss_d) if i % 10 == 0: # print training losses and save logging information to the disk if i > 0: visualizer.plot_current_losses(epoch, i/(len(training_data_loader)*opt.batchSize), errors_ret) print("===> Epoch[{}]: Loss_D: {:.4f} Loss_G: {:.4f} Loss Seg: {:.4f} ".format(epoch, loss_epoch_dis,loss_epoch_gen, total_segloss)) checkpoint_current() MSE=val() if MSE < bon: bon = MSE checkpoint() checkpoint_seg() print("BEST EPOCH SAVED") testing()
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) }
['epoch', 'loss', 'acc', 'lr']) train_batch_logger = Logger( os.path.join(opt.result_path, 'train_batch.log'), ['epoch', 'batch', 'iter', 'loss', 'acc', 'lr']) if opt.nesterov: dampening = 0 else: dampening = opt.dampening optimizer = optim.Adam(model.parameters(), lr=opt.learning_rate, betas=(0.9, 0.999)) if opt.two_step: optimizer = None netG = define_G(3, 3, 64, 'batch', False) if opt.end: optimizerG = optim.Adam(list(model.parameters()) + list(netG.parameters()), opt.learning_rate, betas=(0.9, 0.999)) else: optimizerG = optim.Adam(netG.parameters(), opt.learning_rate, betas=(0.9, 0.999)) print(netG) if opt.use_gan: netD = define_D(3 + 3, 64, 'batch', use_sigmoid=False) optimizerD = optim.Adam(netD.parameters(), opt.learning_rate,
def image_loader(image_name): image = cv2.imread(image_name) image = cv2.copyMakeBorder(image, 16, 16, 16, 16, cv2.BORDER_REPLICATE) image = cv2.resize(image, (256, 256)) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = torch.from_numpy(image.transpose((2, 0, 1))) image = image.float().div(255) image = image.mul_(2).add_(-1) return image.unsqueeze(0) def imsave(tensor, title): tensor = tensor[0, :, :, :].data.cpu().numpy() tensor = tensor.transpose((1, 2, 0)) tensor = cv2.resize(tensor, (288, 288)) image = tensor[16:-16, 16:-16, :] / 2 + 0.5 cv2.imwrite('%s.jpg' % (title), image * 255) netSeg = networks.define_G(6, 3, ngf, 'unet_256').to(device) netSeg.load_state_dict( torch.load('Segmentation/netSeg.pth', map_location=lambda storage, loc: storage)) input_content = image_loader(opt.img_content).to(device) input_style = image_loader(opt.img).to(device) result = netSeg(torch.cat([input_content, input_style], 1)) imsave(result, title='temp/mask_ori')
parser.add_argument( '--glyph_path', help='path to the corresponding glyph of the input style image') parser.add_argument('--content_path', help='path to the target content image') parser.add_argument('--save_name', default='name to save') plt.switch_backend('agg') opt = parser.parse_args() opt.cuda = (opt.gpu != -1) cudnn.benchmark = True device = torch.device("cuda:%d" % (opt.gpu) if opt.cuda else "cpu") ############### Model #################### netG = networks.define_G(9, 3).to(device) netG.load_state_dict( torch.load('cache/%s_netG.pth' % (opt.style_name), map_location=lambda storage, loc: storage)) netG.eval() for p in netG.parameters(): p.requires_grad = False ############### Processing #################### def image_loader(image_name): image = cv2.imread(image_name) image = cv2.resize(image, (256, 256)) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = torch.from_numpy(image.transpose((2, 0, 1))) image = image.float().div(255)
train_set = get_training_set() test_set = get_test_set() training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) train_logger = Logger(opt.nEpochs, len(training_data_loader), opt.date) test_logger = Logger(opt.nEpochs, len(testing_data_loader), opt.date) print('===> Building model') netG = define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.batch_mode, False, [0]) netD = define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.batch_mode, False, [0]) criterionGAN = GANLoss() criterionL1 = nn.L1Loss() criterionMSE = nn.MSELoss() # setup optimizer optimizerG = optim.Adam(netG.parameters(), lr=opt.glr, betas=(opt.beta1, 0.999)) optimizerD = optim.Adam(netD.parameters(), lr=opt.dlr, betas=(opt.beta1, 0.999))
root_path = './dataset/' train_set = get_training_set(root_path) test_set = get_test_set(root_path) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) test_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False) cprint('==> Preparing Data Set: Complete\n', 'green') ################################################################################ cprint('==> Building Models', 'yellow') netG = define_G(opt.input_nc, opt.output_nc, opt.ngf, norm='batch', use_dropout=False, gpu_ids=gpu_ids) netD = define_D(opt.input_nc + opt.output_nc, opt.ndf, norm='batch', use_sigmoid=False, gpu_ids=gpu_ids) print('---------- Networks initialized -------------') print_network(netG) print_network(netD) print('-----------------------------------------------\n') cprint('==> Building Models: Complete\n', 'green') ################################################################################ criterionGAN = GANLoss() criterionL1 = nn.L1Loss() criterionMSE = nn.MSELoss() # setup optimizer optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
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))