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)
