def generate(**kwargs):
'''
随机生成动漫头像，并根据netd的分数选择较好的
'''
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
D = NetD(opt)
G = NetG(opt)
noises = torch.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
D.load_state_dict(torch.load(opt.netd_path, map_location=map_location))
G.load_state_dict(torch.load(opt.netg_path, map_location=map_location))
if torch.cuda.is_available():
D.cuda()
G.cuda()
noises = noises.cuda()
# 生成图片，并计算图片在判别器的分数
fake_img = G(noises)
scores = D(fake_img).data
# 挑选最好的某几张
indexs = scores.topk(opt.gen_num)[1]
result = []
for idx in indexs:
result.append(fake_img.data[idx])
# 保存图片
torchvision.utils.save_image(torch.stack(result),
opt.gen_img,
normalize=True,
range=(-1, 1))
transform = transforms.Compose([
transforms.Scale(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 = 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)
if opt.netd_path:
netd.load_state_dict(torch.load(opt.netd_path, map_location=lambda storage, loc: storage))
if opt.netg_path:
netg.load_state_dict(torch.load(opt.netg_path, map_location=lambda storage, loc: storage))
optimizer_g = Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999))
optimizer_d = Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999))
criterion = nn.BCELoss()
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))
if opt.use_gpu:
netd.cuda()
def train(**kwargs):
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(opt.image_size),
torchvision.transforms.CenterCrop(opt.image_size),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5))
])
dataset = torchvision.datasets.ImageFolder(opt.data_path,
transform=transforms)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
drop_last=True)
# 1、定义神经网络
D = NetD(opt)
G = NetG(opt)
map_location = lambda storage, loc: storage
if opt.netd_path:
D.load_state_dict(torch.load(opt.netd_path, map_location=map_location))
if opt.netg_path:
G.load_state_dict(torch.load(opt.netg_path, map_location=map_location))
# 2、定义优化器和损失
d_optim = torch.optim.Adam(D.parameters(),
opt.d_learning_rate,
betas=(opt.optim_beta1, 0.999))
g_optim = torch.optim.Adam(G.parameters(),
opt.g_learning_rate,
betas=(opt.optim_beta1, 0.999))
criterion = torch.nn.BCELoss()
# 真图片label为1，假图片label为0
real_labels = Variable(torch.ones(opt.batch_size))
fake_labels = Variable(torch.zeros(opt.batch_size))
if torch.cuda.is_available():
D.cuda()
G.cuda()
criterion.cuda()
real_labels, fake_labels = real_labels.cuda(), fake_labels.cuda()
# 3、可视化训练过程
for epoch in range(opt.num_epochs):
for step, (images, _) in tqdm.tqdm(enumerate(dataloader)):
if step % opt.d_every == 0:
# 1、训练判别器
d_optim.zero_grad()
## 尽可能的把真图片判别为正确
d_real_data = Variable(images)
d_real_data = d_real_data.cuda() if torch.cuda.is_available(
) else d_real_data
d_real_decision = D(d_real_data)
d_real_error = criterion(d_real_decision, real_labels)
d_real_error.backward()
## 尽可能把假图片判别为错误
d_gen_input = Variable(
torch.randn(opt.batch_size, opt.nz, 1, 1))
d_gen_input = d_gen_input.cuda() if torch.cuda.is_available(
) else d_gen_input
d_fake_data = G(d_gen_input).detach()
d_fake_decision = D(d_fake_data)
d_fake_error = criterion(d_fake_decision, fake_labels)
d_fake_error.backward()
d_optim.step(
) # Only optimizes D's parameters; changes based on stored gradients from backward()
if step % opt.g_every == 0:
# 2、训练生成器
g_optim.zero_grad()
## 尽可能让判别器把假图片判别为正确
g_gen_input = Variable(
torch.randn(opt.batch_size, opt.nz, 1, 1))
g_gen_input = g_gen_input.cuda() if torch.cuda.is_available(
) else g_gen_input
g_fake_data = G(g_gen_input)
g_fake_decision = D(g_fake_data)
g_fake_error = criterion(g_fake_decision, real_labels)
g_fake_error.backward()
g_optim.step()
if step % opt.epoch_every == 0:
print("%s, %s, D: %s/%s G: %s" %
(step, g_fake_decision.cpu().data.numpy().mean(),
d_real_error.cpu().data[0], d_fake_error.cpu().data[0],
g_fake_error.cpu().data[0]))
# 保存模型、图片
torchvision.utils.save_image(g_fake_data.data[:36],
'%s/%s.png' %
(opt.save_img_path, epoch),
normalize=True,
range=(-1, 1))
torch.save(D.state_dict(),
'%s/netd_%s.pth' % (opt.checkpoints_path, epoch))
torch.save(G.state_dict(),
'%s/netg_%s.pth' % (opt.checkpoints_path, epoch))
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = 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)
if opt.netd_path:
netd.load_state_dict(
torch.load(opt.netd_path, map_location=lambda storage, loc: storage))
if opt.netg_path:
netg.load_state_dict(
torch.load(opt.netg_path, map_location=lambda storage, loc: storage))
optimizer_g = Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999))
optimizer_d = Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999))
criterion = nn.BCELoss()
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))
if opt.use_gpu: