errD_real = criterion(output, label)
errD_real.backward()
label.data.fill_(fake_label)
noise = torch.randn(opt.batchSize, opt.nz, 1, 1)
noise = noise.to(device)
fake = Generator(noise)
output = Discriminator(fake.detach()) #截断G反向传播的梯度流
errD_fake = criterion(output, label)
errD_fake.backward()
errD = errD_fake + errD_real
optimizerD.step()
optimizerG.zero_grad()
label.data.fill_(real_label)
label = label.to(device)
output = Discriminator(fake)
errG = criterion(output, label)
errG.backward()
optimizerG.step()
print('[%d/%d][%d/%d] Loss_D: %.3f Loss_G %.3f' %
(epoch, opt.epoch, i, len(dataloader), errD.item(), errG.item()))
vutils.save_image(fake.data,
'%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch),
normalize=True)
torch.save(Generator.state_dict(),
'%s/netG_%03d.pth' % (opt.outf, epoch))
torch.save(Discriminator.state_dict(),
'%s/netD_%03d.pth' % (opt.outf, epoch))
def train(Generator, Discriminator, dataloader, criterion):
# setup optimizer
g_optim = torch.optim.Adam(Generator.parameters(),
lr=learning_rate,
betas=(0.5, 0.999))
d_optim = torch.optim.Adam(Discriminator.parameters(),
lr=learning_rate,
betas=(0.5, 0.999))
# initialize noise
noise = torch.randn(batch_size, feature_dim, 1, 1)
test_noise = torch.randn(batch_size, feature_dim, 1, 1)
true_label = torch.FloatTensor(batch_size).fill_(1 - smooth_label)
true_label_g = torch.FloatTensor(batch_size).fill_(1)
fake_label = torch.FloatTensor(batch_size).fill_(0 + smooth_label)
if args.cuda:
noise = noise.cuda()
test_noise = test_noise.cuda()
true_label = true_label.cuda()
true_label_g = true_label_g.cuda()
fake_label = fake_label.cuda()
# train!
Generator.train()
Discriminator.train()
for epoch in range(num_epochs):
for i, data in enumerate(dataloader):
img = data['img']
# label = data['label']
if args.cuda:
img = img.cuda()
# label = label.cuda()
img_real = Variable(img)
# label = Variable(label)
#------------------------#
# Train Discriminator
#------------------------#
Discriminator.zero_grad()
# learn real data as real
true_label_var = Variable(true_label)
out_real = Discriminator(img_real)
loss_d_real = criterion(out_real, true_label_var)
# learn fake data as fake
noise_var = Variable(noise)
# generate fake data with generator
img_fake = Generator(noise_var)
fake_label_var = Variable(fake_label)
out_fake = Discriminator(img_fake.detach())
loss_d_fake = criterion(out_fake, fake_label_var)
loss_d = loss_d_real + loss_d_fake
loss_d.backward()
d_optim.step()
#-------------------#
# Train Generator
#-------------------#
Generator.zero_grad()
# fool discriminator to learn as real with fake data
true_label_g_var = Variable(true_label_g)
out_real_but_fake = Discriminator(img_fake)
loss_g = criterion(out_real_but_fake, true_label_g_var)
loss_g.backward()
g_optim.step()
print(
"Epoch [%d/%d] Iter [%d/%d] Loss D : %.4f, Loss G : %.4f, D(x) : %.4f, D(z) : %.4f, g : %.4f"
% (epoch + 1, num_epochs, i + 1, len(dataloader),
loss_d.data[0], loss_g.data[0], loss_d_real.data.mean(),
loss_d_fake.data.mean(), loss_g.data.mean()))
niter = epoch * len(dataloader) + i + 1
writer.add_scalar('Loss/D', loss_d.data[0], niter)
writer.add_scalar('Loss/G', loss_g.data[0], niter)
writer.add_scalar('D/D(x)', loss_d_real.data.mean(), niter)
writer.add_scalar('D/D(z)', loss_d_fake.data.mean(), niter)
writer.add_scalar('D/g', loss_g.data.mean(), niter)
# generate on the way
if (i + 1) % 100 == 0:
test_noise_var = Variable(test_noise)
test_img = Generator(test_noise_var)
vutils.save_image(test_img.data,
base_dir + 'fake_img_epoch_%d_iter_%d.png' %
(epoch + 1, i + 1),
normalize=True)
writer.add_image(
'fake_images',
vutils.make_grid(test_img.data, normalize=True), niter)
# save model
torch.save(Generator.state_dict(), 'Generator.pth')
torch.save(Discriminator.state_dict(), 'Discriminator.pth')
writer.close()
def train():
opt = parse_args()
os.makedirs("images/%s" % (opt.dataset), exist_ok=True)
os.makedirs("checkpoints/%s" % (opt.dataset), exist_ok=True)
cuda = True if torch.cuda.is_available() else False
FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
# get dataloader
train_loader = celeba_loader(opt, mode='train')
val_loader = celeba_loader(opt, mode='val')
# Dimensionality
c_dim = len(opt.selected_attrs)
# Initialize generator and discriminator
generator = Generator(opt.channels, opt.residual_blocks, c_dim)
discriminator = Discriminator(opt.channels, opt.img_height, c_dim)
# Initialize weights
generator.apply(weights_init_normal)
discriminator.apply(weights_init_normal)
# Loss function
cycle_loss = torch.nn.L1Loss()
if cuda:
generator = generator.cuda()
discriminator = discriminator.cuda()
cycle_loss.cuda()
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(),
lr=opt.lr,
betas=(opt.b1, opt.b2))
optimizer_D = torch.optim.Adam(discriminator.parameters(),
lr=opt.lr,
betas=(opt.b1, opt.b2))
# ------------
# Training
# ------------
prev_time = time.time()
for epoch in range(opt.epochs):
for i, (imgs, labels) in enumerate(train_loader):
# Model inputs
imgs = Variable(imgs.type(FloatTensor))
labels = Variable(labels.type(FloatTensor))
# Sample label as generator inputs and Generate fake batch of images
sampled_c = Variable(
FloatTensor(np.random.randint(0, 2, (imgs.size(0), c_dim))))
fake_imgs = generator(imgs, sampled_c)
# ----------------------
# Train Discriminator
# ----------------------
optimizer_D.zero_grad()
real_validity, pred_cls = discriminator(imgs)
fake_validity, _ = discriminator(fake_imgs.detach())
gradient_penalty = compute_gradient_penalty(
discriminator, imgs.data, fake_imgs.data, FloatTensor)
d_adv_loss = -torch.mean(real_validity) + torch.mean(
fake_validity) + opt.lambda_gp * gradient_penalty
d_cls_loss = criterion_cls(pred_cls, labels)
D_loss = d_adv_loss + opt.lambda_cls * d_cls_loss
D_loss.backward()
optimizer_D.step()
# -----------------------------
# Train Generators
# -----------------------------
optimizer_G.zero_grad()
if i % opt.n_critic == 0:
gen_imgs = generator(imgs, sampled_c)
recov_imgs = generator(gen_imgs, labels)
fake_validity, pred_cls = discriminator(gen_imgs)
g_adv_loss = -torch.mean(fake_validity)
g_cls_loss = criterion_cls(pred_cls, sampled_c)
g_rec_loss = cycle_loss(recov_imgs, imgs)
G_loss = g_adv_loss + opt.lambda_cls * g_cls_loss + opt.lambda_rec * g_rec_loss
G_loss.backward()
optimizer_G.step()
# ------------------
# Log Information
# ------------------
batches_done = epoch * len(train_loader) + i
batches_left = opt.epochs * len(train_loader) - batches_done
time_left = datetime.timedelta(seconds=batches_left *
(time.time() - prev_time))
prev_time = time.time()
print(
"[Epoch %d/%d] [Batch %d/%d] [D loss: %f, aux: %f] [G loss: %f, aux: %f, cycle: %f] ETA: %s"
% (epoch, opt.epochs, i, len(train_loader), D_loss.item(),
d_cls_loss.item(), G_loss.item(), g_cls_loss.item(),
g_rec_loss, time_left))
if batches_done % opt.sample_interval == 0:
save_sample(opt.dataset, val_loader, batches_done,
generator, FloatTensor)
if batches_done % opt.checkpoint_interval == 0:
torch.save(
Generator.state_dict(),
"checkpoints/%s/G_%d.pth" % (opt.dataset, epoch))
torch.save(Generator.state_dict(),
"checkpoints/%s/shared_E_done.pth" % opt.dataset)
print("Training Process has been Done!")
