def train(opt, images): gennet_g = model.Generator() gennet_f = model.Generator() vgg = model.VGG() discrinet_c = model.Discriminator(3) discrinet_t = model.Discriminator(1) g_optimizer = optimizer.Adam(params=gennet_g.parameters(), lr=opt.lr) f_optimizer = optimizer.Adam(params=gennet_f.parameters(), lr=opt.lr) c_optimizer = optimizer.Adam(params=discrinet_c.parameters(), lr=opt.lr) t_optimizer = optimizer.Adam(params=discrinet_t.parameters(), lr=opt.lr) optimizers = dict() optimizers['g'] = g_optimizer optimizers['f'] = f_optimizer optimizers['c'] = c_optimizer optimizers['t'] = t_optimizer content_criterion = nn.L1Loss() texture_criterion = model.GANLoss() color_criterion = model.GANLoss() tv_criterion = model.TVLoss(1.0) num_samples = 0 for i in xrange(opt.epoches): for j in xrange(num_samples / opt.batch_size): batch = generate_batches(images, opt.batch_size) loss = train_step(vgg=vgg, gennet_g=gennet_g, gennet_f=gennet_f, discrinet_c=discrinet_c, discrinet_t=discrinet_t, content_criterion=content_criterion, color_criterion=color_criterion, texture_critetion=texture_criterion, tv_criterion=tv_criterion, x=batch['x'], y=batch['y'], optimizers=optimizers) print("\nEpoch: %s\n" % i) print("Total Loss: %s \n" % loss['total_loss']) print("Color Loss: %s \n" % loss['color_loss']) print("Content Loss: %s \n" % loss['content_loss']) print("TV Loss: %s \n" % loss['tv_loss']) print("Texture Loss: %s \n" % loss['texture_loss'])
def __init__(self, args, network_config): self.start_epoch = 0 self.num_epochs = network_config["num_epochs"] self.batch_size = network_config["batch_size"] self.noisy_dir = args.noisy_dir self.clean_dir = args.clean_dir self.wandb = args.wandb if self.wandb == None: from torch.utils.tensorboard import SummaryWriter self.writer = SummaryWriter(args.log) self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") print(self.device) self.generator_noisy2clean = model.Generator().to(self.device) self.generator_clean2noisy = model.Generator().to(self.device) self.discriminator_noisy = model.Discriminator().to(self.device) self.discriminator_clean = model.Discriminator().to(self.device) cudnn.benchmark = True cudnn.fastest = True mse_loss = torch.nn.MSELoss() g_params = list(self.generator_noisy2clean.parameters()) + list( self.generator_clean2noisy.parameters()) d_params = list(self.discriminator_noisy.parameters()) + list( self.discriminator_clean.parameters()) self.generator_lr = network_config["generator_lr"] self.discriminator_lr = network_config["discriminator_lr"] self.decay_iter = network_config["decay_iter"] self.generator_lr_decay = self.generator_lr / self.decay_iter self.discriminator_lr_decay = self.discriminator_lr / self.decay_iter self.start_decay = network_config["start_decay"] self.cycle_loss_lambda = network_config["cycle_loss_lambda"] self.identity_loss_lambda = network_config["identity_loss_lambda"] self.identity_loss_stop = network_config["identity_loss_stop"] self.generator_optim = torch.optim.Adam(g_params, lr=self.generator_lr, betas=(0.5, 0.999)) self.discriminator_optim = torch.optim.Adam(d_params, lr=self.discriminator_lr, betas=(0.5, 0.999)) self.model_dir = args.model_save self.output_dir = args.output_save self.generator_loss = [] self.discriminator_loss = []
def initialize(self): if self.imsize == 256: num_down = 8 num_layer = 3 blocks = 5 elif self.imsize == 128: num_down = 7 num_layer = 2 blocks = 4 else: raise NotImplementedError('not support other size') if self.where_add == 'input' : self.generator = model.Generator_Unet_in( self.in_dim,self.out_dim, self.nz, self.ngf ,self.dropout, norm = self.norm, nl = self.nl , upsample = self.upsample , num_down = num_down) elif self.where_add == 'all' : self.generator = model.Generator_Unet_all(self.in_dim,self.out_dim, self.nz, self.ngf ,self.dropout, norm = self.norm, nl = self.nl , upsample = self.upsample , num_down = num_down ) self.discriminator1 = model.Discriminator( self.in_dim , self.imsize , self.nz , self.ndf , dropout = self.dropout , num_d = self.num_d, norm = self.norm , gan_mode = self.gan_mode , num_layer = num_layer) self.discriminator2 = model.Discriminator( self.in_dim , self.imsize , self.nz , self.ndf , dropout = self.dropout , num_d = self.num_d, norm = self.norm , gan_mode = self.gan_mode , num_layer = num_layer) self.encoder = model.Encoder_resnet( ndf = self.nef , nz = self.nz ,blocks = blocks ,in_dim = self.in_dim , out_dim = self.out_dim , norm = self.norm , nl = self.nl) self.generator = model.init_net(self.generator,gpu_ids = self.gpu_ids) self.discriminator1 = model.init_net(self.discriminator1,gpu_ids = self.gpu_ids) self.discriminator2 = model.init_net(self.discriminator2, gpu_ids = self.gpu_ids) self.encoder = model.init_net(self.encoder ,gpu_ids = self.gpu_ids) self.optimizers = [] self.opt_g = torch.optim.Adam(self.generator.parameters(), lr = self.lr , betas = (self.beta1,0.999)) self.opt_e = torch.optim.Adam(self.encoder.parameters(), lr = self.lr , betas = (self.beta1,0.999)) self.opt_d1 = torch.optim.Adam(self.discriminator1.parameters(), lr = self.lr , betas = (self.beta1,0.999)) self.opt_d2 = torch.optim.Adam(self.discriminator2.parameters(), lr = self.lr , betas = (self.beta1,0.999)) self.optimizers.append(self.opt_g) self.optimizers.append(self.opt_d1) self.optimizers.append(self.opt_d2) self.optimizers.append(self.opt_e) self.set_gan_criterion(gan_mode = self.gan_mode) self.reconstruction_loss = nn.L1Loss() self.setup()
def __init__(self): self.config = utils.Config('./config.yml') self.device = None self.summary = {} self.dumpPath = None self.sysConfig() self.setSummary() self.pipeRaw = self.loadDataset() self.pipeLen = self.pipeRaw['train'].__len__() self.pipe = None self.pipeIter() self.gen = model.Generator(self.config) self.dis = model.Discriminator(self.config) if self.config.GPU == -1 and self.config.CUDA: print('Using MultiGPU') self.gen = nn.parallel.DataParallel(self.gen).to(self.device) else: self.gen = self.gen.to(self.device) self.optGen = torch.optim.Adam(self.gen.parameters(), lr=eval(self.config.LR), betas=self.config.BETA) self.optDis = torch.optim.Adam(self.dis.parameters(), lr=eval(self.config.LR), betas=self.config.BETA)
def main(opt): dataset = data.Dataset(dataset=opt.dataset, pool_size=opt.pool_size, sample_size=opt.sample_size) dataset.show_inf() feature_size, att_size = dataset.feature_size, dataset.att_size discriminator = model.Discriminator(feature_size, att_size).cuda() generator = model.Generator(feature_size, att_size).cuda() for epoch in range(opt.epochs): # d_loss = train.train_disciminator(discriminator, generator, dataset, opt.lr, opt.batch_size, epoch) # g_loss = train.train_generator(discriminator, generator, dataset, opt.lr, opt.batch_size, epoch) d_loss, g_loss = train.train_together(discriminator, generator, dataset, opt.lr, opt.batch_size, epoch) D_zsl_acc = test.compute_acc(discriminator, dataset, opt1='zsl', opt2='test_unseen') D_seen_acc = test.compute_acc(discriminator, dataset, opt1='gzsl', opt2='test_seen') D_unseen_acc = test.compute_acc(discriminator, dataset, opt1='gzsl', opt2='test_unseen') D_harmonic_mean = (2 * D_seen_acc * D_unseen_acc) / (D_seen_acc + D_unseen_acc) print("Epoch {}/{}...".format(epoch + 1, opt.epochs)) print("D_Loss: {:.4f}".format(d_loss), "zsl_acc: {:.4f}".format(D_zsl_acc), "seen_acc: {:.4f}".format(D_seen_acc), "unseen_acc: {:.4f}".format(D_unseen_acc), "harmonic_mean: {:.4f}".format(D_harmonic_mean) ) G_zsl_acc = test.compute_acc(generator, dataset, opt1='zsl', opt2='test_unseen') G_seen_acc = test.compute_acc(generator, dataset, opt1='gzsl', opt2='test_seen') G_unseen_acc = test.compute_acc(generator, dataset, opt1='gzsl', opt2='test_unseen') G_harmonic_mean = (2 * G_seen_acc * G_unseen_acc) / (G_seen_acc + G_unseen_acc) print("G_Loss: {:.4f}".format(g_loss), "zsl_acc: {:.4f}".format(G_zsl_acc), "seen_acc: {:.4f}".format(G_seen_acc), "unseen_acc: {:.4f}".format(G_unseen_acc), "harmonic_mean: {:.4f}".format(G_harmonic_mean) )
def main(opt): dataset = data.Dataset(dataset=opt.dataset, pool_size=opt.pool_size, sample_size=opt.sample_size) dataset.show_inf() feature_size, att_size = dataset.feature_size, dataset.att_size discriminator = model.Discriminator(feature_size, att_size, opt.t1).cuda() generator = model.Generator(feature_size, att_size, opt.t2).cuda() train2.train(discriminator, generator, dataset, d_lr=opt.d_lr, g_lr=opt.g_lr,\ batch_size=opt.batch_size, alpha=opt.alpha, epochs=opt.epochs)
def main(num_block, epochs_list, batch_size, is_train, is_continue, is_save) : ##################################### # 환경변수 지정, 추후 parser로 변환 예정 ##################################### #path_image = os.path.join(os.getcwd(), 'train_image/') path_image = os.path.join(os.getcwd(), '../datasets/DogData/') path_model = os.path.join(os.getcwd(), 'save_model/') print(f' Path of Image : {path_image}') model_name = 'model.pth' # 블록이 하나 이하일시 종료 if(num_block <= 1) : print('Not enough block, Terminated') return # 원하는 갯수의 블록을 가진 Generator 와 Discriminator 생성 할까 했는데 무조건 맥스로 생성 generator = model.Generator(batch_size, 9) discriminator = model.Discriminator(9) if torch.cuda.is_available() == True : generator = generator.cuda() discriminator = discriminator.cuda() if is_continue : file_model = os.path.join(path_model, model_name) if os.path.exists(file_model) : model_dict = load_model(file_model) generator.load_state_dict(model_dict['generator']) discriminator.load_state_dict(model_dict['discriminator']) # 학습 시작 if is_train : train(num_block, generator, discriminator, batch_size, epochs_list, path_image) print(f'Train End') if is_save : if os.path.exists(path_model) == False : os.mkdir(path_model) file_model = os.path.join(path_model, model_name) save_model(file_model, generator, discriminator) for i in range(5) : ###############################33 # 임시 생성 테스트용 z = torch.rand(100) if torch.cuda.is_available() : z = z.cuda() image = generator(z, num_block).cpu().detach().numpy()[0] image = image.transpose((1,2,0)) img = Image.fromarray(np.uint8(image*255)) img.save(os.path.join('save_image/',f'save{i}.png'), format='png')
def main(): voc = util.Voc(init_from_file="data/voc_b.txt") netR_path = 'output/rf_dis.pkg' netG_path = 'output/net_p' netD_path = 'output/net_d' agent_path = 'output/net_gan_%d_%d_%dx%d' % (SIGMA * 10, BL * 10, BATCH_SIZE, MC) netR = util.Environment(netR_path) agent = model.Generator(voc) agent.load_state_dict(T.load(netG_path + '.pkg')) df = pd.read_table('data/CHEMBL251.txt') df = df[df['PCHEMBL_VALUE'] >= 6.5] data = util.MolData(df, voc) loader = DataLoader(data, batch_size=BATCH_SIZE, shuffle=True, drop_last=True, collate_fn=data.collate_fn) netD = model.Discriminator(VOCAB_SIZE, EMBED_DIM, FILTER_SIZE, NUM_FILTER) if not os.path.exists(netD_path + '.pkg'): Train_dis_BCE(netD, agent, loader, epochs=100, out=netD_path) netD.load_state_dict(T.load(netD_path + '.pkg')) best_score = 0 log = open(agent_path + '.log', 'w') for epoch in range(1000): print('\n--------\nEPOCH %d\n--------' % (epoch + 1)) print('\nPolicy Gradient Training Generator : ') Train_GAN(agent, netD, netR) print('\nAdversarial Training Discriminator : ') Train_dis_BCE(netD, agent, loader, epochs=1) seqs = agent.sample(1000) ix = util.unique(seqs) smiles, valids = util.check_smiles(seqs[ix], agent.voc) scores = netR(smiles) scores[valids == False] = 0 unique = (scores >= 0.5).sum() / 1000 if best_score < unique: T.save(agent.state_dict(), agent_path + '.pkg') best_score = unique print("Epoch+: %d average: %.4f valid: %.4f unique: %.4f" % (epoch, scores.mean(), valids.mean(), unique), file=log) for i, smile in enumerate(smiles): print('%f\t%s' % (scores[i], smile), file=log) for param_group in agent.optim.param_groups: param_group['lr'] *= (1 - 0.01) log.close()
def GetModelAndLoss(device, opt): generator = model.Generator(opt.img_size_h, opt.latent_dim, opt.channels) # generator generator.apply(model.weights_init_normal) discriminator = model.Discriminator(opt.img_size_h, opt.channels) # discriminator discriminator.apply(model.weights_init_normal) adversarial_loss = torch.nn.BCELoss() # loss return generator.to(device), discriminator.to(device), adversarial_loss.to(device)
def __init__(self, trial, step, size, batch_size, learning_rate, max_epoch, tfrecord_path, checkpoint_dir, scale,num_of_data, conf): print('Initialize Training') self.trial=trial self.step=step self.HEIGHT = size[0] self.WIDTH = size[1] self.CHANNEL = size[2] self.BATCH_SIZE = batch_size self.learning_rate=learning_rate self.EPOCH = max_epoch self.tfrecord_path = tfrecord_path self.checkpoint_dir=checkpoint_dir self.scale= scale self.num_of_data=num_of_data self.conf=conf self.input = tf.placeholder(dtype=tf.float32,shape=[None,self.HEIGHT//self.scale,self.WIDTH//self.scale,self.CHANNEL]) self.label = tf.placeholder(dtype=tf.float32,shape=[None,self.HEIGHT,self.WIDTH,self.CHANNEL]) self.GEN = model.FRACTAL(self.input,self.scale) self.NAT = model.NMD(self.GEN.output) self.DIS_fake = model.Discriminator(self.GEN.output, reuse=False) self.DIS_real = model.Discriminator(self.label, reuse=True)
def load_model(imsize, nz, nc, ndf, ngf, n_extra_layers, savepath): # load net models print(savepath) assert (os.path.exists(savepath)==True) netD_name = 'netD.pth' netG_name = 'netG.pth' print(netD_name, netG_name) netD = dcgan.Discriminator(imsize = imsize, nz=nz, nc=nc, ndf=ndf, n_extra_layers=n_extra_layers) netG = dcgan.Generator(imsize = imsize, nz=nz, nc=nc, ngf=ngf, n_extra_layers=n_extra_layers) netD.load_state_dict(torch.load(savepath + netD_name)) netG.load_state_dict(torch.load(savepath + netG_name)) losses = np.load(savepath + "losses.npy") return netD, netG, list(losses)
def main(): source_train_loader = mnist.mnist_train_loader target_train_loader = mnistm.mnistm_train_loader if torch.cuda.is_available(): get_free_gpu() print('Running GPU : {}'.format(torch.cuda.current_device())) encoder = model.Extractor().cuda() classifier = model.Classifier().cuda() discriminator = model.Discriminator().cuda() train.source_only(encoder, classifier, discriminator, source_train_loader, target_train_loader, save_name) train.dann(encoder, classifier, discriminator, source_train_loader, target_train_loader, save_name) else: print("There is no GPU -_-!")
def setup_model(latent_dim, input_dim, choice_activation_func='PReLU', gpu=True, checkpoint=None): generator = model.Generator(latent_dim=latent_dim, output_dim=input_dim, choice_activation_func=choice_activation_func) discriminator = model.Discriminator( input_dim=input_dim, choice_activation_func=choice_activation_func) if checkpoint is not None: generator.load_state_dict(checkpoint['generator_state_dict']) discriminator.load_state_dict(checkpoint['discriminator_state_dict']) if gpu: generator = generator.cuda() discriminator = discriminator.cuda() return generator, discriminator
print('Building models...') torch.cuda.set_device(args.gpu) if not (os.path.isfile(args.agent_file) and os.path.isfile(args.agent_file) and os.path.isfile(args.agent_file)): print("need an agent file") exit() args.agent_file = args.env + ".model.80.tar" agent = model.Agent(action_size, args.agent_latent).cuda() agent.load_state_dict(torch.load(args.agent_file, map_location=map_loc)) Z_dim = args.latent wae_z_dim = args.wae_latent encoder = model.Encoder(Z_dim).cuda() generator = model.Generator(Z_dim, action_size).cuda() discriminator = model.Discriminator(Z_dim, action_size).cuda() Q = model.Q_net(args.wae_latent).cuda() P = model.P_net(args.wae_latent).cuda() Q.load_state_dict(torch.load(args.Q, map_location=map_loc)) P.load_state_dict(torch.load(args.P, map_location=map_loc)) encoder.train() generator.train() discriminator.train() Q.eval() P.eval() optim_gen = optim.Adam(generator.parameters(), lr=args.lr, betas=(0.0,0.9)) optim_enc = optim.Adam(filter(lambda p: p.requires_grad, encoder.parameters()), lr=args.lr, betas=(0.0,0.9)) optim_disc = optim.Adam(filter(lambda p: p.requires_grad, discriminator.parameters()), lr=args.lr, betas=(0.0,0.9))
def train(pre_trained=None): # create folder to save models and loss graphs reference = hp['net_type'] + str(time.strftime("_%Y%m%d_%H%M%S")) checkpoints_folder = hp["output_dir"] + '/checkpoints/' + reference os.makedirs(checkpoints_folder, exist_ok=True) # save hyper parameter settings pickle_file_location = checkpoints_folder + "/hp.pkl" pickle_file = open(pickle_file_location, "wb") pickle.dump(hp, pickle_file) pickle_file.close() # create data iterator dataset = DataGenerator(hp) iterator = DataLoader(dataset=dataset, batch_size=hp['batch_size'], num_workers=hp['num_workers'], pin_memory=True, shuffle=False, drop_last=True) # create model and loss generator = model.Generator().to(device) discriminator = model.Discriminator().to(device) loss = model.BCELoss().to(device) # optimizer g_optimizer = torch.optim.Adam(params=generator.parameters(), lr=hp['learning_rate'], betas=(hp['beta1'], hp['beta2'])) d_optimizer = torch.optim.Adam(params=discriminator.parameters(), lr=hp['learning_rate'], betas=(hp['beta1'], hp['beta2'])) start_epoch = 0 # load pre trained model if pre_trained is not None: ckpt = torch.load(pre_trained) generator.load_state_dict(ckpt['G']) discriminator.load_state_dict(ckpt['D']) g_optimizer.load_state_dict(ckpt['G_opt']) d_optimizer.load_state_dict(ckpt['D_opt']) start_epoch = ckpt['epoch'] + 1 # init loss arrays generator_loss = np.zeros(hp['num_epochs']) discriminator_loss = np.zeros(hp['num_epochs']) # training loop for epoch in range (start_epoch, hp['num_epochs']): g_loss = 0 d_loss = 0 for i, img in enumerate(iterator): print("--------------------------------------------------") img = img.to(device) noise_input = torch.randn(hp['batch_size'], 100, 1, 1) noise_input = noise_input.to(device) g_optimizer.zero_grad() d_optimizer.zero_grad() gen_out = generator(noise_input) disc_real_out = discriminator(img.float()) disc_fake_out = discriminator(gen_out) # loss and back prop # generator_loss target = torch.ones(hp['batch_size'], 1).to(device) gen_loss = loss(disc_fake_out, target) gen_loss.backward() g_optimizer.step() g_loss += gen_loss.item() # discriminator real loss target = torch.ones(hp['batch_size'], 1).to(device) d_loss_real = loss(disc_real_out, target) # discriminator fake loss target = torch.zeros(hp['batch_size'], 1).to(device) disc_fake_out = discriminator(gen_out.detach()) d_loss_fake = loss(disc_fake_out, target) disc_loss = d_loss_fake + d_loss_real disc_loss.backward(retain_graph=True) d_optimizer.step() d_loss += disc_loss.item() print("epoch = {}, Training_sample={}, generator_loss ={}".format(epoch, i, gen_loss)) print("epoch = {}, Training_sample={}, discriminator_loss ={}".format(epoch, i, disc_loss)) # average loss per epoch generator_loss[epoch] = g_loss/(i+1) discriminator_loss[epoch] = d_loss/(i+1) print("epoch = {}, average generator_loss ={}".format(epoch, generator_loss[epoch])) print("epoch = {}, average discriminator_loss ={}".format(epoch, discriminator_loss[epoch])) # plot loss curves and save model plt.plot(range(1, len(generator_loss)+1), generator_loss, 'b-', label="Generator loss") plt.xlabel("epochs") plt.ylabel("Error") plt.legend(loc='best') plt.savefig(checkpoints_folder + "/generator_loss.jpeg", bbox_inches="tight") plt.clf() plt.plot(range(1, len(discriminator_loss)+1), discriminator_loss, 'b-', label="Discriminator loss") plt.xlabel("epochs") plt.ylabel("Error") plt.legend(loc='best') plt.savefig(checkpoints_folder + "/discriminator_loss.jpeg", bbox_inches="tight") plt.clf() net_save = {'G': generator.state_dict(), 'D': discriminator.state_dict(), 'G_opt': g_optimizer.state_dict(), 'D_opt': d_optimizer.state_dict(), 'epoch': epoch} torch.save(net_save, checkpoints_folder + "/dc_gan_weights_epoch{}.pth".format(epoch))
cudnn.benchmark = True # For speed i.e, cudnn autotuner ######################################################## if torch.cuda.is_available() and not opt.cuda: print( "WARNING: You have a CUDA device, so you should probably run with --cuda" ) #calling the dataloader data = util.DATA_LOADER(opt) print("training samples: ", data.ntrain) ############## MODEL INITIALIZATION ############# netE = model.Encoder(opt) netG = model.HYBRID_FUSION_ATTENTION(opt) netD = model.Discriminator(opt) print(netE) print(netG) print(netD) ################################################ #init tensors input_res = torch.FloatTensor(opt.batch_size, opt.resSize) input_test_labels = torch.LongTensor(opt.fake_batch_size, opt.nclass_all) input_labels = torch.LongTensor(opt.batch_size, opt.nseen_class) input_train_early_fusion_att = torch.FloatTensor(opt.batch_size, opt.attSize) input_test_early_fusion_att = torch.FloatTensor(opt.fake_batch_size, opt.attSize) noise = torch.FloatTensor(opt.batch_size, opt.attSize) one = torch.FloatTensor([1])
def main(args): #writer = SummaryWriter(log_dir='/content/cycleGAN_seismic_noise/runs/'+ datetime.now().strftime('%b%d_%H-%M-%S')) writer = "" wandb.login() wandb.init(project="cycleGAN_seismic_noise") wandb.wath_called = False ##=== run with model package ====# G_12 = model.Generator(args.batch_size) G_21 = model.Generator(args.batch_size) D_1 = model.Discriminator(args.batch_size) D_2 = model.Discriminator(args.batch_size) G_12.weight_init(mean=0.0, std=0.02) G_21.weight_init(mean=0.0, std=0.02) D_1.weight_init(mean=0.0, std=0.02) D_2.weight_init(mean=0.0, std=0.02) if torch.cuda.is_available(): G_12 = G_12.cuda() G_21 = G_21.cuda() D_1 = D_1.cuda() D_2 = D_2.cuda() #=== optimizer & learning rate schedulers optimizer_G = torch.optim.Adam(itertools.chain(G_12.parameters(), G_21.parameters()), lr=args.lr, betas=(args.beta1, 0.999)) optimizer_D_A = torch.optim.Adam(D_1.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) optimizer_D_B = torch.optim.Adam(D_1.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR( optimizer_G, lr_lambda=residual_model.LambdaLR(args.num_epochs, 0, round(args.num_epochs / 2)).step) lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR( optimizer_D_A, lr_lambda=residual_model.LambdaLR(args.num_epochs, 0, round(args.num_epochs / 2)).step) lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR( optimizer_D_B, lr_lambda=residual_model.LambdaLR(args.num_epochs, 0, round(args.num_epochs / 2)).step) # load seismic dataset A_data, B_data = data.train_dataset(args.dir, args.batch_size, args.image_size, args.num_iter_train) model_state = args.state_dict if model_state != "": checkpoint = torch.load(model_state) G_12.load_state_dict(checkpoint['G_12_state_dict']) G_21.load_state_dict(checkpoint['G_21_state_dict']) D_1.load_state_dict(checkpoint['D_1_state_dict']) D_2.load_state_dict(checkpoint['D_2_state_dict']) optimizer_G.load_state_dict(checkpoint['optimizer_G']) optimizer_D_A.load_state_dict(checkpoint['optimizer_D_A']) optimizer_D_B.load_state_dict(checkpoint['optimizer_D_B']) lr_scheduler_G.load_state_dict(checkpoint['lr_scheduler_G']) lr_scheduler_D_A.load_state_dict(checkpoint['lr_scheduler_D_A']) lr_scheduler_D_B.load_state_dict(checkpoint['lr_scheduler_D_B']) cur_epoch = checkpoint['epoch'] else: cur_epoch = 0 train(G_12, G_21, D_1, D_2, optimizer_G, optimizer_D_A, optimizer_D_B, lr_scheduler_G, lr_scheduler_D_A, lr_scheduler_D_B, args.batch_size, cur_epoch, args.num_epochs, A_data, B_data, writer, args.num_iter_train)
torch.cuda.manual_seed(args.seed) # ############################################################################## # Load data ################################################################################ from data_loader import Data_loader real_datas = Data_loader('data/', args.img_size, args.batch_size, use_cuda) # ############################################################################## # Build model # ############################################################################## import model G = model.Generator(args.img_size, args.img_size, args.channel_dims, args.z_dim) D = model.Discriminator(args.img_size, args.img_size, args.channel_dims, args.relu_leak) if use_cuda: G, D = G.cuda(), D.cuda() optimizer_D = torch.optim.Adam(D.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) optimizer_G = torch.optim.Adam(G.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) criterion = torch.nn.BCELoss() # ############################################################################## # Training # ############################################################################## real_label, fake_label = 1, 0 def train(): loss_g = loss_d = 0 for real_data in tqdm(real_datas, mininterval=1,
T.Resize(32), T.ToTensor(), ]) trainset = tvsets.MNIST(root='data', train=True, transform=transform, download=True) trainloader = D.DataLoader(trainset, batch_size=1, shuffle=True, num_workers=1) size = 32 * 32 #filename = print_freq = 100 #save_freq = generator = nets.Generator(size) discriminator = nets.Discriminator(size) def noise(): return torch.randn(1, size) def train(): g_optimizer = optim.Adam(generator.parameters(), lr=0.01) d_optimizer = optim.Adam(discriminator.parameters(), lr=0.01) loss_func = nn.BCELoss() sum_g_loss = 0 sum_d_loss = 0 print("Beginning training...")
def main(args): if args.dataset == 'cifar10': test_dataloader = data.DataLoader(datasets.CIFAR10( args.data_path, download=True, transform=cifar_transformer(), train=False), batch_size=args.batch_size, drop_last=False) train_dataset = CIFAR10(args.data_path) args.num_images = 50000 args.budget = 2500 args.initial_budget = 5000 args.num_classes = 10 elif args.dataset == 'cifar100': test_dataloader = data.DataLoader(datasets.CIFAR100( args.data_path, download=True, transform=cifar_transformer(), train=False), batch_size=args.batch_size, drop_last=False) train_dataset = CIFAR100(args.data_path) args.num_images = 50000 args.budget = 2500 args.initial_budget = 5000 args.num_classes = 100 elif args.dataset == 'imagenet': test_dataloader = data.DataLoader(datasets.ImageFolder( args.data_path, transform=imagenet_transformer()), drop_last=False, batch_size=args.batch_size) train_dataset = ImageNet(args.data_path) args.num_images = 1281167 args.budget = 64060 args.initial_budget = 128120 args.num_classes = 1000 elif args.dataset == 'semeval': train_dataset = SemEvalRes('train') test_dataloader = data.DataLoader(SemEvalRes('test', mlb=train_dataset.mlb), batch_size=args.batch_size, drop_last=False) N_samples = len(train_dataset) # print(N_samples) args.num_images = N_samples args.budget = int(0.05 * N_samples) args.initial_budget = int(0.1 * N_samples) args.num_classes = train_dataset.n_classes emb_size = train_dataset.emb_size else: raise NotImplementedError all_indices = set(np.arange(args.num_images)) initial_indices = random.sample(all_indices, args.initial_budget) sampler = data.sampler.SubsetRandomSampler(initial_indices) # dataset with labels available querry_dataloader = data.DataLoader(train_dataset, sampler=sampler, batch_size=args.batch_size, drop_last=True) args.cuda = args.cuda and torch.cuda.is_available() solver = Solver(args, test_dataloader) # splits = [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4] splits = [] split = len(querry_dataloader) / N_samples current_indices = list(initial_indices) accuracies = [] f1s = [] while split < N_samples: splits.append(split) # need to retrain all the models on the new images # re initialize and retrain the models # task_model = vgg.vgg16_bn(num_classes=args.num_classes) task_model = linear.LinearModel(emb_size, args.num_classes) vae = model.VAE(args.latent_dim, emb_size) discriminator = model.Discriminator(args.latent_dim) unlabeled_indices = np.setdiff1d(list(all_indices), current_indices) unlabeled_sampler = data.sampler.SubsetRandomSampler(unlabeled_indices) unlabeled_dataloader = data.DataLoader(train_dataset, sampler=unlabeled_sampler, batch_size=args.batch_size, drop_last=False) # train the models on the current data acc, f1, vae, discriminator = solver.train(querry_dataloader, task_model, vae, discriminator, unlabeled_dataloader) print('Final accuracy with {}% of data is: {:.2f}'.format( int(split * 100), acc)) print('Final f1 micro with {}% of data is: {:.2f}'.format( int(split * 100), f1)) accuracies.append(acc) f1s.append(f1) sampled_indices = solver.sample_for_labeling(vae, discriminator, unlabeled_dataloader) current_indices = list(current_indices) + list(sampled_indices) sampler = data.sampler.SubsetRandomSampler(current_indices) querry_dataloader = data.DataLoader(train_dataset, sampler=sampler, batch_size=args.batch_size, drop_last=True) perf = {'labeled_percentage': splits, 'accuracies': accuracies, 'f1': f1s} torch.save(perf, os.path.join(args.out_path, args.log_name))
gpus = (0, 1, 2, 3) print(f'GPUs: {gpus}') generator = model.Globalgenerator(2, 2, 3, 3, 71, 3) generator = model.Enhancer(generator, './work/generator_final.pth.tar', 2, 0.5, 71, 3, True) generator = model.Generator_enhance(generator, resume_path='./vggface.pth', consistency_iter=args.con_iter) generator = nn.DataParallel(generator, device_ids=gpus).cuda() print(f'Generator loaded, consistency iter: {args.con_iter}') multidis = model.MultiscaleDiscriminator(71) lossfunction_dis = criterion.LSGanLoss() discriminator = model.Discriminator(multidis, lossfunction_dis) discriminator = nn.DataParallel(discriminator, device_ids=gpus).cuda() print('Discriminator loaded') if args.resume: g_ckpt = torch.load(f'{args.ckpt_path}/generator_final.pth.tar') d_ckpt = torch.load(f'{args.ckpt_path}/discriminator_final.pth.tar') generator.module.model.load_state_dict(g_ckpt) discriminator.module.load_state_dict(d_ckpt) print(f'Checkpoints Loaded: {args.ckpt_path}') # trainset = dataset.Reenactset( # pkl_path=args.pkl_path, # img_path=args.img_path, # max_iter=args.max_iter, # consistency_iter = args.con_iter,
val_set = loader.ValDataset(PATHS[opt.dataset][opt.artifacts]['hr']['valid'], lr_dir=PATHS[opt.dataset][opt.artifacts]['lr']['valid'], **vars(opt)) val_loader = DataLoader(dataset=val_set, num_workers=1, batch_size=1, shuffle=False) # prepare neural networks if opt.generator.lower() == 'dsgan': model_g = model.Generator(n_res_blocks=opt.num_res_blocks) elif opt.generator.lower() == 'deresnet': model_g = model.De_resnet(n_res_blocks=opt.num_res_blocks, scale=opt.upscale_factor) else: raise NotImplementedError('Generator model [{:s}] not recognized'.format(opt.generator)) print('# Initializing {}'.format(opt.generator)) print('# generator parameters:', sum(param.numel() for param in model_g.parameters()), '\n') model_d = model.Discriminator(kernel_size=opt.kernel_size, wgan=opt.wgan, highpass=opt.highpass, D_arch=opt.discriminator, norm_layer=opt.norm_layer, filter_type=opt.filter, cs=opt.cat_or_sum) print('# discriminator parameters:', sum(param.numel() for param in model_d.parameters())) g_loss_module = loss.GeneratorLoss(**vars(opt)) # filters are used for generating validation images filter_low_module = model.FilterLow(kernel_size=opt.kernel_size, gaussian=opt.filter == 'gau', include_pad=False) filter_high_module = model.FilterHigh(kernel_size=opt.kernel_size, gaussian=opt.filter == 'gau', include_pad=False) print('# FS type: {}, kernel size={}'.format(opt.filter, opt.kernel_size)) if torch.cuda.is_available(): model_g = model_g.cuda() model_d = model_d.cuda() filter_low_module = filter_low_module.cuda() filter_high_module = filter_high_module.cuda()
import model import helper import itertools # inspired by https://gist.github.com/LMescheder/b4d57e094cc8522497fdb5cf8fb44676 # Hyperparameters batch_size = 128 sigma = 1.0 z_dim = 64 learning_rate = 1e-4 reg_param = 10. # Define networks generator = model.Generator(z_dim) discriminator = model.Discriminator(2) optim_gen = torch.optim.RMSprop(generator.parameters(), lr=learning_rate) optim_dis = torch.optim.RMSprop(discriminator.parameters(), lr=learning_rate) # Print Real distribution mu = np.vstack([np.cos(2*np.pi*k/8), np.sin(2*np.pi*k/8)] for k in range(batch_size)) x_real = mu + sigma * np.random.normal(0.0, 0.0, [batch_size, 2]) helper.display_result(x_real, cmap='Reds') plt.savefig("image/real.png") for epoch in range(10000): # Getting the target distribution mu = np.vstack([np.cos(2*np.pi*k/8), np.sin(2*np.pi*k/8)] for k in range(batch_size)) x_real = mu + sigma * np.random.normal(0.0, 0.0, [batch_size, 2])
def main(): # Set random seem for reproducibility manualSeed = 999 # manualSeed = random.randint(1, 10000) # use if you want new results print("Random Seed: ", manualSeed) random.seed(manualSeed) torch.manual_seed(manualSeed) dataroot = "/home/shunan/Data/celeba" workers = 2 # Number of workers for dataloader batch_size = 128 # Batch size during training image_size = 64 # Spatial size of training images. nc = 3 # Number of channels in the training images. For color images this is 3 nz = 100 # Size of z latent vector (i.e. size of generator input) ngf = 64 # Size of z latent vector (i.e. size of generator input) ndf = 64 # Size of feature maps in discriminator num_epochs = 5 # Number of training epochs lr = 0.0002 # Learning rate for optimizers beta1 = 0.5 # Beta1 hyperparam for Adam optimizers ngpu = 1 # Number of GPUs available. Use 0 for CPU mode. # We can use an image folder dataset the way we have it setup. Create the dataset dataset = dset.ImageFolder(root=dataroot, transform=transforms.Compose([ transforms.Resize(image_size), transforms.CenterCrop(image_size), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) # Create the dataloader dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=workers) # Decide which device we want to run on device = torch.device("cuda:0" if ( torch.cuda.is_available() and ngpu > 0) else "cpu") # Plot some training images real_batch = next(iter(dataloader)) plt.figure(figsize=(8, 8)) plt.axis("off") plt.title("Training Images") plt.imshow( np.transpose( vutils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(), (1, 2, 0))) plt.show() # Create the generator netG = model.Generator(ngpu, nz, ngf, nc).to(device) # Apply the weights_init function to randomly initialize all weights to mean=0, stdev=0.2. netG.apply(model.weights_init) # Create the Discriminator netD = model.Discriminator(ngpu, ndf, nc).to(device) # Handle multi-gpu if desired if (device.type == 'cuda') and (ngpu > 1): netD = nn.DataParallel(netD, list(range(ngpu))) # Apply the weights_init function to randomly initialize all weights to mean=0, stdev=0.2. netD.apply(model.weights_init) # Initialize BCELoss function criterion = nn.BCELoss() # Create batch of latent vectors that we will use to visualize the progression of the generator fixed_noise = torch.randn(64, nz, 1, 1, device=device) # Establish convention for real and fake labels during training real_label = 1 fake_label = 0 # Setup Adam optimizers for both G and D optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999)) optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999)) ################################################## # Training Loop # ################################################## # Lists to keep track of progress img_list = [] G_losses = [] D_losses = [] iters = 0 print("Starting Training Loop...") # For each epoch for epoch in range(num_epochs): # For each batch in the dataloader for i, data in enumerate(dataloader, 0): ############################ # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) ########################### ## Train with all-real batch netD.zero_grad() # Format batch real_cpu = data[0].to(device) b_size = real_cpu.size(0) label = torch.full((b_size, ), real_label, device=device) # Forward pass real batch through D output = netD(real_cpu).view(-1) # Calculate loss on all-real batch errD_real = criterion(output, label) # Calculate gradients for D in backward pass errD_real.backward() D_x = output.mean().item() ## Train with all-fake batch # Generate batch of latent vectors noise = torch.randn(b_size, nz, 1, 1, device=device) # Generate fake image batch with G fake = netG(noise) label.fill_(fake_label) # Classify all fake batch with D output = netD(fake.detach()).view(-1) # Calculate D's loss on the all-fake batch errD_fake = criterion(output, label) # Calculate the gradients for this batch errD_fake.backward() D_G_z1 = output.mean().item() # Add the gradients from the all-real and all-fake batches errD = errD_real + errD_fake # Update D optimizerD.step() ############################ # (2) Update G network: maximize log(D(G(z))) ########################### netG.zero_grad() label.fill_(real_label) # fake labels are real for generator cost # Since we just updated D, perform another forward pass of all-fake batch through D output = netD(fake).view(-1) # Calculate G's loss based on this output errG = criterion(output, label) # Calculate gradients for G errG.backward() D_G_z2 = output.mean().item() # Update G optimizerG.step() # Output training stats if i % 50 == 0: print( '[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f' % (epoch, num_epochs, i, len(dataloader), errD.item(), errG.item(), D_x, D_G_z1, D_G_z2)) # Save Losses for plotting later G_losses.append(errG.item()) D_losses.append(errD.item()) # Check how the generator is doing by saving G's output on fixed_noise if (iters % 500 == 0) or ((epoch == num_epochs - 1) and (i == len(dataloader) - 1)): with torch.no_grad(): fake = netG(fixed_noise).detach().cpu() img_list.append( vutils.make_grid(fake, padding=2, normalize=True)) iters += 1 # plotting the loss. plt.figure(figsize=(10, 5)) plt.title("Generator and Discriminator Loss During Training") plt.plot(G_losses, label="G") plt.plot(D_losses, label="D") plt.xlabel("iterations") plt.ylabel("Loss") plt.legend() plt.show() # Show some of the results. # Grab a batch of real images from the dataloader real_batch = next(iter(dataloader)) # Plot the real images plt.figure(figsize=(15, 15)) plt.subplot(1, 2, 1) plt.axis("off") plt.title("Real Images") plt.imshow( np.transpose( vutils.make_grid(real_batch[0].to(device)[:64], padding=5, normalize=True).cpu(), (1, 2, 0))) # Plot the fake images from the last epoch plt.subplot(1, 2, 2) plt.axis("off") plt.title("Fake Images") plt.imshow(np.transpose(img_list[-1], (1, 2, 0))) plt.show()
datasets.CIFAR10('../data/', train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])), #What does transform do? batch_size=args.batch_size, shuffle=True, num_workers=1, pin_memory=True) Z_dim = 128 #number of updates to discriminator for every update to generator disc_iters = 5 # discriminator = torch.nn.DataParallel(Discriminator()).cuda() # TODO: try out multi-gpu training if args.model == 'resnet': discriminator = model_resnet.Discriminator().cuda() generator = model_resnet.Generator(Z_dim).cuda() else: discriminator = model.Discriminator().cuda() generator = model.Generator(Z_dim).cuda() # because the spectral normalization module creates parameters that don't require gradients (u and v), we don't want to # optimize these using sgd. We only let the optimizer operate on parameters that _do_ require gradients # TODO: replace Parameters with buffers, which aren't returned from .parameters() method. optim_disc = optim.Adam(filter(lambda p: p.requires_grad, discriminator.parameters()), lr=args.lr, betas=(0.0,0.9)) optim_gen = optim.Adam(generator.parameters(), lr=args.lr, betas=(0.0,0.9)) # use an exponentially decaying learning rate scheduler_d = optim.lr_scheduler.ExponentialLR(optim_disc, gamma=0.99) scheduler_g = optim.lr_scheduler.ExponentialLR(optim_gen, gamma=0.99) def main(): def train(epoch): for batch_idx, (data, target) in enumerate(loader):
def __init__(self): self.opt = get_args() self.models = {} self.weight = {} self.weight["static"] = self.opt.static_weight self.weight["dynamic"] = self.opt.dynamic_weight self.device = "cuda" self.criterion_d = nn.BCEWithLogitsLoss() self.parameters_to_train = [] self.parameters_to_train_D = [] # Initializing models self.models["encoder"] = model.Encoder(18, self.opt.height, self.opt.width, True) if self.opt.type == "both": self.models["static_decoder"] = model.Decoder( self.models["encoder"].resnet_encoder.num_ch_enc) self.models["static_discr"] = model.Discriminator() self.models["dynamic_decoder"] = model.Discriminator() self.models["dynamic_decoder"] = model.Decoder( self.models["encoder"].resnet_encoder.num_ch_enc) else: self.models["decoder"] = model.Decoder( self.models["encoder"].resnet_encoder.num_ch_enc) self.models["discriminator"] = model.Discriminator() for key in self.models.keys(): self.models[key].to(self.device) if "discr" in key: self.parameters_to_train_D += list( self.models[key].parameters()) else: self.parameters_to_train += list(self.models[key].parameters()) # Optimization self.model_optimizer = optim.Adam(self.parameters_to_train, self.opt.lr) self.model_lr_scheduler = optim.lr_scheduler.StepLR( self.model_optimizer, self.opt.scheduler_step_size, 0.1) self.model_optimizer_D = optim.Adam(self.parameters_to_train_D, self.opt.lr) self.model_lr_scheduler_D = optim.lr_scheduler.StepLR( self.model_optimizer_D, self.opt.scheduler_step_size, 0.1) self.patch = (1, self.opt.occ_map_size // 2**4, self.opt.occ_map_size // 2**4) self.valid = Variable(torch.Tensor( np.ones((self.opt.batch_size, *self.patch))), requires_grad=False).float().cuda() self.fake = Variable(torch.Tensor( np.zeros((self.opt.batch_size, *self.patch))), requires_grad=False).float().cuda() ## Data Loaders dataset_dict = { "3Dobject": dataloader.KITTIObject, "odometry": dataloader.KITTIOdometry, "argo": dataloader.Argoverse } self.dataset = dataset_dict[self.opt.split] fpath = os.path.join(os.path.dirname(__file__), "splits", self.opt.split, "{}_files.txt") train_filenames = readlines(fpath.format("train")) val_filenames = readlines(fpath.format("val")) self.val_filenames = val_filenames self.train_filenames = train_filenames img_ext = '.png' if self.opt.ext == "png" else '.jpg' train_dataset = self.dataset(self.opt, train_filenames) val_dataset = self.dataset(self.opt, val_filenames) self.train_loader = DataLoader(train_dataset, self.opt.batch_size, True, num_workers=self.opt.num_workers, pin_memory=True, drop_last=True) self.val_loader = DataLoader(val_dataset, 1, True, num_workers=self.opt.num_workers, pin_memory=True, drop_last=True) print("Using split:\n ", self.opt.split) print( "There are {:d} training items and {:d} validation items\n".format( len(train_dataset), len(val_dataset)))
def train(): if len(args.gpu_idx) > 1: multi_gpu = True gpu_list = [int(i) for i in args.gpu_idx.split(',')] else: multi_gpu = False writer = SummaryWriter(log_dir=log_path) if args.dataset == 'cifar10': dataset = datasets.CIFAR10(args.data_dir, download=True, transform=transforms.Compose([ transforms.Resize(args.imsize), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ])) elif args.dataset == 'mnist': dataset = datasets.MNIST(args.data_dir, train=True, download=True, transform=transforms.Compose([ transforms.Scale(args.imsize), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])) args.in_dim = 1 args.out_dim = 1 else: dataset = datasets.ImageFolder(args.data_dir, transform=transforms.Compose([ transforms.Resize(args.imsize), transforms.ToTensor(), transforms.Normalize( (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])) dataloader = torch.utils.data.DataLoader( dataset , batch_size = args.batch , \ shuffle = True , num_workers = args.worker) #device = torch.device() generator = model.Generator(args) discriminator = model.Discriminator(args) generator.apply(weights_init) discriminator.apply(weights_init) gan_criterion = nn.BCELoss() aux_criterion = nn.CrossEntropyLoss() #input_noise = torch.from_numpy( np.random.normal(0,1,[args.batch , args.dim_embed]) ) #input_label = torch.from_numpy( np.random.randint(0,args.num_class, [args.batch,1 ]) ) if args.l_smooth: # training strategy stated in improved GAN real_label = 0.9 fake_label = 0.1 else: real_label = 1.0 fake_label = 0.0 step = 0 if args.gpu: # acutally do nothing? because bce and cce don't have paramters gan_criterion = gan_criterion.cuda() aux_criterion = aux_criterion.cuda() generator = generator.cuda() discriminator = discriminator.cuda() if multi_gpu: print('multi gpu') generator = nn.DataParallel(generator, device_ids=gpu_list) discriminator = nn.DataParallel(discriminator, device_ids=gpu_list) opt_d = optim.Adam(discriminator.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) opt_g = optim.Adam(generator.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) if os.path.isfile(os.path.join(ckpt_path, args.run_name + '.ckpt')): print('found ckpt file' + os.path.join(ckpt_path, args.run_name + '.ckpt')) ckpt = torch.load(os.path.join(ckpt_path, args.run_name + '.ckpt')) generator.load_state_dict(ckpt['generator']) discriminator.load_state_dict(ckpt['discriminator']) opt_d.load_state_dict(ckpt['opt_d']) opt_g.load_state_dict(ckpt['opt_g']) step = ckpt['step'] for i in range(args.epoch): for j, data in enumerate(dataloader): images, labels = data[0], data[1] batch = images.shape[0] input_noise = torch.from_numpy( np.random.normal(0, 1, [batch, args.dim_embed]).astype(np.float32)) input_label = torch.from_numpy( np.random.randint(0, args.num_class, [batch])) input_noise = np.random.normal( 0, 1, [batch, args.dim_embed]).astype(np.float32) class_onehot = np.zeros((batch, args.num_class)) class_onehot[np.arange(batch), input_label] = 1 input_noise[np.arange(batch), :args.num_class] = class_onehot[ np.arange(batch)] input_noise = torch.from_numpy(input_noise) real_target = torch.full((batch, 1), real_label) fake_target = torch.full((batch, 1), fake_label) aux_target = torch.autograd.Variable(labels) if args.gpu: input_noise = input_noise.cuda() input_label = input_label.cuda() images = images.cuda() labels = labels.cuda() real_target = real_target.cuda() fake_target = fake_target.cuda() aux_target = aux_target.cuda() # train generator # 好像不call也沒關係 opt_g.zero_grad() fake = generator(input_noise, input_label) gan_out_g, aux_out_g = discriminator(fake) if args.wgan: gan_loss_g = -torch.mean(gan_out_g) else: gan_loss_g = gan_criterion(gan_out_g, real_target) aux_loss_g = aux_criterion(aux_out_g, input_label) g_loss = (gan_loss_g + args.aux_weight * aux_loss_g) * 0.5 g_loss.backward() #opt_g.step() opt_g.step() # train discriminator with real samples opt_d.zero_grad() gan_out_r, aux_out_r = discriminator(images) if args.wgan: gan_loss_r = -torch.mean(gan_out_r) else: gan_loss_r = gan_criterion(gan_out_r, real_target) aux_loss_r = aux_criterion(aux_out_r, aux_target) d_real_loss = (gan_loss_r + args.aux_weight * aux_loss_r) / 2.0 # train discriminator with fake samples #fake = generator(input_noise , input_label).detach() gan_out_f, aux_out_f = discriminator(fake.detach()) if args.wgan: gan_loss_f = torch.mean(gan_out_f) else: gan_loss_f = gan_criterion(gan_out_f, fake_target) aux_loss_f = aux_criterion(aux_out_f, input_label) d_fake_loss = (gan_loss_f + args.aux_weight * aux_loss_f) / 2.0 if args.wgan and args.gp: gp = model.gradient_penalty(discriminator, images, fake, args.num_class, args.gpu) d_loss = 0.5 * (d_real_loss + d_fake_loss) + args.gp_weight * gp else: d_loss = (d_real_loss + d_fake_loss) / 2.0 d_loss.backward() opt_d.step() step = step + 1 if step % 100 == 0: writer.add_scalar('losses/g_loss', g_loss, step) writer.add_scalar('losses/d_loss', d_loss, step) grid = vutils.make_grid(fake.detach(), normalize=True) writer.add_image('generated', grid, step) if args.wgan and not args.gp: for p in discriminator.parameters(): p.data.clamp_(-args.clip, args.clip) if step % args.save_freq == 0: torch.save( { 'step': step, 'generator': generator.state_dict(), 'discriminator': discriminator.state_dict(), 'opt_d': opt_d.state_dict(), 'opt_g': opt_g.state_dict() }, os.path.join(ckpt_path, args.run_name + '.ckpt')) if step % args.sample_freq == 0: sample_generator(generator, input_noise, input_label, step) pred = np.concatenate( [aux_out_r.data.cpu().numpy(), aux_out_f.data.cpu().numpy()], axis=0) gt = np.concatenate( [labels.data.cpu().numpy(), input_label.data.cpu().numpy()], axis=0) d_acc = np.mean(np.argmax(pred, axis=1) == gt) print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f, acc: %d%%] [G loss: %f]" % \ (i, args.epoch , j, len(dataloader), d_loss.item(), 100.0 * d_acc, g_loss.item()))
formatter = logging.Formatter( '%(levelname).1s %(asctime)s [%(name)s] %(message)s', '%H:%M:%S') handler.setFormatter(formatter) logging.basicConfig(level=logging.DEBUG, handlers=[handler]) logger = logging.getLogger(__name__) import data, model from gan import WassersteinGAN def parse_arguments(): parser = argparse.ArgumentParser('HW4') parser.add_argument('--gpus', type=str, default='0') parser.add_argument('--epochs', type=int, default=300) parser.add_argument('--batch_size', type=int, default=64) return parser.parse_args() if __name__ == '__main__': args = parse_arguments() os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus data_sampler = data.DataSampler(os.path.join('data', 'tags.csv')) noise_sampler = data.NoiseSampler() g_net = model.Generator() d_net = model.Discriminator() wgan = WassersteinGAN(g_net, d_net, data_sampler, noise_sampler) wgan.train(epochs=args.epochs, batch_size=args.batch_size)
batch_size = 32 learning_rate = 0.0001 beta1 = 0.5 z_size = 5 save_interval = 10 ### input variables z = tf.placeholder(tf.float32, [batch_size, z_size]) a = tf.placeholder(tf.float32, [batch_size, 32, 32, 32, 1]) rgba = tf.placeholder(tf.float32, [batch_size, 32, 32, 32, 4]) train = tf.placeholder(tf.bool) ### build models G = model.Generator(z_size) D = model.Discriminator() rgba_ = G(a, z, train) y_ = D(rgba_, train) y = D(rgba, train) label_real = np.zeros([batch_size, 2], dtype=np.float32) label_fake = np.zeros([batch_size, 2], dtype=np.float32) label_real[:, 0] = 1 label_fake[:, 1] = 1 loss_G = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_, labels=tf.constant(label_real))) loss_D = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_, labels=tf.constant(label_fake))) loss_D += tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y, labels=tf.constant(label_real))) var_G = [v for v in tf.trainable_variables() if 'g_' in v.name]
save_interval = myconfig.save_interval sample_interval = 10 ##################################################. ## build graph ################################################## z = tf.placeholder(tf.float32, [batch_size, z_size]) a = tf.placeholder(tf.float32, [batch_size, 64, 64, 64, 1]) rgba = tf.placeholder(tf.float32, [batch_size, 64, 64, 64, 4]) train = tf.placeholder(tf.bool) k_t = tf.placeholder(tf.float32, shape=[]) ### build models G = model.Generator(z_size) D = model.Discriminator(z_size) rgba_ = G(a, z, train) rgba_out_ = D(rgba_, train) rgba_out = D(rgba, train) loss_D,loss_G,k_tp,convergence_measure = model.loss(rgba, rgba_out, rgba_, rgba_out_, k_t=k_t) var_G = [v for v in tf.trainable_variables() if 'g_' in v.name] var_D = [v for v in tf.trainable_variables() if 'd_' in v.name] opt_g = tf.train.AdamOptimizer(learning_rate,beta1).minimize(loss_G, var_list = var_G) opt_d = tf.train.AdamOptimizer(learning_rate,beta1).minimize(loss_D, var_list=var_D) print 'var_G' for v in var_G: