def initialize_networks(self, opt, end2end=False, triple=False): if opt.end2endtri: netG_1 = networks.define_G(opt, triple) netD_1 = networks.define_D(opt, triple) netG = networks.define_G(opt) netD = networks.define_D(opt) if opt.isTrain else None netE = networks.define_E(opt) if opt.use_vae else None if not opt.isTrain or opt.continue_train: if opt.end2endtri: netG_1 = util.load_network(netG_1, 'G', opt.which_triple_epoch, opt, triple) netG = util.load_network(netG, 'G', opt.which_epoch, opt) if opt.isTrain: # and end2end: netD = util.load_network(netD, 'D', opt.which_epoch, opt) if opt.end2endtri: netD_1 = util.load_network(netD_1, 'D', opt.which_triple_epoch, opt, triple) if opt.use_vae: netE = util.load_network(netE, 'E', opt.which_epoch, opt) if not opt.end2endtri: netG_1 = None netD_1 = None return netG, netD, netE, netG_1, netD_1
def get_nets(path, which_epoch='latest', def_opt=None): gpu_ids = [0] Tensor = torch.cuda.FloatTensor opt = util.load_opt(path, def_opt) # assume caffe style model opt.not_caffe = False netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, gpu_ids, opt=opt) util.load_network_with_path(netG_A, 'G_A', path) netG_B = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, gpu_ids, opt=opt) util.load_network_with_path(netG_B, 'G_B', path) netG_A.cuda() netG_B.cuda() return {'A': netG_A, 'B': netG_B}
def initialize_networks(self, opt): self.netGA = networks.define_G(opt, opt['netGA']) self.netGB = networks.define_G(opt, opt['netGB']) self.netDA = networks.define_D(opt, opt['netDA']) self.netDB = networks.define_D(opt, opt['netDB']) self.netEA, self.netHairA = networks.define_RES( opt, opt['input_nc_A'], opt['netEDA']) self.netEB, self.netHairB = networks.define_RES( opt, opt['input_nc_B'], opt['netEDB']) if self.opt['pretrain']: self.train_nets = [ self.netGA, self.netGB, self.netDA, self.netDB, self.netEA, self.netHairA, self.netEB, self.netHairB ] else: self.train_nets = [self.netEA, self.netHairA] # set require gradients if self.isTrain: self.set_requires_grad(self.train_nets, True) else: self.set_requires_grad(self.train_nets, False) if self.use_gpu: for i in range(len(self.train_nets)): self.train_nets[i] = DataParallelWithCallback( self.train_nets[i], device_ids=opt['gpu_ids']) if self.opt['pretrain']: self.netGA, self.netGB, self.netDA, self.netDB, self.netEA, \ self.netHairA, self.netEB, self.netHairB = self.train_nets else: self.netEA, self.netHairA = self.train_nets
def __init__(self, opt, device): super(CycleGAN, self).__init__() self.device = device self.opt = opt self.netG_A = networks.define_G(self.opt.input_nc, self.opt.output_nc, self.opt.ngf, self.opt.netG, self.opt.norm, self.opt.dropout, self.opt.init_type, self.opt.init_gain, self.opt.task_num, self.opt.netG_A_filter_list) self.netG_B = networks.define_G(self.opt.input_nc, self.opt.output_nc, self.opt.ngf, self.opt.netG, self.opt.norm, self.opt.dropout, self.opt.init_type, self.opt.init_gain, self.opt.task_num, self.opt.netG_B_filter_list) if opt.train: self.netD_A = networks.define_D(self.opt.input_nc, self.opt.ndf, self.opt.netD, self.opt.norm, self.opt.init_type, self.opt.init_gain) self.netD_B = networks.define_D(self.opt.input_nc, self.opt.ndf, self.opt.netD, self.opt.norm, self.opt.init_type, self.opt.init_gain) self.fake_A_pool = ImageBuffer( self.opt.pool_size ) # create image buffer to store previously generated images self.fake_B_pool = ImageBuffer( self.opt.pool_size ) # create image buffer to store previously generated images self.criterionGAN = networks.GANLoss(self.opt.gan_mode).to( self.device) # define GAN loss. self.criterionCycle = torch.nn.L1Loss() self.criterionIdt = torch.nn.L1Loss() self.optimizer_G = torch.optim.Adam(itertools.chain( self.netG_A.parameters(), self.netG_B.parameters()), lr=self.opt.lr, betas=(self.opt.beta1, 0.999)) self.optimizer_D = torch.optim.Adam(itertools.chain( self.netD_A.parameters(), self.netD_B.parameters()), lr=self.opt.lr, betas=(self.opt.beta1, 0.999)) self.optimizers = [] self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D) self.schedulers = [ networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers ]
def get_stack_nets(path, which_epoch='latest'): gpu_ids = [0] Tensor = torch.cuda.FloatTensor opt = util.load_opt(path) netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.norm, False, opt.init_type, gpu_ids, n_upsampling=3, opt=opt) netG_B = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.which_model_netG, opt.norm, False, opt.init_type, gpu_ids, n_upsampling=3, opt=opt) netG_A_pre = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.norm, False, opt.init_type, gpu_ids, n_downsampling=3, opt=opt) netG_B_pre = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.which_model_netG, opt.norm, False, opt.init_type, gpu_ids, n_downsampling=3, opt=opt) load_network_with_path(netG_A, 'G_A', which_epoch, path) load_network_with_path(netG_B, 'G_B', which_epoch, path) load_network_with_path(netG_A_pre, 'G_A_pre', which_epoch, path) load_network_with_path(netG_B_pre, 'G_B_pre', which_epoch, path) netG_A.cuda() netG_B.cuda() netG_A_pre.cuda() netG_B_pre.cuda() return {'A': netG_A, 'B': netG_B, 'A_pre': netG_A_pre, 'B_pre': netG_B_pre}
def __init__(self, opt): assert opt.isTrain opt = copy.deepcopy(opt) if len(opt.gpu_ids) > 0: opt.gpu_ids = opt.gpu_ids[:1] self.gpu_ids = opt.gpu_ids super(SPADEModelModules, self).__init__() self.opt = opt self.model_names = ['G_student', 'G_teacher', 'D'] teacher_opt = self.create_option('teacher') self.netG_teacher = networks.define_G(opt.teacher_netG, gpu_ids=self.gpu_ids, opt=teacher_opt) student_opt = self.create_option('student') self.netG_student = networks.define_G(opt.student_netG, init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=student_opt) if hasattr(opt, 'distiller'): pretrained_opt = self.create_option('pretrained') self.netG_pretrained = networks.define_G(opt.pretrained_netG, gpu_ids=self.gpu_ids, opt=pretrained_opt) self.netD = networks.define_D(opt.netD, init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=opt) self.mapping_layers = ['head_0', 'G_middle_1', 'up_1'] self.netAs = nn.ModuleList() for i, mapping_layer in enumerate(self.mapping_layers): if mapping_layer != 'up_1': fs, ft = opt.student_ngf * 16, opt.teacher_ngf * 16 else: fs, ft = opt.student_ngf * 4, opt.teacher_ngf * 4 if hasattr(opt, 'distiller'): netA = nn.Conv2d(in_channels=fs, out_channels=ft, kernel_size=1) else: netA = SuperConv2d(in_channels=fs, out_channels=ft, kernel_size=1) networks.init_net(netA, opt.init_type, opt.init_gain, self.gpu_ids) self.netAs.append(netA) self.criterionGAN = GANLoss(opt.gan_mode) self.criterionFeat = nn.L1Loss() self.criterionVGG = VGGLoss() self.optimizers = [] self.netG_teacher.eval() self.config = None
def define_networks(self, start_epoch): opt = self.opt # Generator network input_nc = opt.label_nc if (opt.label_nc != 0 and not self.pose) else opt.input_nc netG_input_nc = input_nc opt.for_face = False self.netG = networks.define_G(opt) if self.refine_face: opt_face = copy.deepcopy(opt) opt_face.n_downsample_G -= 1 if opt_face.n_adaptive_layers > 0: opt_face.n_adaptive_layers -= 1 opt_face.input_nc = opt.output_nc opt_face.fineSize = self.faceRefiner.face_size opt_face.aspect_ratio = 1 opt_face.for_face = True self.netGf = networks.define_G(opt_face) # Discriminator network if self.isTrain or opt.finetune: netD_input_nc = input_nc + opt.output_nc + (1 if self.concat_fg_mask_for_D else 0) if self.concat_ref_for_D: netD_input_nc *= 2 self.netD = networks.define_D(opt, netD_input_nc, opt.ndf, opt.n_layers_D, opt.norm_D, opt.netD_subarch, opt.num_D, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids) if self.add_face_D: self.netDf = networks.define_D(opt, opt.output_nc * 2, opt.ndf, opt.n_layers_D, opt.norm_D, 'n_layers', 1, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids) else: self.netDf = None self.temporal = False self.netDT = None print('---------- Networks initialized -------------') # initialize optimizers if self.isTrain: # optimizer G params = list(self.netG.parameters()) if self.refine_face: params += list(self.netGf.parameters()) self.optimizer_G = self.get_optimizer(params, for_discriminator=False) # optimizer D params = list(self.netD.parameters()) if self.add_face_D: params += list(self.netDf.parameters()) self.optimizer_D = self.get_optimizer(params, for_discriminator=True) print('---------- Optimizers initialized -------------') # make model temporal by generating multiple frames if (not opt.isTrain or start_epoch > opt.niter_single) and opt.n_frames_G > 1: self.make_temporal_model()
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) # 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.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 = ImagePool(opt.pool_size) # create image buffer to store previously generated images self.fake_B_pool = ImagePool(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)
def initialize_networks(self, opt): netG_for_CT = networks.define_G(opt) netD_aligned = networks.define_D(opt) if opt.isTrain else None netG_for_MR = networks.define_G(opt) netD_unaligned = networks.define_D(opt) if opt.isTrain else None if not opt.isTrain or opt.continue_train: netG_for_CT = util.load_network(netG_for_CT, 'G_for_CT', opt.which_epoch, opt) netG_for_MR = util.load_network(netG_for_MR, 'G_for_MR', opt.which_epoch, opt) if opt.isTrain: netD_aligned = util.load_network(netD_aligned, 'D_aligned', opt.which_epoch, opt) netD_unaligned = util.load_network(netD_unaligned, 'D_unaligned', opt.which_epoch, opt) return netG_for_CT, netD_aligned, netG_for_MR, netD_unaligned
def __init__(self, opt): super(HyperRIMModel, self).__init__(opt) train_opt = opt['train'] # define networks and load pretrained models self.netG = networks.define_G(opt).to(self.device) if self.is_train: self.netG.train() self.load() # store the number of levels and code channel self.num_levels = int(math.log(opt['scale'], 2)) self.code_nc = opt['network_G']['code_nc'] self.map_nc = opt['network_G']['map_nc'] # define losses, optimizer and scheduler self.netF = networks.define_F(opt).to(self.device) self.projections = None if self.is_train: # G wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0 map_network_params = [] core_network_params = [] # can freeze weights for any of the levels freeze_level = train_opt['freeze_level'] for k, v in self.netG.named_parameters(): if v.requires_grad: if freeze_level: if "level_%d" % freeze_level not in k: if 'map' in k: map_network_params.append(v) else: core_network_params.append(v) else: if 'map' in k: map_network_params.append(v) else: core_network_params.append(v) else: print('WARNING: params [{:s}] will not optimize.'.format(k)) self.optimizer_G = torch.optim.Adam([{'params': core_network_params}, {'params': map_network_params, 'lr': 1e-2 * train_opt['lr_G']}], lr=train_opt['lr_G'], weight_decay=wd_G, betas=(train_opt['beta1_G'], 0.999)) self.optimizers.append(self.optimizer_G) # for resume training - load the previous optimizer stats self.load_optimizer() # schedulers if train_opt['lr_scheme'] == 'MultiStepLR': for optimizer in self.optimizers: self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, train_opt['lr_steps'], train_opt['lr_gamma'])) else: raise NotImplementedError('MultiStepLR learning rate scheme is enough.') self.log_dict = OrderedDict() print('---------- Model initialized ------------------') self.print_network() print('-----------------------------------------------')
def initialize_networks(self, opt): netG = networks.define_G(opt) # netD = networks.define_D(opt) if opt.isTrain else None if opt.isTrain: opt.label_nc = opt.label_nc-1 netD = networks.define_D(opt) else: netD_fine = None netE = networks.define_E(opt) if opt.use_vae else None if opt.isTrain: opt.label_nc = (opt.label_nc+1) netD_fine = networks.define_D(opt) else: netD_fine = None if not opt.isTrain or opt.continue_train: netG = util.load_network(netG, 'G', opt.which_epoch, opt) if opt.isTrain: netD = util.load_network(netD, 'D', opt.which_epoch, opt) netD_fine = util.load_network(netD_fine, 'D', opt.which_epoch, opt) else: netD = None netD_fine = None if opt.use_vae: netE = util.load_network(netE, 'E', opt.which_epoch, opt) return netG, netD, netE, netD_fine
def __init__(self, opt): super(SRModel, self).__init__(opt) train_opt = opt['train'] # define network and load pretrained models self.netG = networks.define_G(opt).to(self.device) self.load() if self.is_train: self.netG.train() # loss loss_type = train_opt['pixel_criterion'] if loss_type == 'l1': self.cri_pix = nn.L1Loss().to(self.device) elif loss_type == 'l2': self.cri_pix = nn.MSELoss().to(self.device) else: raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type)) self.l_pix_w = train_opt['pixel_weight'] # G feature loss if 'feature_weight' in train_opt and train_opt['feature_weight'] > 0: l_fea_type = train_opt['feature_criterion'] if l_fea_type == 'l1': self.cri_fea = nn.L1Loss().to(self.device) elif l_fea_type == 'l2': self.cri_fea = nn.MSELoss().to(self.device) else: raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_fea_type)) self.l_fea_w = train_opt['feature_weight'] else: logger.info('Remove feature loss.') self.cri_fea = None if self.cri_fea: # load VGG perceptual loss self.netF = networks.define_F(opt, use_bn=False).to(self.device) # optimizers wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0 optim_params = [] for k, v in self.netG.named_parameters(): # can optimize for a part of the model if v.requires_grad: optim_params.append(v) else: logger.warning('Params [{:s}] will not optimize.'.format(k)) self.optimizer_G = torch.optim.Adam( optim_params, lr=train_opt['lr_G'], weight_decay=wd_G) self.optimizers.append(self.optimizer_G) # schedulers if train_opt['lr_scheme'] == 'MultiStepLR': for optimizer in self.optimizers: self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \ train_opt['lr_steps'], train_opt['lr_gamma'])) else: raise NotImplementedError('MultiStepLR learning rate scheme is enough.') self.log_dict = OrderedDict() # print network self.print_network()
def main(): opt = { 'gpu_ids': [0], 'model': 'NLCSNet', 'network_G': {'which_model_G': 'NLCSGen_ResNet', 'is_train': True, 'norm_type': None, 'act_type': 'leakyrelu', 'mode': 'CNA', 'k': 6, 'in_nc': 1, 'out_nc': 1, 'nf': 16, 'nb': 5, 'upscale': 2, 'ksize_enc': 8, 'patch_size': 4, 'patch_stride': 4, 'group': 1, 'gc': 32, 'upsample_mode': 'pixelshuffle', 'fusion': 'cat', 'enc_enable': False}, 'train': {'lr_G': 0.0001, 'weight_decay_G': 0, 'beta1_G': 0.9, 'lr_scheme': 'MultiStepLR', 'lr_steps': [20000.0, 50000.0], 'lr_gamma': 0.5, 'pixel_criterion': 'l2', 'pixel_weight': 1, 'feature_criterion': 'l2', 'feature_weight': 0, 'manual_seed': 0, 'niter': 100000.0, 'val_freq': 8, 'lr_decay_iter': 10}, } opt['network_G']['enc_enable'] = True device = torch.device('cuda' if opt['gpu_ids'] is not None else 'cpu') x = torch.rand([16, 1, 128, 128], requires_grad=True).double().to(device) z = torch.rand([16, 1, 32, 32], requires_grad=True).double().to(device) # model1 = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=4, stride=4).double().to(device) # y = model1(x) # cri_pix = nn.MSELoss().to(device) # l_g_pix = cri_pix(y, z) # l_g_pix.backward() # for p in model1.parameters(): # print(p.grad) model = networks.define_G(opt).double().to(device) y = model(x) cri_pix = nn.MSELoss().to(device) l_g_pix = cri_pix(y, z) l_g_pix.backward() for p in model.parameters(): print(p.grad)
def initialize_networks(self, opt): net = {} net['netG'] = networks.define_G(opt) net['netD'] = networks.define_D(opt) if opt.isTrain else None net['netCorr'] = networks.define_Corr(opt) net['netDomainClassifier'] = networks.define_DomainClassifier( opt) if opt.weight_domainC > 0 and opt.domain_rela else None if not opt.isTrain or opt.continue_train: net['netG'] = util.load_network(net['netG'], 'G', opt.which_epoch, opt) if opt.isTrain: net['netD'] = util.load_network(net['netD'], 'D', opt.which_epoch, opt) net['netCorr'] = util.load_network(net['netCorr'], 'Corr', opt.which_epoch, opt) if opt.weight_domainC > 0 and opt.domain_rela: net['netDomainClassifier'] = util.load_network( net['netDomainClassifier'], 'DomainClassifier', opt.which_epoch, opt) if (not opt.isTrain) and opt.use_ema: net['netG'] = util.load_network(net['netG'], 'G_ema', opt.which_epoch, opt) net['netCorr'] = util.load_network(net['netCorr'], 'netCorr_ema', opt.which_epoch, opt) return net
def __init__(self, args, logger): super().__init__(args, logger) # specify the training losses you want to print out. The program will call base_model.get_current_losses self.loss_names = ['loss_G', 'loss_D'] # specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks self.model_names = ['G', 'D'] self.sample_names = ['fake_B', 'real_A', 'real_B'] # load/define networks self.G = networks.define_G(args.input_nc, args.output_nc, args.ngf, args.which_model_netG, args.norm, not args.no_dropout, args.init_type, args.init_gain, self.gpu_ids) if not 'continue_train' in args: use_sigmoid = args.no_lsgan self.D = networks.define_D(args.input_nc + args.output_nc, args.ndf, args.which_model_netD, args.n_layers_D, args.norm, use_sigmoid, args.init_type, args.init_gain, self.gpu_ids) self.fake_AB_pool = ImagePool(args.pool_size) # define loss functions self.criterionGAN = networks.GANLoss(use_lsgan=not args.no_lsgan).to(self.device) self.criterionL1 = torch.nn.L1Loss() # initialize optimizers self.optimizers = [] self.optimizer_G = torch.optim.Adam(self.G.parameters(), lr=args.g_lr, betas=(args.beta1, 0.999)) self.optimizer_D = torch.optim.Adam(self.D.parameters(), lr=args.d_lr, betas=(args.beta1, 0.999)) self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D)
def initialize_networks(self, opt): netG = networks.define_G(opt) netD = networks.define_D(opt) if opt.isTrain else None netD_rotate = networks.define_D(opt) if opt.isTrain else None netE = networks.define_E(opt) if opt.use_vae else None pretrained_path = '' if not opt.isTrain or opt.continue_train: self.load_network(netG, 'G', opt.which_epoch, pretrained_path) if opt.isTrain and not opt.noload_D: self.load_network(netD, 'D', opt.which_epoch, pretrained_path) self.load_network(netD_rotate, 'D_rotate', opt.which_epoch, pretrained_path) if opt.use_vae: self.load_network(netE, 'E', opt.which_epoch, pretrained_path) else: if opt.load_separately: netG = self.load_separately(netG, 'G', opt) if not opt.noload_D: netD = self.load_separately(netD, 'D', opt) netD_rotate = self.load_separately(netD_rotate, 'D_rotate', opt) if opt.use_vae: netE = self.load_separately(netE, 'E', opt) return netG, netD, netE, netD_rotate
def __init__(self, opt): super(TestModel, self).__init__(opt) self.visual_names = ['real_A', 'fake_B', 'real_B'] self.model_names = ['G'] self.netG = networks.define_G(opt.netG, input_nc=opt.input_nc, output_nc=opt.output_nc, ngf=opt.ngf, norm=opt.norm, dropout_rate=opt.dropout_rate, gpu_ids=self.gpu_ids, opt=opt) self.netG.eval()
def __init__(self, opt): 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 = ['G_GAN', 'G_L1', 'D_real', 'D_fake'] # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals> self.visual_names = ['real_A', 'fake_B', 'real_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', 'D'] else: # during test time, only load G self.model_names = ['G'] # define networks (both generator and discriminator) self.netG = 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) if self.isTrain: # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids) if self.isTrain: # define loss functions self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device) self.criterionL1 = torch.nn.L1Loss() # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>. self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D)
def __init__(self, opt): BaseModel.__init__(self, opt) self.netP = networks.define_P(opt, weight_path=opt.parse_net_weight) self.netG = networks.define_G(opt, use_norm='spectral_norm') if self.isTrain: self.netD = networks.define_D(opt, opt.Dinput_nc, use_norm='spectral_norm') self.vgg_model = loss.PCPFeat(weight_path='./pretrain_models/vgg19-dcbb9e9d.pth').to(opt.device) if len(opt.gpu_ids) > 0: self.vgg_model = torch.nn.DataParallel(self.vgg_model, opt.gpu_ids, output_device=opt.device) self.model_names = ['G'] self.loss_names = ['Pix', 'PCP', 'G', 'FM', 'D', 'SS'] # Generator loss, fm loss, parsing loss, discriminator loss self.visual_names = ['img_LR', 'img_HR', 'img_SR', 'ref_Parse', 'hr_mask'] self.fm_weights = [1**x for x in range(opt.D_num)] if self.isTrain: self.model_names = ['G', 'D'] self.load_model_names = ['G', 'D'] self.criterionParse = torch.nn.CrossEntropyLoss().to(opt.device) self.criterionFM = loss.FMLoss().to(opt.device) self.criterionGAN = loss.GANLoss(opt.gan_mode).to(opt.device) self.criterionPCP = loss.PCPLoss(opt) self.criterionPix= nn.L1Loss() self.criterionRS = loss.RegionStyleLoss() self.optimizer_G = optim.Adam([p for p in self.netG.parameters() if p.requires_grad], lr=opt.g_lr, betas=(opt.beta1, 0.999)) self.optimizer_D = optim.Adam([p for p in self.netD.parameters() if p.requires_grad], lr=opt.d_lr, betas=(opt.beta1, 0.999)) self.optimizers = [self.optimizer_G, self.optimizer_D]
def __init__(self, opt): super(SRModel, self).__init__(opt) train_opt = opt['train'] # define network and load pretrained models self.netG = networks.define_G(opt).to(self.device) self.load() if self.is_train: self.netG.train() # loss loss_type = train_opt['pixel_criterion'] if loss_type == 'l1': self.cri_pix = nn.L1Loss().to(self.device) elif loss_type == 'l2': self.cri_pix = nn.MSELoss().to(self.device) else: raise NotImplementedError( 'Loss type [{:s}] is not recognized.'.format(loss_type)) self.l_pix_w = train_opt['pixel_weight'] # optimizers wd_G = train_opt['weight_decay_G'] if train_opt[ 'weight_decay_G'] else 0 # find the parameters to optimize if opt['finetune_norm']: optim_params = [] for k, v in self.netG.named_parameters(): v.requires_grad = False if k.find('transformer') >= 0: v.requires_grad = True v.data.zero_() optim_params.append(v) logger.info( 'Params [{:s}] initialized to 0 and will optimize.' .format(k)) else: optim_params = list(self.netG.parameters()) self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'], weight_decay=wd_G) self.optimizers.append(self.optimizer_G) # schedulers if train_opt['lr_scheme'] == 'MultiStepLR': for optimizer in self.optimizers: self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \ train_opt['lr_steps'], train_opt['lr_gamma'])) else: raise NotImplementedError( 'MultiStepLR learning rate scheme is enough.') self.log_dict = OrderedDict() # print network self.print_network()
def __init__(self, opt): super(SRModel, self).__init__(opt) train_opt = opt['train'] self.input_L = self.Tensor() self.input_H = self.Tensor() # define network and load pretrained models self.netG = networks.define_G(opt) self.load() if self.is_train: self.netG.train() # loss loss_type = train_opt['pixel_criterion'] if loss_type == 'l1': self.cri_pix = nn.L1Loss() elif loss_type == 'l2': self.cri_pix = nn.MSELoss() else: raise NotImplementedError('Loss type [%s] is not recognized.' % loss_type) if self.use_gpu: self.cri_pix.cuda() self.l_pix_w = train_opt['pixel_weight'] # optimizers self.optimizers = [] wd_G = train_opt['weight_decay_G'] if train_opt[ 'weight_decay_G'] else 0 optim_params = [] for k, v in self.netG.named_parameters( ): # can optimize for a part of the model if v.requires_grad: optim_params.append(v) else: print('WARNING: params [%s] will not optimize.' % k) self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'], weight_decay=wd_G) self.optimizers.append(self.optimizer_G) # schedulers self.schedulers = [] if train_opt['lr_scheme'] == 'MultiStepLR': for optimizer in self.optimizers: self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \ train_opt['lr_steps'], train_opt['lr_gamma'])) else: raise NotImplementedError( 'MultiStepLR learning rate scheme is enough.') self.log_dict = OrderedDict() print('---------- Model initialized ------------------') self.print_network() print('-----------------------------------------------')
def __init__(self, opt): super(SRModel, self).__init__(opt) train_opt = opt['train'] self.chop = opt['chop'] self.scale = opt['scale'] self.val_lpips = opt['val_lpips'] # define network and load pretrained models self.netG = networks.define_G(opt).to(self.device) self.load() if self.is_train: self.netG.train() # loss loss_type = train_opt['pixel_criterion'] if loss_type == 'l1': self.cri_pix = nn.L1Loss().to(self.device) elif loss_type == 'l2': self.cri_pix = nn.MSELoss().to(self.device) else: raise NotImplementedError( 'Loss type [{:s}] is not recognized.'.format(loss_type)) self.l_pix_w = train_opt['pixel_weight'] # optimizers wd_G = train_opt['weight_decay_G'] if train_opt[ 'weight_decay_G'] else 0 optim_params = [] for k, v in self.netG.named_parameters( ): # can optimize for a part of the model if v.requires_grad: optim_params.append(v) else: logger.warning( 'Params [{:s}] will not optimize.'.format(k)) self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'], weight_decay=wd_G) self.optimizers.append(self.optimizer_G) # schedulers if train_opt['lr_scheme'] == 'MultiStepLR': for optimizer in self.optimizers: self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \ train_opt['lr_steps'], train_opt['lr_gamma'])) else: raise NotImplementedError( 'MultiStepLR learning rate scheme is enough.') self.log_dict = OrderedDict() # print network self.print_network() if self.val_lpips: self.cri_fea_lpips = PerceptualLoss(model='net-lin', net='alex').to(self.device)
def initialize_networks(self, opt): netG = networks.define_G(opt) netD = networks.define_D(opt) if opt.isTrain else None if not opt.isTrain or opt.continue_train: netG = util.load_network(netG, 'G', opt.which_epoch, opt) if opt.isTrain: netD = util.load_network(netD, 'D', opt.which_epoch, opt) return netG, netD
def main(opt): # define the generator with spectral normalization. Only the last argument counts netG = networks.define_G(opt.netG, opt=opt) util.load_network(netG, opt.restore_G_path, True) print(netG) netG.remove_spectral_norm() dirname = os.path.dirname(opt.output_path) os.makedirs(dirname, exist_ok=True) torch.save(netG.cpu().state_dict(), opt.output_path) print('Successfully export the model at [%s]!' % opt.output_path)
def initialize_networks(self, opt): net = {} net['netG'] = networks.define_G(opt) # Get SPADEGenerator net['netCorr'] = networks.define_Corr(opt) # Get NoVGGCorrespondence if not opt.isTrain or opt.continue_train: net['netG'] = util.load_network(net['netG'], 'G', opt.which_epoch, opt) net['netCorr'] = util.load_network(net['netCorr'], 'Corr', opt.which_epoch, opt) return net
def initialize_networks(self, opt): netG = None netD = None netE = None netV = None netA = None netA_sync = None if opt.train_recognition: netV = networks.define_V(opt) elif opt.train_sync: netA_sync = networks.define_A_sync(opt) if opt.use_audio else None netE = networks.define_E(opt) else: netG = networks.define_G(opt) netA = networks.define_A( opt) if opt.use_audio and opt.use_audio_id else None netA_sync = networks.define_A_sync(opt) if opt.use_audio else None netE = networks.define_E(opt) netV = networks.define_V(opt) if opt.isTrain: netD = networks.define_D(opt) if not opt.isTrain or opt.continue_train: self.load_network(netG, 'G', opt.which_epoch) self.load_network(netV, 'V', opt.which_epoch) self.load_network(netE, 'E', opt.which_epoch) if opt.use_audio: if opt.use_audio_id: self.load_network(netA, 'A', opt.which_epoch) self.load_network(netA_sync, 'A_sync', opt.which_epoch) if opt.isTrain and not opt.noload_D: self.load_network(netD, 'D', opt.which_epoch) # self.load_network(netD_rotate, 'D_rotate', opt.which_epoch, pretrained_path) else: if self.opt.pretrain: if opt.netE == 'fan': netE.load_pretrain() netV.load_pretrain() if opt.load_separately: netG = self.load_separately(netG, 'G', opt) netA = self.load_separately( netA, 'A', opt) if opt.use_audio and opt.use_audio_id else None netA_sync = self.load_separately( netA_sync, 'A_sync', opt) if opt.use_audio else None netV = self.load_separately(netV, 'V', opt) netE = self.load_separately(netE, 'E', opt) if not opt.noload_D: netD = self.load_separately(netD, 'D', opt) return netG, netD, netA, netA_sync, netV, netE
def load_weight(self, pathlist: dict): self.net_Gs = [] self.net_Ds = [] for weight in pathlist['net_G']: net_G = define_G(self.opt).to(self.device) net_G.load_state_dict(torch.load(weight, map_location=self.device)) self.net_Gs.append(net_G) for weight in pathlist['net_D']: net_D = define_D(self.opt).to(self.device) net_D.load_state_dict(torch.load(weight, map_location=self.device)) self.net_Ds.append(net_D)
def __init__(self, opt): self.opt = opt self.device = torch.device( 'cuda' if opt['gpu_ids'] is not None else 'cpu') self.is_train = opt['is_train'] self.schedulers = [] self.optimizers = [] # define network and load pretrained models self.netG = networks.define_G(opt).to(self.device) self.print_network() self.load()
def __init__(self, opt): super(SPADEModel, self).__init__(opt) self.model_names = ['G'] self.visual_names = ['labels', 'fake_B', 'real_B'] self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, opt.dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) if opt.isTrain: raise NotImplementedError("Training mode of SPADE is currently not supported!!!") else: self.netG.eval()
def __init__(self, model_path): self.model = define_G() self.transform_list = [ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5, 0.5), (0.5, 0.5, 0.5, 0.5)) ] self.transform = transforms.Compose(self.transform_list) self.model.load_state_dict(torch.load(model_path)) self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.model.eval()