def build_model_and_trainer(config, data_loader): if config.model.type == 'classifier': model_builder = Classifier(config) model, parallel_model = WithLoadWeights(model_builder, model_name='classifier') \ .build_model(model_name='classifier') trainer = ClassifierTrainer(model, parallel_model, data_loader, config) return model, trainer elif config.model.type == 'dcgan': g_model_builder = Generator(config) d_model_builder = Discriminator(config) # c_model_builder = Classifier(config) g = g_model_builder.define_model('generator') d, parallel_d = d_model_builder.build_model('discriminator') # c, _ = c_model_builder.build_model('classifier') # Load weights to classifier # checkpoint_path = './experiments/classifier_mnist/checkpoints/0050-classifier.hdf5' # if os.path.exists(checkpoint_path): # c.load_weights(checkpoint_path) combined_model_builder = GANCombined(config) combined, parallel_combined = WithLoadWeights(combined_model_builder, model_name='combined') \ .build_model(g=g, d=d, model_name='combined') # .build_model(g=g, d=d, c=c, model_name='combined') # trainer = GANTrainer(data_loader, config, g, d, parallel_d, c, combined, parallel_combined) trainer = GANTrainer(data_loader, config, g, d, parallel_d, combined, parallel_combined) return combined, trainer
class GAN(nn.Module): def __init__(self, num_fakes=16, rand_dim=128, size=(32, 32), use_gpu=False): super(GAN, self).__init__() self.generator = Generator(num_fakes, rand_dim, size, use_gpu) self.discriminator = Discriminator((1, ) + size) def forward(self, x=None, mode='D'): if mode == 'D': with torch.no_grad(): fake_imgs = self.generator() self.discriminator.freeze(False) fake_imgs = fake_imgs.detach() imgs = torch.cat([x, fake_imgs], dim=0) out = self.discriminator(imgs) return out else: self.discriminator.freeze(True) fake_imgs = self.generator() out = self.discriminator(fake_imgs) return fake_imgs, out
def init_models(self): self.G_X = Generator(3, 3, nn.InstanceNorm2d) self.D_X = Discriminator(3) self.G_Y = Generator(3, 3, nn.InstanceNorm2d) self.D_Y = Discriminator(3) if self.load_models: self.G_X.load_state_dict( torch.load(self.output_path + "models/G_X", map_location='cpu')) self.G_Y.load_state_dict( torch.load(self.output_path + "models/G_Y", map_location='cpu')) self.D_X.load_state_dict( torch.load(self.output_path + "models/D_X", map_location='cpu')) self.D_Y.load_state_dict( torch.load(self.output_path + "models/D_Y", map_location='cpu')) else: self.G_X.apply(init_func) self.G_Y.apply(init_func) self.D_X.apply(init_func) self.D_Y.apply(init_func) self.G_X.to(self.device) self.G_Y.to(self.device) self.D_X.to(self.device) self.D_Y.to(self.device)
def make_model(args, dataset): N = args.gpu_num class Model(): pass model = Model() netG = Generator(args) netD = Discriminator(args) netRNN = DS_SJE(args, dataset.embed_mat) if N == 0: setattr(model, "netG_0", netG) setattr(model, "netD_0", netD) setattr(model, "netRNN_0", netRNN) else: for i in range(N): temp_netG = netG.copy() temp_netD = netD.copy() temp_netRNN = netRNN.copy() temp_netG.gpu_id = i temp_netD.gpu_id = i temp_netRNN.gpu_id = i setattr(model, f"netG_{i}", temp_netG.to_gpu(i)) setattr(model, f"netD_{i}", temp_netD.to_gpu(i)) setattr(model, f"netRNN_{i}", temp_netRNN.to_gpu(i)) return model
def __init__(self, num_fakes=16, rand_dim=128, size=(32, 32), use_gpu=False): super(GAN, self).__init__() self.generator = Generator(num_fakes, rand_dim, size, use_gpu) self.discriminator = Discriminator((1, ) + size)
def __init__(self): self.env = gym.make('CartPole-v1') self.expert_traj_fpath = os.path.join('GAIL', 'expert_trajectories', 'expert_traj_test.npz') self.save_policy_fpath = os.path.join('GAIL','gail_actor1.pt') self.save_rewards_fig = os.path.join('GAIL', 'gail_rewards.png') # self.state_space = 4 self.action_space = 2 # These values are taken from the env's state and action sizes self.actor_net = ActorNet(self.state_space, self.action_space) self.actor_net.to(device=Device) self.actor_optim = torch.optim.Adam(self.actor_net.parameters(), lr = 0.0001) self.critic_net = CriticNet(self.state_space) self.critic_net.to(Device) self.critic_net_optim = torch.optim.Adam(self.critic_net.parameters(), lr = 0.0001) self.discriminator = Discriminator(self.state_space, self.action_space) self.discriminator.to(Device) self.discriminator_optim = torch.optim.Adam(self.discriminator.parameters(), lr = 0.0001) # Storing all the values used to calculate the model losses self.traj_obs = [] self.traj_actions = [] self.traj_rewards = [] self.traj_dones = [] self.traj_logprobs = [] self.traj_logits = [] self.traj_state_values = [] # Discount factor self.gamma = 0.95 # Bias Variance tradeoff (higher value results in high variance, low bias) self.gae_lambda = 0.95 # These two will be used during the training of the policy self.ppo_batch_size = 500 self.ppo_epochs = 12 self.ppo_eps = 0.2 # Discriminator self.num_expert_transitions = 150 # These will be used for the agent to play using the current policy self.max_eps = 5 # Max steps in mountaincar ex is 200 self.max_steps = 800 # documenting the stats self.avg_over = 5 # episodes self.stats = {'episode': 0, 'ep_rew': []}
def build_model(self): self.G = Generator(self.noise_n, self.G_last_act) self.D = Discriminator(self.D_out_n, self.D_last_act) self.G_optimizer = torch.optim.Adam(self.G.parameters(), self.G_lr, [self.beta1, self.beta2]) self.D_optimizer = torch.optim.Adam(self.D.parameters(), self.D_lr, [self.beta1, self.beta2]) if torch.cuda.is_available(): self.G.cuda() self.D.cuda()
def build_discriminator_trainer(self): """ Build a Keras model for training image and mask discriminators. """ # Mask Discriminator D_Mask = Discriminator(self.conf.input_shape, output='2D', downsample_blocks=3, name='D_M') D_Mask.build() self.MaskDiscriminator = D_Mask.model real_M = Input(self.conf.input_shape) fake_M = Input(self.conf.input_shape) dis_real_M = self.MaskDiscriminator(real_M) dis_fake_M = self.MaskDiscriminator(fake_M) D_Image = Discriminator(self.conf.input_shape, output='2D', downsample_blocks=3, name='D_X') D_Image.build() self.ImageDiscriminator = D_Image.model real_X = Input(self.conf.input_shape) fake_X = Input(self.conf.input_shape) dis_real_X = self.ImageDiscriminator(real_X) dis_fake_X = self.ImageDiscriminator(fake_X) self.D_model = Model( inputs=[real_M, fake_M, real_X, fake_X], outputs=[dis_real_M, dis_fake_M, dis_real_X, dis_fake_X]) self.D_model.compile(Adam(lr=0.0001, beta_1=0.5), loss='mse') log.info('Discriminators Trainer') self.D_model.summary(print_fn=log.info)
def get_modules(opt): modules = {} disc = Discriminator() gen = Generator() clf = Classifier() if opt.cuda: disc = disc.cuda() gen = gen.cuda() clf = clf.cuda() modules['Discriminator'] = disc modules['Generator'] = gen modules['Classifier'] = clf return modules
def test_adapt(): args = parse() print(args) dataset = SepeDataset(args.poses_train,args.images_train,coor_layer_flag =False) dataloader = DataLoader(dataset, batch_size=1,shuffle=True ,num_workers=1,drop_last=True,worker_init_fn=lambda wid:np.random.seed(np.uint32(torch.initial_seed() + wid))) dataset_tgt = SepeDataset(args.poses_target,args.images_target,coor_layer_flag =False) dataloader_tgt = DataLoader(dataset_tgt, batch_size=1,shuffle=True ,num_workers=1,drop_last=True,worker_init_fn=lambda wid:np.random.seed(np.uint32(torch.initial_seed() + wid))) src_extractor = DVOFeature() tgt_extractor = DVOFeature() src_extractor.load_state_dict(torch.load(args.feature_model)) tgt_extractor.load_state_dict(torch.load(args.feature_model)) dvo_discriminator = Discriminator(500,500,2) adapt(src_extractor,tgt_extractor,dvo_discriminator,dataloader,dataloader_tgt,args) torch.save(tgt_extractor.state_dict(),'tgt_feature_'+args.tag+str(args.epoch)+'.pt') torch.save(dvo_discriminator.state_dict(),'dis_'+args.tag+str(args.epoch)+'.pt')
def get_energy(args, input_dims, device): if args.discriminator == 'convolutional': return Discriminator(nn_type=args.d_model, bn=args.bn, no_trunc=args.no_trunc, skipinit=args.skipinit).to(device) elif args.discriminator == 'nvp': return energy_model.NVP([input_dims], device, args.num_blocks, mode='discriminator', with_bn=args.dis_bn).to(device) elif args.discriminator == 'made': return energy_model.MADEGenerator([input_dims], mode='discriminator').to(device) elif args.discriminator == 'maf': return energy_model.FlowGenerator([input_dims], device, args.num_blocks, 'maf', mode='discriminator', with_bn=args.dis_bn).to(device) elif args.discriminator == 'mogmaf': return energy_model.FlowGenerator([input_dims], device, args.num_blocks, 'mogmaf', mode='discriminator', with_bn=args.dis_bn).to(device) elif args.discriminator == 'toy': return tm.Discriminator(3) elif args.discriminator == 'are': return energy_model.Discriminator(input_dims, device).to(device) elif args.discriminator == 'are4': return energy_model.Discriminator4(input_dims, device).to(device)
def __init__(self, args): self.args = args # To write to Tensorboard self.writer = SummaryWriter() # Holds all data classes needed self.data = GANData(args, root='./data') # Instantiate models and load to device self.D = Discriminator(784, args.d_hidden_size).to(args.device) self.G = Generator(784, args.latent_size, args.g_hidden_size).to(args.device) # Instantiate criterion used for both D and G self.criterion = nn.BCELoss() # Instantiate an optimizer for both D and G self.d_optim = optim.Adam(self.D.parameters(), lr=args.lr) self.g_optim = optim.Adam(self.G.parameters(), lr=args.lr)
def build_mask_discriminator(self): # Build a discriminator for masks. D = Discriminator(self.conf.d_mask_params) D.build() log.info('Mask Discriminator D_M') D.model.summary(print_fn=log.info) self.D_Mask = D.model real_M = Input(self.conf.d_mask_params.input_shape) fake_M = Input(self.conf.d_mask_params.input_shape) real = self.D_Mask(real_M) fake = self.D_Mask(fake_M) self.D_Mask_trainer = Model([real_M, fake_M], [real, fake], name='D_Mask_trainer') self.D_Mask_trainer.compile(Adam(lr=self.conf.d_mask_params.lr), loss='mse') self.D_Mask_trainer.summary(print_fn=log.info)
def build_model(args): args.to_train = 'CDG' networks = {} opts = {} if 'C' in args.to_train: networks['C'] = GuidingNet(args.img_size, {'cont': args.sty_dim, 'disc': args.output_k}) networks['C_EMA'] = GuidingNet(args.img_size, {'cont': args.sty_dim, 'disc': args.output_k}) if 'D' in args.to_train: networks['D'] = Discriminator(args.img_size, num_domains=args.output_k) if 'G' in args.to_train: networks['G'] = Generator(args.img_size, args.sty_dim, use_sn=False) networks['G_EMA'] = Generator(args.img_size, args.sty_dim, use_sn=False) if args.distributed: if args.gpu is not None: print('Distributed to', args.gpu) torch.cuda.set_device(args.gpu) args.batch_size = int(args.batch_size / args.ngpus_per_node) args.workers = int(args.workers / args.ngpus_per_node) for name, net in networks.items(): if name in ['inceptionNet']: continue net_tmp = net.cuda(args.gpu) networks[name] = torch.nn.parallel.DistributedDataParallel(net_tmp, device_ids=[args.gpu], output_device=args.gpu) else: for name, net in networks.items(): net_tmp = net.cuda() networks[name] = torch.nn.parallel.DistributedDataParallel(net_tmp) elif args.gpu is not None: torch.cuda.set_device(args.gpu) for name, net in networks.items(): networks[name] = net.cuda(args.gpu) else: for name, net in networks.items(): networks[name] = torch.nn.DataParallel(net).cuda() if 'C' in args.to_train: opts['C'] = torch.optim.Adam( networks['C'].module.parameters() if args.distributed else networks['C'].parameters(), 1e-4, weight_decay=0.001) if args.distributed: networks['C_EMA'].module.load_state_dict(networks['C'].module.state_dict()) else: networks['C_EMA'].load_state_dict(networks['C'].state_dict()) if 'D' in args.to_train: opts['D'] = torch.optim.RMSprop( networks['D'].module.parameters() if args.distributed else networks['D'].parameters(), 1e-4, weight_decay=0.0001) if 'G' in args.to_train: opts['G'] = torch.optim.RMSprop( networks['G'].module.parameters() if args.distributed else networks['G'].parameters(), 1e-4, weight_decay=0.0001) return networks, opts
def test_model(image): dvo_feature_extractor = DVOFeature() dvo_regressor = DVORegression() dvo_discriminator = Discriminator(500,500,2) feature = dvo_feature_extractor(image) print(feature.shape) motion = dvo_regressor(feature) print(motion.shape) dis = dvo_discriminator(feature) print(dis)
def __init__(self, opt): super(Model, self).__init__() # Generator self.gen = Generator(opt).cuda(opt.gpu_id) self.gen_params = self.gen.parameters() num_params = 0 for p in self.gen.parameters(): num_params += p.numel() print(self.gen) print(num_params) # Discriminator self.dis = Discriminator(opt).cuda(opt.gpu_id) self.dis_params = self.dis.parameters() num_params = 0 for p in self.dis.parameters(): num_params += p.numel() print(self.dis) print(num_params) # Regressor if opt.mse_weight: self.reg = torch.load('data/utils/classifier.pth').cuda( opt.gpu_id).eval() else: self.reg = None # Losses self.criterion_gan = GANLoss(opt, self.dis) self.criterion_mse = lambda x, y: l1_loss(x, y) * opt.mse_weight self.loss_mse = Variable(torch.zeros(1).cuda()) self.loss_adv = Variable(torch.zeros(1).cuda()) self.loss = Variable(torch.zeros(1).cuda()) self.path = opt.experiments_dir + opt.experiment_name + '/checkpoints/' self.gpu_id = opt.gpu_id self.noise_channels = opt.in_channels - len(opt.input_idx.split(','))
def __init__(self, n_items, user_layers, lambda_mi=0.1, drop_ratio=0.4, aggregator_type='attention'): super(GroupIM, self).__init__() self.n_items = n_items self.lambda_mi = lambda_mi self.drop = nn.Dropout(drop_ratio) self.embedding_dim = user_layers[-1] self.aggregator_type = aggregator_type self.user_preference_encoder = Encoder(self.n_items, user_layers, self.embedding_dim, drop_ratio) if self.aggregator_type == 'maxpool': self.preference_aggregator = MaxPoolAggregator( self.embedding_dim, self.embedding_dim) elif self.aggregator_type == 'meanpool': self.preference_aggregator = MeanPoolAggregator( self.embedding_dim, self.embedding_dim) elif self.aggregator_type == 'attention': self.preference_aggregator = AttentionAggregator( self.embedding_dim, self.embedding_dim) else: raise NotImplementedError( "Aggregator type {} not implemented ".format( self.aggregator_type)) self.group_predictor = nn.Linear(self.embedding_dim, self.n_items, bias=False) nn.init.xavier_uniform_(self.group_predictor.weight) self.discriminator = Discriminator(embedding_dim=self.embedding_dim) for m in self.modules(): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) if isinstance(m, nn.Embedding): nn.init.xavier_uniform_(m.weight)
def __init__(self, args, training=False): # Define parameters for training self.learning_rate = args["learning_rate"] # Learning rate self.loss_function = args["loss_function"] # loss_function self.metrics = args["metrics"] # metrics self.channels = args['channels'] self.batch_size = args['batch_size'] self.num_steps = args['num_steps'] self.seq_shape = (self.num_steps, self.channels) self.noise_dim = args["noise_dim"] # metrics self.use_random_noise = args["use_random_noise"] self.training_mode = training # Following parameter and optimizer set as recommended in paper self.optimizer = Adam(lr=self.learning_rate) # Build and compile the critic self.critic = Discriminator(args, training) self.critic.model.compile(loss=self.loss_function, optimizer=self.optimizer, metrics=self.metrics) # Build the generator self.generator = Generator(args, training) # Build de combined model (generator => critic) # The generator takes latent space as input and generates fake signal as output z = Input(shape=(self.noise_dim, )) signal = self.generator.model(z) # For the combined model we will only train the generator self.critic.model.trainable = False # The critic takes generated signal as input and determines validity valid = self.critic.model(signal) # The combined model (stacked generator and critic) - try to minimize the loss self.combined = Model(z, valid) self.combined.compile(loss='binary_crossentropy', optimizer=self.optimizer)
def __init__(self, hparams): super().__init__() self.hparams = hparams # Encoder self.encoder = Encoder(ngf=self.hparams.ngf, z_dim=self.hparams.z_dim) self.encoder.apply(weights_init) device = "cuda" if isinstance(self.hparams.gpus, int) else "cpu" # Decoder self.decoder = Decoder(ngf=self.hparams.ngf, z_dim=self.hparams.z_dim) self.decoder.apply(weights_init) # Discriminator self.discriminator = Discriminator() self.discriminator.apply(weights_init) # Losses self.criterionFeat = torch.nn.L1Loss() self.criterionGAN = GANLoss(gan_mode="lsgan") if self.hparams.use_vgg: self.criterion_perceptual_style = [Perceptual_Loss(device)]
def test_test_real(): args = parse() print(args) dataset = SepeDataset(args.poses_train,args.images_train,coor_layer_flag =False) dataloader = DataLoader(dataset, batch_size=10,shuffle=False ,num_workers=1,drop_last=True) dvo_feature_extractor = DVOFeature() dvo_regressor = DVORegression() dvo_discriminator = Discriminator(500,500,2) dvo_feature_extractor.load_state_dict(torch.load(('feature_ntsd_2_10.pt'))) dvo_regressor.load_state_dict(torch.load('regressor_seed_ntsd_2_10.pt')) test(dvo_feature_extractor,dvo_regressor,dataloader,args)
def test_train_real(): args = parse() print(args) dataset = SepeDataset(args.poses_train,args.images_train,coor_layer_flag =False) dataloader = DataLoader(dataset, batch_size=3,shuffle=True ,num_workers=1,drop_last=True,worker_init_fn=lambda wid:np.random.seed(np.uint32(torch.initial_seed() + wid))) dvo_feature_extractor = DVOFeature() dvo_regressor = DVORegression() dvo_discriminator = Discriminator(500,500,2) trained_feature,trained_regressor = train(dvo_feature_extractor,dvo_regressor,dataloader,args) torch.save(trained_feature.state_dict(),'feature_'+args.tag+str(args.epoch)+'.pt') torch.save(trained_regressor.state_dict(),'regressor_'+args.tag+str(args.epoch)+'.pt')
def test_train(): args = parse() print(args) motion_ax_i = [int(i) for i in args.motion_ax.split(' ')] dataset = RandomDataset(20000,motion_ax = motion_ax_i) dataloader = DataLoader(dataset, batch_size=1000,shuffle=False ,num_workers=1,drop_last=True,worker_init_fn=lambda wid:np.random.seed(np.uint32(torch.initial_seed() + wid))) dvo_feature_extractor = DVOFeature() dvo_regressor = DVORegression() dvo_discriminator = Discriminator(500,500,2) trained_feature,trained_regressor = train(dvo_feature_extractor,dvo_regressor,dataloader,args) torch.save(trained_feature.state_dict(),'feature_seed'+args.motion_ax.replace(' ','')+str(args.epoch)+'.pt') torch.save(trained_regressor.state_dict(),'regressor_seed'+args.motion_ax.replace(' ','')+str(args.epoch)+'.pt')
def build_image_discriminator2(self): """ Build a discriminator for images """ params2 = self.conf.d_image_params params2['name'] = 'D_Image2' D = Discriminator(params2) D.build() log.info('Image Discriminator D_I2') D.model.summary(print_fn=log.info) self.D_Image2 = D.model real_x = Input(self.conf.d_image_params.input_shape) fake_x = Input(self.conf.d_image_params.input_shape) real = self.D_Image2(real_x) fake = self.D_Image2(fake_x) self.D_Image2_trainer = Model([real_x, fake_x], [real, fake], name='D_Image2_trainer') self.D_Image2_trainer.compile(Adam(lr=self.conf.d_image_params.lr), loss='mse') self.D_Image2_trainer.summary(print_fn=log.info)
def test_test(): args = parse() print(args) motion_ax_i = [int(i) for i in args.motion_ax.split(' ')] test_motion_ax_i = [int(i) for i in args.test_motion_ax.split(' ')] dataset = RandomDataset(2,motion_ax = test_motion_ax_i) dataloader = DataLoader(dataset, batch_size=1,shuffle=False ,num_workers=1,drop_last=True) dvo_feature_extractor = DVOFeature() dvo_regressor = DVORegression() dvo_discriminator = Discriminator(500,500,2) dvo_feature_extractor.load_state_dict(torch.load('feature'+args.motion_ax.replace(' ','')+'.pt')) dvo_regressor.load_state_dict(torch.load('regressor'+args.motion_ax.replace(' ','')+'.pt')) test(dvo_feature_extractor,dvo_regressor,dataloader,args)
def build_models(conf, n_voc, seq_len): model = dict() netG = Generator(n_voc, conf.emb_dim, conf.hid_dim, seq_len, conf.gpu_num) netD = Discriminator(n_voc, conf.emb_dim, conf.hid_dim, seq_len, conf.gpu_num) if conf.gpu_num == 0: model["netG_0"] = netG model["netD_0"] = netD else: for i in range(conf.gpu_num): copy_netG = netG.copy() copy_netD = netD.copy() copy_netG.gpu_id = i copy_netD.gpu_id = i model[f"netG_{i}"] = copy_netG.to_gpu(i) model[f"netD_{i}"] = copy_netD.to_gpu(i) return model
def main(): data_loader = DataLoader(data_path, batch_size) generator = Generator(vocab_size, g_emb_dim, g_hidden_dim) discriminator = Discriminator(vocab_size, d_hidden_dim) gen_optimizer = optim.Adam(generator.parameters(), lr=0.0001) disc_optimizer = optim.Adam(discriminator.parameters(), lr=0.0001) bce_criterion = nn.BCELoss() gen_criterion = nn.NLLLoss(size_average=False) if (opt.cuda): generator.cuda() pretrain_lstm(generator, data_loader, gen_optimizer, gen_criterion, 10) all_G_losses = [] all_D_losses = [] for i in tqdm(range(total_epochs)): g_losses, d_losses = train_gan_epoch(discriminator, generator, data_loader, gen_optimizer, disc_optimizer, bce_criterion) all_G_losses += g_losses all_D_losses += d_losses sample = generator.sample(batch_size, g_seq_length) print(generator) with open('./data/gumbel_softmax_gan_gen.txt', 'w') as f: for each_str in data_loader.convert_to_char(sample): f.write(each_str + '\n') gen_file_name = 'gen_gumbel_softmax_' + str(total_epochs) + '.pth' disc_file_name = 'disc_gumbel_softmax_' + str(total_epochs) + '.pth' Utils.save_checkpoints(checkpoint_dir, gen_file_name, generator) Utils.save_checkpoints(checkpoint_dir, disc_file_name, discriminator) plt.plot(list(range(len(all_G_losses))), all_G_losses, 'g-', label='gen loss') plt.plot(list(range(len(all_D_losses))), all_D_losses, 'b-', label='disc loss') plt.legend() plt.show()
def __init__(self, config): """ Construct a new GAN trainer :param Config config: The parsed network configuration. """ self.config = config LOG.info("CUDA version: {0}".format(version.cuda)) LOG.info("Creating data loader from path {0}".format(config.FILENAME)) self.data_loader = Data( config.FILENAME, config.BATCH_SIZE, polarisations=config.POLARISATIONS, # Polarisations to use frequencies=config.FREQUENCIES, # Frequencies to use max_inputs=config. MAX_SAMPLES, # Max inputs per polarisation and frequency normalise=config.NORMALISE) # Normalise inputs shape = self.data_loader.get_input_shape() width = shape[1] LOG.info("Creating models with input shape {0}".format(shape)) self._autoencoder = Autoencoder(width) self._discriminator = Discriminator(width) # TODO: Get correct input and output widths for generator self._generator = Generator(width, width) if config.USE_CUDA: LOG.info("Using CUDA") self.autoencoder = self._autoencoder.cuda() self.discriminator = self._discriminator.cuda() self.generator = self._generator.cuda() else: LOG.info("Using CPU") self.autoencoder = self._autoencoder self.discriminator = self._discriminator self.generator = self._generator
def init_swa_models(self): """ Initialise objects for Stochastic Weight Averaging """ self.swa_D_Mask = SWA(40, Discriminator(self.conf.d_mask_params).build, None) self.swa_D_Image1 = SWA(40, Discriminator(self.conf.d_image_params).build, None) self.swa_D_Image2 = SWA(40, Discriminator(self.conf.d_image_params).build, None) self.swa_Enc_Anatomy1 = SWA(40, anatomy_encoder.build, self.conf.anatomy_encoder) self.swa_Enc_Anatomy2 = SWA(40, anatomy_encoder.build, self.conf.anatomy_encoder) self.swa_Enc_Modality = SWA(40, modality_encoder.build, self.conf) self.swa_Anatomy_Fuser = SWA(40, anatomy_fuser.build, self.conf) self.swa_Segmentor = SWA(40, segmentor.build, self.conf) self.swa_Decoder = SWA(40, decoder.build, self.conf) self.swa_Balancer = SWA(40, balancer.build, self.conf) self.set_swa_model_weights()
def __init__(self, device, model, image_nc, box_min, box_max): output_nc = image_nc self.device = device self.model = model self.input_nc = image_nc self.output_nc = output_nc self.box_min = box_min self.box_max = box_max self.gen_input_nc = image_nc self.netG = Generator(self.gen_input_nc, image_nc).to(device) self.netDisc = Discriminator(image_nc).to(device) # initialize all weights self.netG.apply(weights_init) self.netDisc.apply(weights_init) # initialize optimizers self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=0.001) self.optimizer_D = torch.optim.Adam(self.netDisc.parameters(), lr=0.001) if not os.path.exists(models_path): os.makedirs(models_path)
def build_model(args): generator = Generator(args) discriminator = Discriminator(args) mapping_network = MappingNetwork(args) generator_ema = copy.deepcopy(generator) mapping_network_ema = copy.deepcopy(mapping_network) nets = Munch(generator=generator, discriminator=discriminator, mapping_network=mapping_network) nets_ema = Munch(generator=generator_ema, mapping_network=mapping_network_ema) return nets, nets_ema