def generate(**kwargs): ''' 随机生成动漫头像,并根据netd的分数选择较好的 ''' for k_, v_ in kwargs.items(): setattr(opt, k_, v_) netg, netd = NetG(opt).eval(), NetD(opt).eval() noises = t.randn(opt.gen_search_num, opt.nz, 1, 1).normal_(opt.gen_mean, opt.gen_std) noises = Variable(noises, volatile=True) map_location = lambda storage, loc: storage netd.load_state_dict(t.load(opt.netd_path, map_location=map_location)) netg.load_state_dict(t.load(opt.netg_path, map_location=map_location)) if opt.gpu: netd.cuda() netg.cuda() noises = noises.cuda() # 生成图片,并计算图片在判别器的分数 fake_img = netg(noises) scores = netd(fake_img).data # 挑选最好的某几张 indexs = scores.topk(opt.gen_num)[1] result = [] for ii in indexs: result.append(fake_img.data[ii]) # 保存图片 tv.utils.save_image(t.stack(result), opt.gen_img, normalize=True, range=(-1, 1))
def discriminator(**kwargs): """ 随机生成动漫头像,并根据netd的分数选择较好的 """ for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device = t.device('cuda') if opt.gpu else t.device('cpu') netd = NetD(opt).eval() noises = t.randn(opt.gen_search_num, 3, 224, 224).normal_(opt.gen_mean, opt.gen_std) # noises = t.randn(3, 3, 224).normal_(opt.gen_mean, opt.gen_std) noises = noises.to(device) map_location = lambda storage, loc: storage netd.load_state_dict(t.load(opt.netd_path, map_location=map_location)) netd.to(device) # 生成图片,并计算图片在判别器的分数 fake_label = netd(noises).detach() # scores = netd(fake_img).detach() data = pd.DataFrame(fake_label) data.to_csv("score.csv") print (fake_label)
def singletest(**kwargs): """ python main.py singletest --nogpu --vis=False --netd-path=checkpoints/netd_2199.pth :param kwargs: :return: """ for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device = t.device('cuda') if opt.gpu else t.device('cpu') pil_img = Image.open(opt.test_path) pil_img = pil_img.convert('RGB') transforms = tv.transforms.Compose([ tv.transforms.Resize(opt.image_size), tv.transforms.CenterCrop(opt.image_size), tv.transforms.ToTensor(), tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) data = transforms(pil_img) array = numpy.asarray(data) array = t.from_numpy(array) array = array.view(1, 3, 224, 224).to(device) print (array.shape) netd = NetD(opt).eval() map_location = lambda storage, loc: storage netd.load_state_dict(t.load(opt.netd_path, map_location=map_location)) netd.to(device) getlabel = netd(array).double().detach() print (getlabel)
def generate(**kwargs): """ 随机生成动漫头像,并根据netd的分数选择较好的 """ for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device=t.device('cuda') if opt.gpu else t.device('cpu') netg, netd = NetG(opt).eval(), NetD(opt).eval() noises = t.randn(opt.gen_search_num, opt.nz, 1, 1).normal_(opt.gen_mean, opt.gen_std) noises = noises.to(device) map_location = lambda storage, loc: storage netd.load_state_dict(t.load(opt.netd_path, map_location=map_location)) netg.load_state_dict(t.load(opt.netg_path, map_location=map_location)) netd.to(device) netg.to(device) # 生成图片,并计算图片在判别器的分数 fake_img = netg(noises) scores = netd(fake_img).detach() # 挑选最好的某几张 indexs = scores.topk(opt.gen_num)[1] result = [] for ii in indexs: result.append(fake_img.data[ii]) # 保存图片 epoch_num = opt.netd_path[opt.netd_path.find("_") + 1:opt.netd_path.rfind(".")] randid = random.randint(0, 1000) pic_stack_name = "".join(["result_", str(epoch_num), "_", str(randid), ".png"]) tv.utils.save_image(t.stack(result), pic_stack_name, normalize=True, range=(-1, 1)) base_dir = "./imgs/tiny_epoch_{0}_{1}".format(epoch_num, randid) if not os.path.exists(base_dir): os.mkdir(base_dir) for picx in range(len(result)): picx_name = "".join([base_dir, "/picid_", str(picx), ".png"]) tv.utils.save_image(result[picx], picx_name, normalize=True, range=(-1, 1))
def setup(): args = parse_args() if args.cfg_file is not None: cfg_from_file(args.cfg_file) cfg.GPU_ID = args.gpu_id cfg.DATA_DIR = args.data_dir args.manualSeed = 100 random.seed(args.manualSeed) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) if cfg.CUDA: torch.cuda.manual_seed_all(args.manualSeed) now = datetime.datetime.now(dateutil.tz.tzlocal()) timestamp = now.strftime('%Y_%m_%d_%H_%M_%S') output_dir = '../output/%s_%s_%s' % \ (cfg.DATASET_NAME, cfg.CONFIG_NAME, timestamp) torch.cuda.set_device(cfg.GPU_ID) cudnn.benchmark = True device = torch.device("cuda" if torch.cuda.is_available() else "cpu") netG = NetG(cfg.TRAIN.NF, 100).to(device) netD = NetD(cfg.TRAIN.NF).to(device) #blah text_encoder = RNN_ENCODER(27297, nhidden=cfg.TEXT.EMBEDDING_DIM) state_dict = torch.load(cfg.TEXT.DAMSM_NAME, map_location=lambda storage, loc: storage) text_encoder.load_state_dict(state_dict) text_encoder.cuda() for p in text_encoder.parameters(): p.requires_grad = False text_encoder.eval() return text_encoder, netG
def generate(**kwargs): """ 随机生成动漫头像,并根据netd的分数选择较好的 """ for k_, v_ in kwargs.items(): setattr(opt, k_, v_) # 将网络模型置为预测模式 不保存中间结果,加速 netg, netd = NetG(opt).eval(), NetD(opt).eval() # 初始化gen_search_num张噪声,期望生成gen_search_num张预测图像 noises = t.randn(opt.gen_search_num, opt.nz, 1, 1).normal_(opt.gen_mean, opt.gen_std) noises = Variable(noises, volatile=True) # 将模型参数加载到cpu中 map_location = lambda storage, loc: storage netd.load_state_dict(t.load(opt.netd_path, map_location=map_location)) netg.load_state_dict(t.load(opt.netg_path, map_location=map_location)) # 模型和输入噪声转到GPU中 if opt.gpu: netd.cuda() netg.cuda() noises = noises.cuda() # 生成图片,并计算图片在判别器的分数 fake_img = netg(noises) scores = netd(fake_img).data # 挑选最好的某几张 从512章图片中按分数排序,取前64张 的下标 indexs = scores.topk(opt.gen_num)[1] result = [] for ii in indexs: result.append(fake_img.data[ii]) # 保存图片 tv.utils.save_image(t.stack(result), opt.gen_img, normalize=True, range=(-1, 1))
def test(**kwargs): """ python main.py test --nogpu --vis=False --netd-path=checkpoints/netd_2199.pth --gen-num=1 --batch_size=1 """ for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device = t.device('cuda') if opt.gpu else t.device('cpu') # 数据 transforms = tv.transforms.Compose([ tv.transforms.Resize(opt.image_size), tv.transforms.CenterCrop(opt.image_size), tv.transforms.ToTensor(), tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) datasettest = mydataset.AADBDatasetTest(opt.alltestdata_path, opt.alltestlabel_path, transforms=transforms) # dataset = tv.datasets.ImageFolder(opt.data_path, transform=transforms) dataloadertest = t.utils.data.DataLoader(datasettest, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, drop_last=True ) netd = NetD(opt).eval() map_location = lambda storage, loc: storage netd.load_state_dict(t.load(opt.netd_path, map_location=map_location)) netd.to(device) error = 0 criterion = t.nn.MSELoss().to(device) for ii, (img, label) in tqdm.tqdm(enumerate(dataloadertest)): img224 = img.to(device) getlabel = netd(img224).double().detach() label = label.double() error = error + criterion(label, getlabel) print (error / 1000)
def train(): dataset = torchvision.datasets.ImageFolder(conf.data_path, transform=transforms) dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=conf.batch_size, shuffle=True, drop_last=True) netG = NetG(conf.ngf, conf.nz) netD = NetD(conf.ndf) criterion = nn.BCELoss() optimizerG = torch.optim.Adam(netG.parameters(), lr=conf.lr, betas=(conf.beta1, 0.999)) optimizerD = torch.optim.Adam(netD.parameters(), lr=conf.lr, betas=(conf.beta1, 0.999)) label = torch.FloatTensor(conf.batch_size) real_label = 1 fake_label = 0 for epoch in range(1, conf.epoch + 1): for i, (imgs, _) in enumerate(dataloader): #step1:固定G,训练D optimizerD.zero_grad() output = netD(imgs) #让D尽可能把真图片识别为1 label.data.fill_(real_label) errD_real = criterion(output, label) errD_real.backward() #让D尽可能把假图判别为0 label.data.fill_(fake_label) noise = torch.randn(conf.batch_size, conf.nz, 1, 1) fake = netG(noise) #生成假图 output = netD(fake.detach()) #避免梯度传到G,因为G不用更新 errD_fake = criterion(output, label) errD_fake.backward() errD = errD_fake + errD_real optimizerD.step() #step2:固定判别器D,训练生成器G optimizerG.zero_grad() label.data.fill_(real_label) #让D尽可能把G生成的假图判别为1 output = netD(fake) errG = criterion(output, label) errG.backward() optimizerG.step() if i % 4 == 0: rate = i * 1.0 / len(dataloader) * 100 logger.info( "epoch={}, i={}, N={}, rate={}%, errD={}, errG={}".format( epoch, i, len(dataloader), rate, errD, errG)) #end-for save_image(fake.data, '%s/fake_samples_epoch_%03d.png' % (conf.checkpoints, epoch), normalize=True) torch.save(netG.state_dict(), '%s/netG_%03d.pth' % (conf.checkpoints, epoch)) torch.save(netD.state_dict(), '%s/netD_%03d.pth' % (conf.checkpoints, epoch))
def train(**kwargs): for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device = t.device("cuda") if opt.gpu else t.device("cpu") # 数据处理,输出规范为-1~1 transforms = tv.transforms.Compose([ tv.transforms.Resize(opt.image_size), tv.transforms.CenterCrop(opt.image_size), tv.transforms.ToTensor(), tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) dataset = tv.datasets.ImageFolder(opt.data_path, transform=transforms) dataloader = t.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, drop_last=True) # 网络 netg, netd = NetG(opt), NetD(opt) map_location = lambda storage, loc: storage if opt.netd_path: netd.load_state_dict(t.load(opt.netd_path, map_location=map_location)) if opt.netg_path: netg.load_state_dict(t.load(opt.netg_path, map_location=map_location)) netd.to(device) netg.to(device) # 定义优化器和损失 optimizer_g = t.optim.Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999)) optimizer_d = t.optim.Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999)) criterion = t.nn.BCELoss() # 真图片label为1,假图片label为0, noise为生成网络的输入 true_labels = t.ones(opt.batch_size).to(device) fake_labels = t.zeros(opt.batch_size).to(device) fix_noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device) noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device) # 用来结果的均值和标准差 errord_meter = AverageValueMeter() errorg_meter = AverageValueMeter() epochs = range(opt.max_epoch) for epoch in iter(epochs): for ii, (img, _) in tqdm.tqdm(enumerate(dataloader)): real_img = img.to(device) if ii % opt.d_every == 0: # 训练判别器 optimizer_d.zero_grad() ## 尽可能把真图片判别为正 output = netd(real_img) error_d_real = criterion(output, true_labels) error_d_real.backward() ## 尽可能把假图片判断为错误 noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1)) # 使用detach来关闭G求梯度,加速训练 fake_img = netg(noises).detach() output = netd(fake_img) error_d_fake = criterion(output, fake_labels) error_d_fake.backward() optimizer_d.step() error_d = error_d_fake + error_d_real errord_meter.add(error_d.item()) if ii % opt.g_every == 0: # 训练生成器 optimizer_g.zero_grad() noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1)) fake_img = netg(noises) output = netd(fake_img) # 尽可能把假的图片也判别为1 error_g = criterion(output, true_labels) error_g.backward() optimizer_g.step() errorg_meter.add(error_g.item()) # 可视化 # 保存模型、图片 if (epoch + 1) % opt.save_every == 0: fix_fake_imgs = netg(fix_noises) tv.utils.save_image(fix_fake_imgs.data[:64], "%s%s.png" % (opt.save_path, epoch), normalize=True, range=(-1, 1)) t.save(netd.state_dict(), r"./checkpoints/netd_%s.pth" % epoch) t.save(netg.state_dict(), r"./checkpoints/netg_%s.pth" % epoch) errord_meter.reset() errorg_meter.reset()
transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) dataset = dset.ImageFolder(dataroot, transform=transforms) dataloader = torch.utils.data.DataLoader( dataset=dataset, batch_size=batch_size, shuffle=True, drop_last=True, ) netG = NetG(ngf, nz).to(device) # 生成器网络 Generator netG.apply(weights_init) # 生成器网络w初始化 Generator weights initialization netD = NetD(ndf).to(device) # 判别器网络 Discriminator netD.apply(weights_init) # 生成器网络w初始化 Discriminator weights initialization # 打印网络模型 Print the model print(netG) print(netD) criterion = nn.BCELoss() # 初始化损失函数 Initialize BCELoss function # Create batch of latent vectors that we will use to visualize the progression of the generator fixed_noise = torch.randn(100, nz, 1, 1, device=device) # 评估真假的标签 真为1 假为0 Establish convention for real and fake labels during training real_label = 1 fake_label = 0 # Adam优化 Setup Adam optimizers for both G and D
class GANAgent(object): def __init__(self, input_size, output_size, num_env, num_step, gamma, lam=0.95, learning_rate=1e-4, ent_coef=0.01, clip_grad_norm=0.5, epoch=3, batch_size=128, ppo_eps=0.1, update_proportion=0.25, use_gae=True, use_cuda=False, use_noisy_net=False, hidden_dim=512): self.model = CnnActorCriticNetwork(input_size, output_size, use_noisy_net) self.num_env = num_env self.output_size = output_size self.input_size = input_size self.num_step = num_step self.gamma = gamma self.lam = lam self.epoch = epoch self.batch_size = batch_size self.use_gae = use_gae self.ent_coef = ent_coef self.ppo_eps = ppo_eps self.clip_grad_norm = clip_grad_norm self.update_proportion = update_proportion self.device = torch.device('cuda' if use_cuda else 'cpu') self.netG = NetG(z_dim=hidden_dim) #(input_size, z_dim=hidden_dim) self.netD = NetD(z_dim=1) self.netG.apply(weights_init) self.netD.apply(weights_init) self.optimizer_policy = optim.Adam(list(self.model.parameters()), lr=learning_rate) self.optimizer_G = optim.Adam(list(self.netG.parameters()), lr=learning_rate, betas=(0.5, 0.999)) self.optimizer_D = optim.Adam(list(self.netD.parameters()), lr=learning_rate, betas=(0.5, 0.999)) self.netG = self.netG.to(self.device) self.netD = self.netD.to(self.device) self.model = self.model.to(self.device) def reconstruct(self, state): state = torch.Tensor(state).to(self.device) state = state.float() reconstructed = self.vae(state.unsqueeze(0))[0].squeeze(0) return reconstructed.detach().cpu().numpy() def get_action(self, state): state = torch.Tensor(state).to(self.device) state = state.float() policy, value_ext, value_int = self.model(state) action_prob = F.softmax(policy, dim=-1).data.cpu().numpy() action = self.random_choice_prob_index(action_prob) return action, value_ext.data.cpu().numpy().squeeze( ), value_int.data.cpu().numpy().squeeze(), policy.detach() @staticmethod def random_choice_prob_index(p, axis=1): r = np.expand_dims(np.random.rand(p.shape[1 - axis]), axis=axis) return (p.cumsum(axis=axis) > r).argmax(axis=axis) def compute_intrinsic_reward(self, obs): obs = torch.FloatTensor(obs).to(self.device) #embedding = self.vae.representation(obs) #reconstructed_embedding = self.vae.representation(self.vae(obs)[0]) # why use index[0] reconstructed_img, embedding, reconstructed_embedding = self.netG(obs) intrinsic_reward = (embedding - reconstructed_embedding ).pow(2).sum(1) / 2 # Not use reconstructed loss return intrinsic_reward.detach().cpu().numpy() def train_model(self, s_batch, target_ext_batch, target_int_batch, y_batch, adv_batch, next_obs_batch, old_policy): s_batch = torch.FloatTensor(s_batch).to(self.device) target_ext_batch = torch.FloatTensor(target_ext_batch).to(self.device) target_int_batch = torch.FloatTensor(target_int_batch).to(self.device) y_batch = torch.LongTensor(y_batch).to(self.device) adv_batch = torch.FloatTensor(adv_batch).to(self.device) next_obs_batch = torch.FloatTensor(next_obs_batch).to(self.device) sample_range = np.arange(len(s_batch)) #reconstruction_loss = nn.MSELoss(reduction='none')] l_adv = nn.MSELoss(reduction='none') l_con = nn.L1Loss(reduction='none') l_enc = nn.MSELoss(reduction='none') l_bce = nn.BCELoss(reduction='none') with torch.no_grad(): policy_old_list = torch.stack(old_policy).permute( 1, 0, 2).contiguous().view(-1, self.output_size).to(self.device) m_old = Categorical(F.softmax(policy_old_list, dim=-1)) log_prob_old = m_old.log_prob(y_batch) # ------------------------------------------------------------ #recon_losses = np.array([]) #kld_losses = np.array([]) mean_err_g_adv_per_batch = np.array([]) mean_err_g_con_per_batch = np.array([]) mean_err_g_enc_per_batch = np.array([]) mean_err_d_per_batch = np.array([]) for i in range(self.epoch): np.random.shuffle(sample_range) for j in range(int(len(s_batch) / self.batch_size)): sample_idx = sample_range[self.batch_size * j:self.batch_size * (j + 1)] # -------------------------------------------------------------------------------- # for generative curiosity (GAN loss) #gen_next_state, mu, logvar = self.vae(next_obs_batch[sample_idx]) ############### netG forward ############################################## gen_next_state, latent_i, latent_o = self.netG( next_obs_batch[sample_idx]) ############### netD forward ############################################## pred_real, feature_real = self.netD(next_obs_batch[sample_idx]) pred_fake, feature_fake = self.netD(gen_next_state) #d = len(gen_next_state.shape) #recon_loss = reconstruction_loss(gen_next_state, next_obs_batch[sample_idx]).mean(axis=list(range(1, d))) ############### netG backward ############################################# self.optimizer_G.zero_grad() err_g_adv_per_img = l_adv( self.netD(next_obs_batch[sample_idx])[1], self.netD(gen_next_state)[1]).mean( axis=list(range(1, len(feature_real.shape)))) err_g_con_per_img = l_con( next_obs_batch[sample_idx], gen_next_state).mean( axis=list(range(1, len(gen_next_state.shape)))) err_g_enc_per_img = l_enc(latent_i, latent_o).mean(-1) #kld_loss = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(axis=1) # TODO: keep this proportion of experience used for VAE update? # Proportion of experience used for VAE update img_num = len(err_g_con_per_img) mask = torch.rand(img_num).to(self.device) mask = (mask < self.update_proportion).type( torch.FloatTensor).to(self.device) mean_err_g_adv = (err_g_adv_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) mean_err_g_con = (err_g_con_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) mean_err_g_enc = (err_g_enc_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) # hyperparameter weights: w_adv = 1 w_con = 50 w_enc = 1 mean_err_g = mean_err_g_adv * w_adv +\ mean_err_g_con * w_con +\ mean_err_g_enc * w_enc mean_err_g.backward(retain_graph=True) self.optimizer_G.step() mean_err_g_adv_per_batch = np.append( mean_err_g_adv_per_batch, mean_err_g_adv.detach().cpu().numpy()) mean_err_g_con_per_batch = np.append( mean_err_g_con_per_batch, mean_err_g_con.detach().cpu().numpy()) mean_err_g_enc_per_batch = np.append( mean_err_g_enc_per_batch, mean_err_g_enc.detach().cpu().numpy()) ############## netD backward ############################################## self.optimizer_D.zero_grad() real_label = torch.ones_like(pred_real).to(self.device) fake_label = torch.zeros_like(pred_fake).to(self.device) err_d_real_per_img = l_bce(pred_real, real_label) err_d_fake_per_img = l_bce(pred_fake, fake_label) mean_err_d_real = (err_d_real_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) mean_err_d_fake = (err_d_fake_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) mean_err_d = (mean_err_d_real + mean_err_d_fake) / 2 mean_err_d.backward() self.optimizer_D.step() mean_err_d_per_batch = np.append( mean_err_d_per_batch, mean_err_d.detach().cpu().numpy()) if mean_err_d.item() < 1e-5: self.netD.apply(weights_init) print('Reloading net d') ############# policy update ############################################### policy, value_ext, value_int = self.model(s_batch[sample_idx]) m = Categorical(F.softmax(policy, dim=-1)) log_prob = m.log_prob(y_batch[sample_idx]) ratio = torch.exp(log_prob - log_prob_old[sample_idx]) surr1 = ratio * adv_batch[sample_idx] surr2 = torch.clamp(ratio, 1.0 - self.ppo_eps, 1.0 + self.ppo_eps) * adv_batch[sample_idx] actor_loss = -torch.min(surr1, surr2).mean() critic_ext_loss = F.mse_loss(value_ext.sum(1), target_ext_batch[sample_idx]) critic_int_loss = F.mse_loss(value_int.sum(1), target_int_batch[sample_idx]) critic_loss = critic_ext_loss + critic_int_loss entropy = m.entropy().mean() self.optimizer_policy.zero_grad() loss = actor_loss + 0.5 * critic_loss - self.ent_coef * entropy loss.backward() #global_grad_norm_(list(self.model.parameters())+list(self.vae.parameters())) do we need this step #global_grad_norm_(list(self.model.parameter())) or just norm policy self.optimizer_poilicy.step() return mean_err_g_adv_per_batch, mean_err_g_con_per_batch, mean_err_g_enc_per_batch, mean_err_d_per_batch def train_just_vae(self, s_batch, next_obs_batch): s_batch = torch.FloatTensor(s_batch).to(self.device) next_obs_batch = torch.FloatTensor(next_obs_batch).to(self.device) sample_range = np.arange(len(s_batch)) l_adv = nn.MSELoss(reduction='none') l_con = nn.L1Loss(reduction='none') l_enc = nn.MSELoss(reduction='none') l_bce = nn.BCELoss(reduction='none') mean_err_g_adv_per_batch = np.array([]) mean_err_g_con_per_batch = np.array([]) mean_err_g_enc_per_batch = np.array([]) mean_err_d_per_batch = np.array([]) for i in range(self.epoch): np.random.shuffle(sample_range) for j in range(int(len(s_batch) / self.batch_size)): sample_idx = sample_range[self.batch_size * j:self.batch_size * (j + 1)] ############### netG forward ############################################## gen_next_state, latent_i, latent_o = self.netG( next_obs_batch[sample_idx]) ############### netD forward ############################################## pred_real, feature_real = self.netD(next_obs_batch[sample_idx]) pred_fake, feature_fake = self.netD(gen_next_state) #d = len(gen_next_state.shape) #recon_loss = reconstruction_loss(gen_next_state, next_obs_batch[sample_idx]).mean(axis=list(range(1, d))) ############### netG backward ############################################# self.optimizer_G.zero_grad() err_g_adv_per_img = l_adv( self.netD(next_obs_batch[sample_idx])[1], self.netD(gen_next_state)[1]).mean( axis=list(range(1, len(feature_real.shape)))) err_g_con_per_img = l_con( next_obs_batch[sample_idx], gen_next_state).mean( axis=list(range(1, len(gen_next_state.shape)))) err_g_enc_per_img = l_enc(latent_i, latent_o).mean(-1) #kld_loss = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(axis=1) # TODO: keep this proportion of experience used for VAE update? # Proportion of experience used for VAE update img_num = len(err_g_con_per_img) mask = torch.rand(img_num).to(self.device) mask = (mask < self.update_proportion).type( torch.FloatTensor).to(self.device) mean_err_g_adv = (err_g_adv_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) mean_err_g_con = (err_g_con_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) mean_err_g_enc = (err_g_enc_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) # hyperparameter weights: w_adv = 1 w_con = 50 w_enc = 1 mean_err_g = mean_err_g_adv * w_adv +\ mean_err_g_con * w_con +\ mean_err_g_enc * w_enc mean_err_g.backward(retain_graph=True) self.optimizer_G.step() mean_err_g_adv_per_batch = np.append( mean_err_g_adv_per_batch, mean_err_g_adv.detach().cpu().numpy()) mean_err_g_con_per_batch = np.append( mean_err_g_con_per_batch, mean_err_g_con.detach().cpu().numpy()) mean_err_g_enc_per_batch = np.append( mean_err_g_enc_per_batch, mean_err_g_enc.detach().cpu().numpy()) ############## netD backward ############################################## self.optimizer_D.zero_grad() real_label = torch.ones_like(pred_real).to(self.device) fake_label = torch.zeros_like(pred_fake).to(self.device) err_d_real_per_img = l_bce(pred_real, real_label) err_d_fake_per_img = l_bce(pred_fake, fake_label) mean_err_d_real = (err_d_real_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) mean_err_d_fake = (err_d_fake_per_img * mask).sum() / torch.max( mask.sum(), torch.Tensor([1]).to(self.device)) mean_err_d = (mean_err_d_real + mean_err_d_fake) / 2 mean_err_d.backward() self.optimizer_D.step() mean_err_d_per_batch = np.append( mean_err_d_per_batch, mean_err_d.detach().cpu().numpy()) return mean_err_g_adv_per_batch, mean_err_g_con_per_batch, mean_err_g_enc_per_batch, mean_err_d_per_batch
class AnoGAN: """AnoGAN Class """ def __init__(self, opt): # super(AnoGAN, self).__init__(opt, dataloader) # Initalize variables. self.opt = opt self.niter = self.opt.niter self.start_iter = 0 self.netd_niter = 5 self.test_iter = 100 self.lr = self.opt.lr self.batchsize = {'train': self.opt.batchsize, 'test': 1} self.pretrained = False self.phase = 'train' self.outf = self.opt.experiment_group self.algorithm = 'wgan' # LOAD DATA SET self.dataloader = { 'train': provider('train', opt.category, batch_size=self.batchsize['train'], num_workers=4), 'test': provider('test', opt.category, batch_size=self.batchsize['test'], num_workers=4) } self.trn_dir = os.path.join(self.outf, self.opt.experiment_name, 'train') self.tst_dir = os.path.join(self.outf, self.opt.experiment_name, 'test') self.test_img_dir = os.path.join(self.outf, self.opt.experiment_name, 'test', 'images') if not os.path.isdir(self.test_img_dir): os.makedirs(self.test_img_dir) self.best_test_dir = os.path.join(self.outf, self.opt.experiment_name, 'test', 'best_images') if not os.path.isdir(self.best_test_dir): os.makedirs(self.best_test_dir) self.weight_dir = os.path.join(self.trn_dir, 'weights') if not os.path.exists(self.weight_dir): os.makedirs(self.weight_dir) # -- Misc attributes self.epoch = 0 self.l_con = l1_loss self.l_enc = l2_loss ## # Create and initialize networks. self.netg = NetG().cuda() self.netd = NetD().cuda() # Setup optimizer self.optimizer_d = optim.RMSprop(self.netd.parameters(), lr=self.lr) self.optimizer_g = optim.Adam(self.netg.parameters(), lr=self.lr) ## self.weight_path = os.path.join(self.outf, self.opt.experiment_name, 'train', 'weights') if os.path.exists(self.weight_path) and len( os.listdir(self.weight_path)) == 2: print("Loading pre-trained networks...\n") self.netg.load_state_dict( torch.load(os.path.join(self.weight_path, 'netG.pth'))['state_dict']) self.netd.load_state_dict( torch.load(os.path.join(self.weight_path, 'netD.pth'))['state_dict']) self.optimizer_g.load_state_dict( torch.load(os.path.join(self.weight_path, 'netG.pth'))['optimizer']) self.optimizer_d.load_state_dict( torch.load(os.path.join(self.weight_path, 'netD.pth'))['optimizer']) self.start_iter = torch.load( os.path.join(self.weight_path, 'netG.pth'))['epoch'] ## def start(self): """ Train the model """ ## # TRAIN # self.total_steps = 0 best_criterion = -1 #float('inf') best_auc = -1 # Train for niter epochs. # print(">> Training model %s." % self.name) for self.epoch in range(self.start_iter, self.niter): # Train for one epoch mean_wass = self.train() (auc, res, best_rec, best_threshold), res_total = self.test() message = '' # message += 'criterion: (%.3f+%.3f)/2=%.3f ' % (best_rec[0], best_rec[1], res) # message += 'best threshold: %.3f ' % best_threshold message += 'Wasserstein Distance:%.3d ' % mean_wass message += 'AUC: %.3f ' % auc print(message) torch.save( { 'epoch': self.epoch + 1, 'state_dict': self.netg.state_dict(), 'optimizer': self.optimizer_g.state_dict() }, '%s/netG.pth' % (self.weight_dir)) torch.save( { 'epoch': self.epoch + 1, 'state_dict': self.netd.state_dict(), 'optimizer': self.optimizer_d.state_dict() }, '%s/netD.pth' % (self.weight_dir)) if auc > best_auc: best_auc = auc new_message = "******** New optimal found, saving state ********" message = message + '\n' + new_message print(new_message) for img in os.listdir(self.best_test_dir): os.remove(os.path.join(self.best_test_dir, img)) for img in os.listdir(self.test_img_dir): shutil.copyfile(os.path.join(self.test_img_dir, img), os.path.join(self.best_test_dir, img)) shutil.copyfile('%s/netG.pth' % (self.weight_dir), '%s/netg_best.pth' % (self.weight_dir)) log_name = os.path.join(self.outf, self.opt.experiment_name, 'loss_log.txt') message = 'Epoch%3d:' % self.epoch + ' ' + message with open(log_name, "a") as log_file: if self.epoch == 0: log_file.write('\n\n') log_file.write('%s\n' % message) print(">> Training %s Done..." % self.opt.experiment_name) ## def train(self): """ Train the model for one epoch. """ print("\n>>> Epoch %d/%d, Running " % (self.epoch + 1, self.niter) + self.opt.experiment_name) self.netg.train() self.netd.train() # for p in self.netg.parameters(): p.requires_grad = True mean_wass = 0 tk0 = tqdm(self.dataloader['train'], total=len(self.dataloader['train'])) for i, itr in enumerate(tk0): input, _ = itr input = input.cuda() wasserstein_d = None # if self.algorithm == 'wgan': # train NetD for _ in range(self.netd_niter): # for p in self.netd.parameters(): p.requires_grad = True self.optimizer_d.zero_grad() # forward_g latent_i = torch.rand(self.batchsize['train'], 64, 1, 1).cuda() fake = self.netg(latent_i) # forward_d _, pred_real = self.netd(input) _, pred_fake = self.netd(fake) # .detach() TODO # Backward-pass wasserstein_d = (pred_fake.mean() - pred_real.mean()) * 1 wasserstein_d.backward() self.optimizer_d.step() for p in self.netd.parameters(): p.data.clamp_(-0.01, 0.01) #<<<<<<< # train netg # for p in self.netd.parameters(): p.requires_grad = False self.optimizer_g.zero_grad() noise = torch.rand(self.batchsize['train'], 64, 1, 1).cuda() fake = self.netg(noise) _, pred_fake = self.netd(fake) err_g_d = -pred_fake.mean() # negative err_g_d.backward() self.optimizer_g.step() errors = { 'loss_netD': wasserstein_d.item(), 'loss_netG': round(err_g_d.item(), 3), } mean_wass += wasserstein_d.item() tk0.set_postfix(errors) if i % 50 == 0: img_dir = os.path.join(self.outf, self.opt.experiment_name, 'train', 'images') if not os.path.isdir(img_dir): os.makedirs(img_dir) self.save_image_cv2(input.data, '%s/reals.png' % img_dir) self.save_image_cv2(fake.data, '%s/fakes%03d.png' % (img_dir, i)) mean_wass /= len(self.dataloader['train']) return mean_wass ## def test(self): """ Test AnoGAN model. Args: dataloader ([type]): Dataloader for the test set Raises: IOError: Model weights not found. """ self.netg.eval() self.netd.eval() # for p in self.netg.parameters(): p.requires_grad = False # for p in self.netd.parameters(): p.requires_grad = False for img in os.listdir(self.test_img_dir): os.remove(os.path.join(self.test_img_dir, img)) self.phase = 'test' meter = Meter_AnoGAN() tk1 = tqdm(self.dataloader['test'], total=len(self.dataloader['test'])) for i, itr in enumerate(tk1): input, target = itr input = input.cuda() latent_i = torch.rand(self.batchsize['test'], 64, 1, 1).cuda() latent_i.requires_grad = True optimizer_latent = optim.Adam([latent_i], lr=self.lr) test_loss = None for _ in range(self.test_iter): optimizer_latent.zero_grad() fake = self.netg(latent_i) residual_loss = self.l_con(input, fake) latent_o, _ = self.netd(fake) discrimination_loss = self.l_enc(latent_i, latent_o) alpha = 0.1 test_loss = ( 1 - alpha) * residual_loss + alpha * discrimination_loss test_loss.backward() optimizer_latent.step() abnormal_score = test_loss meter.update(abnormal_score, target) #<<<TODO # Save test images. combine = torch.cat([input.cpu(), fake.cpu()], dim=0) self.save_image_cv2(combine, '%s/%05d.jpg' % (self.test_img_dir, i + 1)) criterion, res_total = meter.get_metrics() # rename images for i, res in enumerate(res_total): os.rename('%s/%05d.jpg' % (self.test_img_dir, i + 1), '%s/%05d_%s.jpg' % (self.test_img_dir, i + 1, res)) return criterion, res_total @staticmethod def save_image_cv2(tensor, filename): # return from torchvision.utils import make_grid # tensor = (tensor + 1) / 2 grid = make_grid(tensor, 8, 2, 0, False, None, False) ndarray = grid.mul_(255).clamp_(0, 255).permute(1, 2, 0).to( 'cpu', torch.uint8).numpy() cv2.imwrite(filename, ndarray)
def generate(**kwargs): """ 随机生成动漫头像,并根据netd的分数选择较好的 """ for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device = torch.device('cuda') if opt.gpu else torch.device('cpu') hash_id = hash(opt.gen_id) % 51603 torch.manual_seed(hash_id) netg, netd = NetG(opt).eval(), NetD(opt).eval() noises = torch.randn(opt.gen_search_num, opt.nz, 1, 1).normal_(opt.gen_mean, opt.gen_std) noises = noises.to(device) map_location = lambda storage, loc: storage netd.load_state_dict(torch.load(opt.netd_path, map_location=map_location)) netg.load_state_dict(torch.load(opt.netg_path, map_location=map_location)) netd.to(device) netg.to(device) # 生成图片,并计算图片在判别器的分数 fake_img = netg(noises) scores = netd(fake_img).detach() # 挑选最好的某几张 indexs = scores.topk(opt.gen_num)[1] result = [] for ii in indexs: result.append(fake_img.data[ii]) # 保存图片 # epoch_num = opt.netd_path[opt.netd_path.find("_") + 1:opt.netd_path.rfind(".")] # pic_stack_name = "".join(["result_", str(epoch_num), "_", str(randid), ".png"]) print(result[0]) print(result[0].size()) # np_array = result[0].cpu().numpy() # # trans_array = np.zeros((96,96,3)) # # trans_array[:,:,0] = np_array[0,:,:] # trans_array[:,:,1] = np_array[1,:,:] # trans_array[:,:,2] = np_array[2,:,:] # # print(trans_array.shape) # # img_ = Image.fromarray(np.uint8(trans_array*255)) # img_ = img_.resize((224, 224), Image.BILINEAR) #resize to # t_img = torch.from_numpy(np.array(img_)) # tv.utils.save_image(t_img, "PIL_test_img.png", normalize=True, range=(-0.8, 0.8)) #tv.utils.save_image(torch.stack(result), "save_test.png", normalize=True, range=(-1, 1)) base_dir = opt.gen_dst base_name = opt.base_name if not os.path.exists(base_dir): os.mkdir(base_dir) for picx in range(len(result)): picx_name = "".join([base_dir, base_name, str(picx), ".png"]) tv.utils.save_image(result[picx], picx_name, normalize=True, range=(-0.8, 0.8)) img = Image.open(picx_name) img = img.resize((188, 188), Image.ANTIALIAS) img.save(base_dir + "/" + str(picx) + ".png", "png") os.remove(picx_name)
def train(): # change opt # for k_, v_ in kwargs.items(): # setattr(opt, k_, v_) device = torch.device('cuda') if torch.cuda.is_available else torch.device( 'cpu') if opt.vis: from visualizer import Visualizer vis = Visualizer(opt.env) # rescale to -1~1 transform = transforms.Compose([ transforms.Resize(opt.image_size), transforms.CenterCrop(opt.image_size), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) dataset = datasets.ImageFolder(opt.data_path, transform=transform) dataloader = DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, drop_last=True) netd = NetD(opt) netg = NetG(opt) map_location = lambda storage, loc: storage if opt.netd_path: netd.load_state_dict(torch.load(opt.netd_path), map_location=map_location) if opt.netg_path: netg.load_state_dict(torch.load(opt.netg_path), map_location=map_location) if torch.cuda.is_available(): netd.to(device) netg.to(device) # 定义优化器和损失 optimizer_g = torch.optim.Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999)) optimizer_d = torch.optim.Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999)) criterion = torch.nn.BCELoss().to(device) # 真label为1, noises是输入噪声 true_labels = Variable(torch.ones(opt.batch_size)) fake_labels = Variable(torch.zeros(opt.batch_size)) fix_noises = Variable(torch.randn(opt.batch_size, opt.nz, 1, 1)) noises = Variable(torch.randn(opt.batch_size, opt.nz, 1, 1)) errord_meter = AverageValueMeter() errorg_meter = AverageValueMeter() if torch.cuda.is_available(): netd.cuda() netg.cuda() criterion.cuda() true_labels, fake_labels = true_labels.cuda(), fake_labels.cuda() fix_noises, noises = fix_noises.cuda(), noises.cuda() for epoch in range(opt.max_epoch): print("epoch:", epoch, end='\r') # sys.stdout.flush() for ii, (img, _) in enumerate(dataloader): real_img = Variable(img) if torch.cuda.is_available(): real_img = real_img.cuda() # 训练判别器, real -> 1, fake -> 0 if (ii + 1) % opt.d_every == 0: # real optimizer_d.zero_grad() output = netd(real_img) # print(output.shape, true_labels.shape) error_d_real = criterion(output, true_labels) error_d_real.backward() # fake noises.data.copy_(torch.randn(opt.batch_size, opt.nz, 1, 1)) fake_img = netg(noises).detach() # 随机噪声生成假图 fake_output = netd(fake_img) error_d_fake = criterion(fake_output, fake_labels) error_d_fake.backward() # update optimizer optimizer_d.step() error_d = error_d_fake + error_d_real errord_meter.add(error_d.item()) # 训练生成器, 让生成器得到的图片能够被判别器判别为真 if (ii + 1) % opt.g_every == 0: optimizer_g.zero_grad() noises.data.copy_(torch.randn(opt.batch_size, opt.nz, 1, 1)) fake_img = netg(noises) fake_output = netd(fake_img) error_g = criterion(fake_output, true_labels) error_g.backward() optimizer_g.step() errorg_meter.add(error_g.item()) if opt.vis and ii % opt.plot_every == opt.plot_every - 1: # 进行可视化 # if os.path.exists(opt.debug_file): # import ipdb # ipdb.set_trace() fix_fake_img = netg(fix_noises) vis.images( fix_fake_img.detach().cpu().numpy()[:opt.batch_size] * 0.5 + 0.5, win='fixfake') vis.images(real_img.data.cpu().numpy()[:opt.batch_size] * 0.5 + 0.5, win='real') vis.plot('errord', errord_meter.value()[0]) vis.plot('errorg', errorg_meter.value()[0]) if (epoch + 1) % opt.save_every == 0: # 保存模型、图片 tv.utils.save_image(fix_fake_img.data[:opt.batch_size], '%s/%s.png' % (opt.save_path, epoch), normalize=True, range=(-1, 1)) torch.save(netd.state_dict(), 'checkpoints/netd_%s.pth' % epoch) torch.save(netg.state_dict(), 'checkpoints/netg_%s.pth' % epoch) errord_meter.reset() errorg_meter.reset()
def train(**kwargs): for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device = t.device('cuda') if opt.gpu else t.device('cpu') # if opt.vis: # from visualize import Visualizer # vis = Visualizer(opt.env) # 数据 transforms = tv.transforms.Compose([ tv.transforms.Resize(opt.image_size), tv.transforms.CenterCrop(opt.image_size), tv.transforms.ToTensor(), tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) dataset = tv.datasets.ImageFolder(root=opt.data_path, transform=transforms) dataloader = t.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, drop_last=True) # 网络 netg, netd = NetG(opt), NetD(opt) map_location = lambda storage, loc: storage if opt.netd_path: netd.load_state_dict(t.load(opt.netd_path, map_location=map_location)) if opt.netg_path: netg.load_state_dict(t.load(opt.netg_path, map_location=map_location)) netd.to(device) netg.to(device) # 定义优化器和损失 optimizer_g = t.optim.Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999)) optimizer_d = t.optim.Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999)) criterion = t.nn.BCELoss().to(device) # 真图片label为1,假图片label为0 # noises为生成网络的输入 true_labels = t.ones(opt.batch_size).to(device) fake_labels = t.zeros(opt.batch_size).to(device) fix_noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device) noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device) errord_meter = AverageValueMeter() errorg_meter = AverageValueMeter() epochs = range(opt.max_epoch) for epoch in iter(epochs): for ii, (img, _) in tqdm.tqdm(enumerate(dataloader)): real_img = img.to(device) if ii % opt.d_every == 0: # 训练判别器 optimizer_d.zero_grad() ## 尽可能的把真图片判别为正确 output = netd(real_img) error_d_real = criterion(output, true_labels) error_d_real.backward() ## 尽可能把假图片判别为错误 noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1)) fake_img = netg(noises).detach() # 根据噪声生成假图 output = netd(fake_img) error_d_fake = criterion(output, fake_labels) error_d_fake.backward() optimizer_d.step() error_d = error_d_fake + error_d_real errord_meter.add(error_d.item()) if ii % opt.g_every == 0: # 训练生成器 optimizer_g.zero_grad() noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1)) fake_img = netg(noises) output = netd(fake_img) error_g = criterion(output, true_labels) error_g.backward() optimizer_g.step() errorg_meter.add(error_g.item()) if opt.vis and ii % opt.plot_every == opt.plot_every - 1: ## 可视化 if os.path.exists(opt.debug_file): ipdb.set_trace() fix_fake_imgs = netg(fix_noises) vis.images(fix_fake_imgs.detach().cpu().numpy()[:64] * 0.5 + 0.5, win='fixfake') vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5, win='real') vis.plot('errord', errord_meter.value()[0]) vis.plot('errorg', errorg_meter.value()[0]) if (epoch + 1) % opt.save_every == 0: # 保存模型、图片 tv.utils.save_image(fix_fake_imgs.data[:64], '%s/%s.png' % (opt.save_path, epoch), normalize=True, range=(-1, 1)) t.save(netd.state_dict(), 'checkpoints/netd_%s.pth' % epoch) t.save(netg.state_dict(), 'checkpoints/netg_%s.pth' % epoch) errord_meter.reset() errorg_meter.reset()
# else: # dataset = TextDataset(cfg.DATA_DIR, 'train', # base_size=cfg.TREE.BASE_SIZE, # transform=image_transform) # print(dataset.n_words, dataset.embeddings_num) # assert dataset # dataloader = torch.utils.data.DataLoader( # dataset, batch_size=batch_size, drop_last=True, # shuffle=True, num_workers=int(cfg.WORKERS)) # # validation data # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") netG = NetG_V2(cfg.TRAIN.NF, 100,cfg.TEXT.EMBEDDING_DIM).cuda() netD = NetD(cfg.TRAIN.NF,cfg.TEXT.EMBEDDING_DIM).cuda() #text_encoder = RNN_ENCODER(dataset.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM) #state_dict = torch.load(cfg.TEXT.DAMSM_NAME, map_location=lambda storage, loc: storage) #text_encoder.load_state_dict(state_dict) #text_encoder.cuda() #for p in text_encoder.parameters(): # p.requires_grad = False #text_encoder.eval() text_encoder = None state_epoch = args.resume_epoch optimizerG = torch.optim.Adam(netG.parameters(), lr=0.0001, betas=(0.0, 0.9)) optimizerD = torch.optim.Adam(netD.parameters(), lr=0.0004, betas=(0.0, 0.9))
transform=transforms.Compose([ transforms.Scale(imgSize), transforms.CenterCrop(imgSize), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])) dataLoader = torch.utils.data.DataLoader(dataSet, batch_size=batchSize, shuffle=True, num_workers=2) dataIter = iter(dataLoader) """model""" netG = NetG(genLayerNum, ngf, nz, up_type=opt.type) print netG netD = NetD(deLayerNum, ndf) print netD z = torch.FloatTensor(batchSize, nz, 1, 1) realData = torch.FloatTensor(batchSize, 3, imgSize, imgSize) label = torch.FloatTensor(batchSize) criterion = nn.BCELoss() realData = Variable(realData) z = Variable(z) label = Variable(label) netG = netG.cuda() netD = netD.cuda() z = torch.FloatTensor(batchSize, nz, 8, 8) realData = torch.FloatTensor(batchSize, 3, imgSize, imgSize) one = torch.FloatTensor([1])
tv.transforms.Resize(opt.image_size), tv.transforms.CenterCrop(opt.image_size), tv.transforms.ToTensor(), tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) dataset = tv.datasets.ImageFolder(opt.data_path, transform=transforms) dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, drop_last=True) netg = NetG(opt) netd = NetD(opt) # 加载已有的网络参数 if opt.netd_path: print('Loading netd...', end='') netd.load_state_dict(torch.load(opt.netd_path)) print('Successful!') if opt.netg_path: print('Loading netg...', end='') netg.load_state_dict(torch.load(opt.netg_path)) print('Successful!') optimizer_g = torch.optim.Adam(netg.parameters(), opt.lr_netg, betas=(opt.beta1, 0.999)) optimizer_d = torch.optim.Adam(netd.parameters(),
project="mnist78", jobtype="training", ) mean = (0.5, ) std = (0.5, ) train_dataset = Dataset(csv_file=CONFIG.train_csv_file, transform=ImageTransform(mean, std)) train_dataloader = DataLoader(train_dataset, batch_size=CONFIG.batch_size, shuffle=True) G = NetG(CONFIG) D = NetD(CONFIG) E = NetE(CONFIG) G.apply(weights_init) D.apply(weights_init) E.apply(weights_init) if not args.no_wandb: # Magic wandb.watch(G, log="all") wandb.watch(D, log="all") wandb.watch(E, log="all") G_update, D_update, E_update = train( G, D,
torchvision.transforms.Scale(opt.imageSize), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) dataset = torchvision.datasets.ImageFolder(opt.data_path, transform=transforms) dataloader = torch.utils.data.DataLoader( dataset=dataset, batch_size=opt.batchSize, shuffle=True, drop_last=True, ) #默认ngf是64,nz是100,ndf是64 netG = NetG(opt.ngf, opt.nz).to(device) netD = NetD(opt.ndf).to(device) criterion = nn.BCELoss() optimizerG = torch.optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizerD = torch.optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) label = torch.FloatTensor(opt.batchSize) real_label = 1 fake_label = 0 for epoch in range(1, opt.epoch + 1): for i, (imgs, _) in enumerate(dataloader):
else: dataset = TextDataset(cfg.DATA_DIR, 'train', base_size=cfg.TREE.BASE_SIZE, transform=image_transform) print(dataset.n_words, dataset.embeddings_num) assert dataset dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, drop_last=True, shuffle=True, num_workers=int(cfg.WORKERS)) # # validation data # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") netG = NetG(cfg.TRAIN.NF, 100).to(device) netD = NetD(cfg.TRAIN.NF).to(device) text_encoder = RNN_ENCODER(dataset.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM) state_dict = torch.load(cfg.TEXT.DAMSM_NAME, map_location=lambda storage, loc: storage) text_encoder.load_state_dict(state_dict) text_encoder.cuda() for p in text_encoder.parameters(): p.requires_grad = False text_encoder.eval() state_epoch=0 optimizerG = torch.optim.Adam(netG.parameters(), lr=0.0001, betas=(0.0, 0.9)) optimizerD = torch.optim.Adam(netD.parameters(), lr=0.0004, betas=(0.0, 0.9))
transforms = torchvision.transforms.Compose([ torchvision.transforms.Scale((opt.imageSize,opt.imageSize)), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) dataset = torchvision.datasets.ImageFolder(opt.data_path, transform=transforms) dataloader = torch.utils.data.DataLoader( dataset=dataset, batch_size=opt.batchSize, shuffle=True, drop_last=True, ) netG = NetG(opt.ngf, opt.nz).to(device) netD = NetD(opt.ndf).to(device) criterion = nn.BCELoss() optimizerG = torch.optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizerD = torch.optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) for epoch in range(1, opt.epoch + 1): for i, (imgs, _) in enumerate(dataloader): # 每次epoch,遍历所有图片,共800个batch # 1,固定生成器G,训练鉴别器D real_label = Variable(torch.ones(opt.batchSize)).cuda() fake_label = Variable(torch.zeros(opt.batchSize)).cuda() netD.zero_grad() # 让D尽可能的把真图片判别为1
base_size=cfg.TREE.BASE_SIZE, transform=image_transform) print(dataset.n_words, dataset.embeddings_num) assert dataset dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, drop_last=True, shuffle=True, num_workers=int(cfg.WORKERS)) # # validation data # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") netG = NetG(cfg.TRAIN.NF, 100).cuda() netD = NetD(cfg.TRAIN.NF).cuda() text_encoder = RNN_ENCODER(dataset.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM) state_dict = torch.load(cfg.TEXT.DAMSM_NAME, map_location=lambda storage, loc: storage) text_encoder.load_state_dict(state_dict) text_encoder.cuda() for p in text_encoder.parameters(): p.requires_grad = False text_encoder.eval() state_epoch = args.resume_epoch optimizerG = torch.optim.Adam(netG.parameters(), lr=0.0001,
def train(**kwargs): for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device = t.device('cuda') if opt.gpu else t.device('cpu') # 可视化 if opt.vis: from visualize import Visualizer vis = Visualizer(opt.env) # 数据 transforms = tv.transforms.Compose([ tv.transforms.Resize(opt.image_size), tv.transforms.CenterCrop(opt.image_size), tv.transforms.ToTensor(), tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) dataset = tv.datasets.ImageFolder(opt.data_path, transforms) dataloader = t.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers, drop_last=True) netg, netd = NetG(opt), NetD(opt) if opt.netd_path: netd.load_state_dict( t.load(opt.netd_path, map_location=t.device('cpu'))) if opt.netg_path: netg.load_state_dict( t.load(opt.netg_path, map_location=t.device('cpu'))) netd.to(device) netg.to(device) # 定义优化器和损失 optimizer_g = t.optim.Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999)) optimizer_d = t.optim.Adam(netd.parameters(), opt.lr1, betas=(opt.beta1, 0.999)) criterion = t.nn.BCELoss().to(device) # 真图片label为1,假图片label为0 # noise为生成网络的输入 true_labels = t.ones(opt.batch_size).to(device) fake_labels = t.zeros(opt.batch_size).to(device) fix_noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device) noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device) errord = 0 errorg = 0 epochs = range(opt.max_epoch) for epoch in epochs: for ii, (img, _) in tqdm(enumerate(dataloader)): real_img = img.to(device) if (ii + 1) % opt.d_every == 0: # 训练判别器 optimizer_d.zero_grad() # 把真图片判断为正确 output = netd(real_img) error_d_real = criterion(output, true_labels) error_d_real.backward() # 把假图片(netg通过噪声生成的图片)判断为错误 noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1)) fake_img = netg(noises).detach() output = netd(fake_img) error_d_fake = criterion(output, fake_labels) error_d_fake.backward() optimizer_d.step() errord += (error_d_fake + error_d_real).item() if (ii + 1) % opt.g_every == 0: # 训练生成器 optimizer_g.zero_grad() noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1)) fake_img = netg(noises) output = netd(fake_img) error_g = criterion(output, true_labels) error_g.backward() optimizer_g.step() errorg += error_g.item() if opt.vis and (ii + 1) % opt.plot_every == 0: fix_fake_imgs = netg(fix_noises) vis.images(fix_fake_imgs.detach().cpu().numpy()[:64] * 0.5 + 0.5, win='fixfake') vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5, win='real') vis.plot('errord', errord / (opt.plot_every)) vis.plot('errorg', errorg / (opt.plot_every)) errord = 0 errorg = 0 if (epoch + 1) % opt.save_every == 0: tv.utils.save_image(fix_fake_imgs.data[:64], '%s%s.png' % (opt.save_path, epoch), normalize=True, range=(-1, 1)) t.save(netd.state_dict(), 'checkpoints/netd_%s.pth' % (epoch + 1)) t.save(netg.state_dict(), 'checkpoints/netg_%s.pth' % (epoch + 1))
def __init__(self, opt): # super(AnoGAN, self).__init__(opt, dataloader) # Initalize variables. self.opt = opt self.niter = self.opt.niter self.start_iter = 0 self.netd_niter = 5 self.test_iter = 100 self.lr = self.opt.lr self.batchsize = {'train': self.opt.batchsize, 'test': 1} self.pretrained = False self.phase = 'train' self.outf = self.opt.experiment_group self.algorithm = 'wgan' # LOAD DATA SET self.dataloader = { 'train': provider('train', opt.category, batch_size=self.batchsize['train'], num_workers=4), 'test': provider('test', opt.category, batch_size=self.batchsize['test'], num_workers=4) } self.trn_dir = os.path.join(self.outf, self.opt.experiment_name, 'train') self.tst_dir = os.path.join(self.outf, self.opt.experiment_name, 'test') self.test_img_dir = os.path.join(self.outf, self.opt.experiment_name, 'test', 'images') if not os.path.isdir(self.test_img_dir): os.makedirs(self.test_img_dir) self.best_test_dir = os.path.join(self.outf, self.opt.experiment_name, 'test', 'best_images') if not os.path.isdir(self.best_test_dir): os.makedirs(self.best_test_dir) self.weight_dir = os.path.join(self.trn_dir, 'weights') if not os.path.exists(self.weight_dir): os.makedirs(self.weight_dir) # -- Misc attributes self.epoch = 0 self.l_con = l1_loss self.l_enc = l2_loss ## # Create and initialize networks. self.netg = NetG().cuda() self.netd = NetD().cuda() # Setup optimizer self.optimizer_d = optim.RMSprop(self.netd.parameters(), lr=self.lr) self.optimizer_g = optim.Adam(self.netg.parameters(), lr=self.lr) ## self.weight_path = os.path.join(self.outf, self.opt.experiment_name, 'train', 'weights') if os.path.exists(self.weight_path) and len( os.listdir(self.weight_path)) == 2: print("Loading pre-trained networks...\n") self.netg.load_state_dict( torch.load(os.path.join(self.weight_path, 'netG.pth'))['state_dict']) self.netd.load_state_dict( torch.load(os.path.join(self.weight_path, 'netD.pth'))['state_dict']) self.optimizer_g.load_state_dict( torch.load(os.path.join(self.weight_path, 'netG.pth'))['optimizer']) self.optimizer_d.load_state_dict( torch.load(os.path.join(self.weight_path, 'netD.pth'))['optimizer']) self.start_iter = torch.load( os.path.join(self.weight_path, 'netG.pth'))['epoch']
base_size=cfg.TREE.BASE_SIZE, transform=image_transform) print(dataset.n_words, dataset.embeddings_num) assert dataset dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, drop_last=True, shuffle=True, num_workers=int(cfg.WORKERS)) # # validation data # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") netG = NetG(cfg.TRAIN.NF, 100).to(device) netD = NetD(cfg.TRAIN.NF).to(device) text_encoder = RNN_ENCODER(dataset.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM) state_dict = torch.load(cfg.TEXT.DAMSM_NAME, map_location=lambda storage, loc: storage) text_encoder.load_state_dict(state_dict) text_encoder.cuda() for p in text_encoder.parameters(): p.requires_grad = False text_encoder.eval() state_epoch = 0 optimizerG = torch.optim.Adam(netG.parameters(), lr=0.0001,
transforms = torchvision.transforms.Compose([ torchvision.transforms.Resize(opt.imageSize), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) dataset = torchvision.datasets.ImageFolder(opt.data_path, transform=transforms) dataloader = torch.utils.data.DataLoader( dataset=dataset, batch_size=opt.batchSize, shuffle=True, drop_last=True, ) netG = NetG(opt.ngf, opt.nz).to(device) netD = NetD(opt.ndf).to(device) criterion = nn.BCELoss() optimizerG = torch.optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizerD = torch.optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) label = torch.FloatTensor(opt.batchSize) real_label = 1 fake_label = 0 if not os.path.exists(opt.outf): os.makedirs(opt.outf) if not os.path.exists(opt.model_path): os.makedirs(opt.model_path) for epoch in range(1, opt.epoch + 1):
torchvision.transforms.Scale(opt.imageSize), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) dataset = torchvision.datasets.ImageFolder(opt.data_path, transform=transforms) dataloader = torch.utils.data.DataLoader( dataset=dataset, batch_size=opt.batchSize, shuffle=True, drop_last=True, ) netG = NetG(opt.ngf, opt.nz).to(device) netD = NetD(opt.ndf).to(device) criterion = nn.BCELoss() optimizerG = torch.optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) optimizerD = torch.optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) label = torch.FloatTensor(opt.batchSize) real_label = 1 fake_label = 0 for epoch in range(1, opt.epoch + 1): for i, (imgs, _) in enumerate(dataloader):
import cv2 import torch as th from model import NetG, NetD, NetA from my_dataSet import CASIABDatasetGenerate netg = NetG(nc=1) netd = NetD(nc=1) neta = NetA(nc=1) device = th.device("cuda:0") netg = netg.to(device) netd = netd.to(device) neta = neta.to(device) fineSize = 64 checkpoint = '/home/mg/code/my_GAN_dataSet/snapshots/snapshot_449.t7' checkpoint = th.load(checkpoint) neta.load_state_dict(checkpoint['netA']) netg.load_state_dict(checkpoint['netG']) netd.load_state_dict(checkpoint['netD']) neta.eval() netg.eval() netd.eval() angles = [ '000', '018', '036', '054', '072', '090', '108', '126', '144', '162', '180' ] for cond in ['nm-01', 'nm-02', 'nm-03', 'nm-04', 'cl-01', 'cl-02']: dataset = CASIABDatasetGenerate( data_dir='/home/mg/code/data/GEI_CASIA_B/gei/', cond=cond)
import torch as th import itertools import torch.nn as nn import torch.nn.functional as F import torch.nn.init as init from model import NetG, NetD, NetA from data_set import CASIABDataset import torch.optim as optim import visdom from torchvision.utils import make_grid vis = visdom.Visdom(port=5274) win = None win1 = None netg = NetG(nc=1) netd = NetD(nc=1) neta = NetA(nc=1) device = th.device("cuda:2") # weights init all_mods = itertools.chain() all_mods = itertools.chain(all_mods, [ list(netg.children())[0].children(), list(netd.children())[0].children(), list(neta.children())[0].children() ]) for mod in all_mods: if isinstance(mod, nn.Conv2d) or isinstance(mod, nn.ConvTranspose2d): init.normal_(mod.weight, 0.0, 0.02) elif isinstance(mod, nn.BatchNorm2d): init.normal_(mod.weight, 1.0, 0.02)