def __init__(self, args):
print("WGAN_CP init model.")
self.G = Generator(args.channels)
self.D = Discriminator(args.channels)
self.C = args.channels
# check if cuda is available
self.check_cuda(args.cuda)
# WGAN values from paper
self.learning_rate = 0.00005
self.batch_size = 64
self.weight_cliping_limit = 0.01
# WGAN with gradient clipping uses RMSprop instead of ADAM
self.d_optimizer = torch.optim.RMSprop(self.D.parameters(), lr=self.learning_rate)
self.g_optimizer = torch.optim.RMSprop(self.G.parameters(), lr=self.learning_rate)
# Set the logger
self.logger = Logger('./logs')
self.logger.writer.flush()
self.number_of_images = 10
self.generator_iters = args.generator_iters
self.critic_iter = 5
def __init__(self, args):
print("VLGAN model initalization.")
self.G = Generator(args.channels)
self.D = Discriminator(args.channels)
self.E = Encoder(args.channels)
self.C = args.channels
self.recon_weight = 5
# binary cross entropy loss and optimizer
self.loss = nn.BCELoss()
self.v_loss = Variational_loss
self.cuda = "False"
self.cuda_index = 0
# check if cuda is available
self.check_cuda(args.cuda)
# Using lower learning rate than suggested by (ADAM authors) lr=0.0002 and Beta_1 = 0.5 instead od 0.9 works better [Radford2015]
self.d_optimizer = torch.optim.Adam(self.D.parameters(),
lr=0.0002,
betas=(0.5, 0.999))
self.g_optimizer = torch.optim.Adam(self.G.parameters(),
lr=0.0002,
betas=(0.5, 0.999))
self.e_optimizer = torch.optim.Adam(self.E.parameters(),
lr=0.0002,
betas=(0.5, 0.999))
self.epochs = args.epochs
self.batch_size = args.batch_size
# Set the logger
self.logger = Logger('./logs')
self.number_of_images = 10
def __init__(self, args):
# Generator architecture
self.G = nn.Sequential(nn.Linear(100, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 512), nn.LeakyReLU(0.2),
nn.Linear(512, 1024), nn.LeakyReLU(0.2),
nn.Tanh())
# Discriminator architecture
self.D = nn.Sequential(nn.Linear(1024, 512), nn.LeakyReLU(0.2),
nn.Linear(512, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 1), nn.Sigmoid())
# Check if cuda is available
self.check_cuda(args.cuda)
# Binary cross entropy loss and optimizer
self.loss = nn.BCELoss()
self.d_optimizer = torch.optim.Adam(self.D.parameters(),
lr=0.0002,
weight_decay=0.00001)
self.g_optimizer = torch.optim.Adam(self.G.parameters(),
lr=0.0002,
weight_decay=0.00001)
# Set the logger
self.logger = Logger('./logs')
self.number_of_images = 10
self.epochs = args.epochs
self.batch_size = args.batch_size
def __init__(self, args):
print("WGAN_GradientPenalty init model.")
self.G = Generator(args.channels)
self.D = Discriminator(args.channels)
self.C = args.channels
# Check if cuda is available
self.check_cuda(args.cuda)
# WGAN values from paper
self.learning_rate = 1e-4
self.b1 = 0.5
self.b2 = 0.999
self.batch_size = args.batch_size
# WGAN_gradient penalty uses ADAM
self.d_optimizer = optim.Adam(self.D.parameters(), lr=self.learning_rate, betas=(self.b1, self.b2))
self.g_optimizer = optim.Adam(self.G.parameters(), lr=self.learning_rate, betas=(self.b1, self.b2))
# Set the logger
self.logger = Logger('./logs')
self.logger.writer.flush()
self.number_of_images = 10
self.generator_iters = args.generator_iters
self.critic_iter = 5
self.lambda_term = 10
class WGAN_CP(object):
def __init__(self, args):
print("WGAN_CP init model.")
self.G = Generator(args.channels)
self.D = Discriminator(args.channels)
self.C = args.channels
# check if cuda is available
self.check_cuda(args.cuda)
# WGAN values from paper
self.learning_rate = 0.00005
self.batch_size = 64
self.weight_cliping_limit = 0.01
# WGAN with gradient clipping uses RMSprop instead of ADAM
self.d_optimizer = torch.optim.RMSprop(self.D.parameters(), lr=self.learning_rate)
self.g_optimizer = torch.optim.RMSprop(self.G.parameters(), lr=self.learning_rate)
# Set the logger
self.logger = Logger('./logs')
self.logger.writer.flush()
self.number_of_images = 10
self.generator_iters = args.generator_iters
self.critic_iter = 5
def check_cuda(self, cuda_flag=False):
if cuda_flag:
self.cuda_index = 0
self.cuda = True
self.D.cuda()
self.G.cuda()
print("Cuda enabled flag: {}".format(self.cuda))
def train(self, train_loader):
self.t_begin = t.time()
#self.file = open("inception_score_graph.txt", "w")
# Now batches are callable self.data.next()
self.data = self.get_infinite_batches(train_loader)
one = torch.FloatTensor([1])
mone = one * -1
if self.cuda:
one = one.cuda()
mone = mone.cuda()
for g_iter in range(self.generator_iters):
# Requires grad, Generator requires_grad = False
for p in self.D.parameters():
p.requires_grad = True
# Train Dicriminator forward-loss-backward-update self.critic_iter times while 1 Generator forward-loss-backward-update
for d_iter in range(self.critic_iter):
self.D.zero_grad()
# Clamp parameters to a range [-c, c], c=self.weight_cliping_limit
for p in self.D.parameters():
p.data.clamp_(-self.weight_cliping_limit, self.weight_cliping_limit)
images = self.data.__next__()
# Check for batch to have full batch_size
if (images.size()[0] != self.batch_size):
continue
z = torch.rand((self.batch_size, 100, 1, 1))
if self.cuda:
images, z = Variable(images.cuda()), Variable(z.cuda())
else:
images, z = Variable(images), Variable(z)
# Train discriminator
# WGAN - Training discriminator more iterations than generator
# Train with real images
d_loss_real = self.D(images)
d_loss_real = d_loss_real.mean(0).view(1)
d_loss_real.backward(one)
# Train with fake images
if self.cuda:
z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda()
else:
z = Variable(torch.randn(self.batch_size, 100, 1, 1))
fake_images = self.G(z)
d_loss_fake = self.D(fake_images)
d_loss_fake = d_loss_fake.mean(0).view(1)
d_loss_fake.backward(mone)
d_loss = d_loss_fake - d_loss_real
Wasserstein_D = d_loss_real - d_loss_fake
self.d_optimizer.step()
# Generator update
for p in self.D.parameters():
p.requires_grad = False # to avoid computation
self.G.zero_grad()
# Train generator
# Compute loss with fake images
z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda()
fake_images = self.G(z)
g_loss = self.D(fake_images)
g_loss = g_loss.mean().mean(0).view(1)
g_loss.backward(one)
g_cost = -g_loss
self.g_optimizer.step()
# Saving model and sampling images every 1000th generator iterations
if (g_iter) % 1000 == 0:
self.save_model()
# Workaround because graphic card memory can't store more than 830 examples in memory for generating image
# Therefore doing loop and generating 800 examples and stacking into list of samples to get 8000 generated images
# This way Inception score is more correct since there are different generated examples from every class of Inception model
# sample_list = []
# for i in range(10):
# z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index)
# samples = self.G(z)
# sample_list.append(samples.data.cpu().numpy())
#
# # Flattening list of list into one list
# new_sample_list = list(chain.from_iterable(sample_list))
# print("Calculating Inception Score over 8k generated images")
# # Feeding list of numpy arrays
# inception_score = get_inception_score(new_sample_list, cuda=True, batch_size=32,
# resize=True, splits=10)
if not os.path.exists('training_result_images/'):
os.makedirs('training_result_images/')
# Denormalize images and save them in grid 8x8
z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index)
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()[:64]
grid = utils.make_grid(samples)
utils.save_image(grid, 'training_result_images/img_generatori_iter_{}.png'.format(str(g_iter).zfill(3)))
# Testing
time = t.time() - self.t_begin
#print("Inception score: {}".format(inception_score))
print("Generator iter: {}".format(g_iter))
print("Time {}".format(time))
# Write to file inception_score, gen_iters, time
#output = str(g_iter) + " " + str(time) + " " + str(inception_score[0]) + "\n"
#self.file.write(output)
# ============ TensorBoard logging ============#
# (1) Log the scalar values
info = {
'Wasserstein distance': Wasserstein_D.data[0],
'Loss D': d_loss.data[0],
'Loss G': g_cost.data[0],
'Loss D Real': d_loss_real.data[0],
'Loss D Fake': d_loss_fake.data[0]
}
for tag, value in info.items():
self.logger.scalar_summary(tag, value, g_iter + 1)
# (3) Log the images
info = {
'real_images': self.real_images(images, self.number_of_images),
'generated_images': self.generate_img(z, self.number_of_images)
}
for tag, images in info.items():
self.logger.image_summary(tag, images, g_iter + 1)
self.t_end = t.time()
print('Time of training-{}'.format((self.t_end - self.t_begin)))
#self.file.close()
# Save the trained parameters
self.save_model()
def evaluate(self, test_loader, D_model_path, G_model_path):
self.load_model(D_model_path, G_model_path)
z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda()
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()
grid = utils.make_grid(samples)
print("Grid of 8x8 images saved to 'dgan_model_image.png'.")
utils.save_image(grid, 'dgan_model_image.png')
def real_images(self, images, number_of_images):
if (self.C == 3):
return self.to_np(images.view(-1, self.C, 32, 32)[:self.number_of_images])
else:
return self.to_np(images.view(-1, 32, 32)[:self.number_of_images])
def generate_img(self, z, number_of_images):
samples = self.G(z).data.cpu().numpy()[:number_of_images]
generated_images = []
for sample in samples:
if self.C == 3:
generated_images.append(sample.reshape(self.C, 32, 32))
else:
generated_images.append(sample.reshape(32, 32))
return generated_images
def to_np(self, x):
return x.data.cpu().numpy()
def save_model(self):
torch.save(self.G.state_dict(), './generator.pkl')
torch.save(self.D.state_dict(), './discriminator.pkl')
print('Models save to ./generator.pkl & ./discriminator.pkl ')
def load_model(self, D_model_filename, G_model_filename):
D_model_path = os.path.join(os.getcwd(), D_model_filename)
G_model_path = os.path.join(os.getcwd(), G_model_filename)
self.D.load_state_dict(torch.load(D_model_path))
self.G.load_state_dict(torch.load(G_model_path))
print('Generator model loaded from {}.'.format(G_model_path))
print('Discriminator model loaded from {}-'.format(D_model_path))
def get_infinite_batches(self, data_loader):
while True:
for i, (images, _) in enumerate(data_loader):
yield images
def generate_latent_walk(self, number):
if not os.path.exists('interpolated_images/'):
os.makedirs('interpolated_images/')
number_int = 10
# interpolate between twe noise(z1, z2).
z_intp = torch.FloatTensor(1, 100, 1, 1)
z1 = torch.randn(1, 100, 1, 1)
z2 = torch.randn(1, 100, 1, 1)
if self.cuda:
z_intp = z_intp.cuda()
z1 = z1.cuda()
z2 = z2.cuda()
z_intp = Variable(z_intp)
images = []
alpha = 1.0 / float(number_int + 1)
print(alpha)
for i in range(1, number_int + 1):
z_intp.data = z1*alpha + z2*(1.0 - alpha)
alpha += alpha
fake_im = self.G(z_intp)
fake_im = fake_im.mul(0.5).add(0.5) #denormalize
images.append(fake_im.view(self.C,32,32).data.cpu())
grid = utils.make_grid(images, nrow=number_int )
utils.save_image(grid, 'interpolated_images/interpolated_{}.png'.format(str(number).zfill(3)))
print("Saved interpolated images.")
class VLGAN_MODEL(object):
def __init__(self, args):
print("VLGAN model initalization.")
self.G = Generator(args.channels)
self.D = Discriminator(args.channels)
self.E = Encoder(args.channels)
self.C = args.channels
self.recon_weight = 5
# binary cross entropy loss and optimizer
self.loss = nn.BCELoss()
self.v_loss = Variational_loss
self.cuda = "False"
self.cuda_index = 0
# check if cuda is available
self.check_cuda(args.cuda)
# Using lower learning rate than suggested by (ADAM authors) lr=0.0002 and Beta_1 = 0.5 instead od 0.9 works better [Radford2015]
self.d_optimizer = torch.optim.Adam(self.D.parameters(),
lr=0.0002,
betas=(0.5, 0.999))
self.g_optimizer = torch.optim.Adam(self.G.parameters(),
lr=0.0002,
betas=(0.5, 0.999))
self.e_optimizer = torch.optim.Adam(self.E.parameters(),
lr=0.0002,
betas=(0.5, 0.999))
self.epochs = args.epochs
self.batch_size = args.batch_size
# Set the logger
self.logger = Logger('./logs')
self.number_of_images = 10
# cuda support
def check_cuda(self, cuda_flag=False):
if cuda_flag:
self.cuda = True
self.D.cuda(self.cuda_index)
self.G.cuda(self.cuda_index)
self.E.cuda(self.cuda_index)
self.loss = nn.BCELoss().cuda(self.cuda_index)
L1_loss = nn.L1Loss().cuda(self.cuda_index)
def Variational_loss(input, target, mu, logvar):
alpha = 1
beta = 1
recon_loss = L1_loss(input, target)
batch_size = logvar.data.shape[0]
nz = logvar.data.shape[1]
KLD_loss = (-0.5 * torch.sum(1 + logvar - mu.pow(2) -
logvar.exp())) / (nz * batch_size)
return alpha * recon_loss, beta * KLD_loss
self.v_loss = Variational_loss
print("Cuda enabled flag: ")
print(self.cuda)
def train(self, train_loader):
self.t_begin = t.time()
generator_iter = 0
self.file = open("inception_score_graph.txt", "w")
for epoch in range(self.epochs):
self.epoch_start_time = t.time()
for i, (images, _) in enumerate(train_loader):
# Check if round number of batches
if i == train_loader.dataset.__len__() // self.batch_size:
break
z = torch.rand((self.batch_size, 100, 1, 1))
real_labels = torch.ones(self.batch_size)
fake_labels = torch.zeros(self.batch_size)
if self.cuda:
images, z = Variable(images).cuda(
self.cuda_index), Variable(z).cuda(self.cuda_index)
real_labels, fake_labels = Variable(real_labels).cuda(
self.cuda_index), Variable(fake_labels).cuda(
self.cuda_index)
else:
images, z = Variable(images), Variable(z)
real_labels, fake_labels = Variable(real_labels), Variable(
fake_labels)
self.e_optimizer.zero_grad()
self.g_optimizer.zero_grad()
mu, logvar = self.E(images)
std = torch.exp(0.5 * logvar)
eps = Variable(torch.randn(std.size()),
requires_grad=False).cuda(self.cuda_index)
z = eps.mul(std).add_(mu)
recon_image = self.G(z)
err_recon, err_KLD = self.v_loss(recon_image, images.detach(),
mu, logvar)
err = self.recon_weight * (err_recon + err_KLD)
err.backward()
self.e_optimizer.step()
self.g_optimizer.step()
if self.cuda:
images, z = Variable(images).cuda(
self.cuda_index), Variable(z).cuda(self.cuda_index)
real_labels, fake_labels = Variable(real_labels).cuda(
self.cuda_index), Variable(fake_labels).cuda(
self.cuda_index)
else:
images, z = Variable(images), Variable(z)
real_labels, fake_labels = Variable(real_labels), Variable(
fake_labels)
# Train discriminator
# Compute BCE_Loss using real images
outputs = self.D(images)
d_loss_real = self.loss(outputs, real_labels)
real_score = outputs
# Compute BCE Loss using fake images
if self.cuda:
z = Variable(torch.randn(self.batch_size, 100, 1,
1)).cuda(self.cuda_index)
else:
z = Variable(torch.randn(self.batch_size, 100, 1, 1))
fake_images = self.G(z)
outputs = self.D(fake_images)
d_loss_fake = self.loss(outputs, fake_labels)
fake_score = outputs
# Optimize discriminator
d_loss = d_loss_real + d_loss_fake
self.D.zero_grad()
d_loss.backward()
self.d_optimizer.step()
# Train generator
# Compute loss with fake images
if self.cuda:
z = Variable(torch.randn(self.batch_size, 100, 1,
1)).cuda(self.cuda_index)
else:
z = Variable(torch.randn(self.batch_size, 100, 1, 1))
fake_images = self.G(z)
outputs = self.D(fake_images)
g_loss = self.loss(outputs, real_labels)
# Optimize generator
self.D.zero_grad()
self.G.zero_grad()
g_loss.backward()
self.g_optimizer.step()
generator_iter += 1
if generator_iter % 1000 == 0:
# Workaround because graphic card memory can't store more than 800+ examples in memory for generating image
# Therefore doing loop and generating 800 examples and stacking into list of samples to get 8000 generated images
# This way Inception score is more correct since there are different generated examples from every class of Inception model
sample_list = []
for i in range(10):
z = Variable(torch.randn(800, 100, 1,
1)).cuda(self.cuda_index)
samples = self.G(z)
sample_list.append(samples.data.cpu().numpy())
# Flattening list of lists into one list of numpy arrays
new_sample_list = list(chain.from_iterable(sample_list))
print(
"Calculating Inception Score over 8k generated images")
# Feeding list of numpy arrays
inception_score = get_inception_score(new_sample_list,
cuda=True,
batch_size=16,
resize=True,
splits=10)
print('Epoch-{}'.format(epoch + 1))
self.save_model()
if not os.path.exists('training_result_images/'):
os.makedirs('training_result_images/')
# Denormalize images and save them in grid 8x8
z = Variable(torch.randn(800, 100, 1,
1)).cuda(self.cuda_index)
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()[:64]
grid = utils.make_grid(samples)
utils.save_image(
grid,
'training_result_images/img_generatori_iter_{}.png'.
format(str(generator_iter).zfill(3)))
time = t.time() - self.t_begin
print("Inception score: {}".format(inception_score))
print("Generator iter: {}".format(generator_iter))
print("Time {}".format(time))
# Write to file inception_score, gen_iters, time
output = str(generator_iter) + " " + str(time) + " " + str(
inception_score[0]) + "\n"
self.file.write(output)
if ((i + 1) % 100) == 0:
print("Epoch: [%2d] [%4d/%4d] D_loss: %.8f, G_loss: %.8f" %
((epoch + 1),
(i + 1), train_loader.dataset.__len__() //
self.batch_size, d_loss.data[0], g_loss.data[0]))
z = Variable(
torch.randn(self.batch_size, 100, 1,
1).cuda(self.cuda_index))
# TensorBoard logging
# Log the scalar values
info = {'d_loss': d_loss.data[0], 'g_loss': g_loss.data[0]}
for tag, value in info.items():
self.logger.scalar_summary(tag, value, generator_iter)
# Log values and gradients of the parameters
for tag, value in self.D.named_parameters():
tag = tag.replace('.', '/')
self.logger.histo_summary(tag, self.to_np(value),
generator_iter)
self.logger.histo_summary(tag + '/grad',
self.to_np(value.grad),
generator_iter)
# Log the images while training
info = {
'real_images':
self.real_images(images, self.number_of_images),
'generated_images':
self.generate_img(z, self.number_of_images)
}
for tag, images in info.items():
self.logger.image_summary(tag, images, generator_iter)
self.t_end = t.time()
print('Time of training-{}'.format((self.t_end - self.t_begin)))
#self.file.close()
# Save the trained parameters
self.save_model()
def evaluate(self, test_loader, D_model_path, G_model_path):
self.load_model(D_model_path, G_model_path)
z = Variable(torch.randn(self.batch_size, 100, 1,
1)).cuda(self.cuda_index)
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()
grid = utils.make_grid(samples)
print("Grid of 8x8 images saved to 'dgan_model_image.png'.")
utils.save_image(grid, 'dgan_model_image.png')
def real_images(self, images, number_of_images):
if (self.C == 3):
return self.to_np(
images.view(-1, self.C, 32, 32)[:self.number_of_images])
else:
return self.to_np(images.view(-1, 32, 32)[:self.number_of_images])
def generate_img(self, z, number_of_images):
samples = self.G(z).data.cpu().numpy()[:number_of_images]
generated_images = []
for sample in samples:
if self.C == 3:
generated_images.append(sample.reshape(self.C, 32, 32))
else:
generated_images.append(sample.reshape(32, 32))
return generated_images
def to_np(self, x):
return x.data.cpu().numpy()
def save_model(self):
torch.save(self.G.state_dict(), './generator.pkl')
torch.save(self.D.state_dict(), './discriminator.pkl')
print('Models save to ./generator.pkl & ./discriminator.pkl ')
def load_model(self, D_model_filename, G_model_filename):
D_model_path = os.path.join(os.getcwd(), D_model_filename)
G_model_path = os.path.join(os.getcwd(), G_model_filename)
self.D.load_state_dict(torch.load(D_model_path))
self.G.load_state_dict(torch.load(G_model_path))
print('Generator model loaded from {}.'.format(G_model_path))
print('Discriminator model loaded from {}-'.format(D_model_path))
def generate_latent_walk(self, number):
if not os.path.exists('interpolated_images/'):
os.makedirs('interpolated_images/')
# Interpolate between twe noise(z1, z2) with number_int steps between
number_int = 10
z_intp = torch.FloatTensor(1, 100, 1, 1)
z1 = torch.randn(1, 100, 1, 1)
z2 = torch.randn(1, 100, 1, 1)
if self.cuda:
z_intp = z_intp.cuda()
z1 = z1.cuda()
z2 = z2.cuda()
z_intp = Variable(z_intp)
images = []
alpha = 1.0 / float(number_int + 1)
print(alpha)
for i in range(1, number_int + 1):
z_intp.data = z1 * alpha + z2 * (1.0 - alpha)
alpha += alpha
fake_im = self.G(z_intp)
fake_im = fake_im.mul(0.5).add(0.5) #denormalize
images.append(fake_im.view(self.C, 32, 32).data.cpu())
grid = utils.make_grid(images, nrow=number_int)
utils.save_image(
grid, 'interpolated_images/interpolated_{}.png'.format(
str(number).zfill(3)))
print(
"Saved interpolated images to interpolated_images/interpolated_{}."
.format(str(number).zfill(3)))
class GAN(object):
def __init__(self, args):
# Generator architecture
self.G = nn.Sequential(nn.Linear(100, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 512), nn.LeakyReLU(0.2),
nn.Linear(512, 1024), nn.LeakyReLU(0.2),
nn.Tanh())
# Discriminator architecture
self.D = nn.Sequential(nn.Linear(1024, 512), nn.LeakyReLU(0.2),
nn.Linear(512, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 1), nn.Sigmoid())
# Check if cuda is available
self.check_cuda(args.cuda)
# Binary cross entropy loss and optimizer
self.loss = nn.BCELoss()
self.d_optimizer = torch.optim.Adam(self.D.parameters(),
lr=0.0002,
weight_decay=0.00001)
self.g_optimizer = torch.optim.Adam(self.G.parameters(),
lr=0.0002,
weight_decay=0.00001)
# Set the logger
self.logger = Logger('./logs')
self.number_of_images = 10
self.epochs = args.epochs
self.batch_size = args.batch_size
# Cuda support
def check_cuda(self, cuda_flag=False):
if cuda_flag:
self.cuda_index = 0
self.cuda = True
self.D.cuda(self.cuda_index)
self.G.cuda(self.cuda_index)
self.loss = nn.BCELoss().cuda(self.cuda_index)
print("Cuda enabled flag: ")
print(self.cuda)
def train(self, train_loader):
self.t_begin = time.time()
generator_iter = 0
for epoch in range(self.epochs + 1):
for i, (images, _) in enumerate(train_loader):
# Check if round number of batches
if i == train_loader.dataset.__len__() // self.batch_size:
break
# Flatten image 1,32x32 to 1024
images = images.view(self.batch_size, -1)
z = torch.rand((self.batch_size, 100))
real_labels = Variable(torch.ones(self.batch_size)).cuda(
self.cuda_index)
fake_labels = Variable(torch.zeros(self.batch_size)).cuda(
self.cuda_index)
if self.cuda:
images, z = Variable(images.cuda(
self.cuda_index)), Variable(z.cuda(self.cuda_index))
else:
images, z = Variable(images), Variable(z)
# Train discriminator
# compute BCE_Loss using real images where BCE_Loss(x, y): - y * log(D(x)) - (1-y) * log(1 - D(x))
# [Training discriminator = Maximizing discriminator being correct]
outputs = self.D(images)
d_loss_real = self.loss(outputs, real_labels)
real_score = outputs
# Compute BCELoss using fake images
fake_images = self.G(z)
outputs = self.D(fake_images)
d_loss_fake = self.loss(outputs, fake_labels)
fake_score = outputs
# Optimizie discriminator
d_loss = d_loss_real + d_loss_fake
self.D.zero_grad()
d_loss.backward()
self.d_optimizer.step()
# Train generator
if self.cuda:
z = Variable(
torch.randn(self.batch_size,
100).cuda(self.cuda_index))
else:
z = Variable(torch.randn(self.batch_size, 100))
fake_images = self.G(z)
outputs = self.D(fake_images)
# We train G to maximize log(D(G(z))[maximize likelihood of discriminator being wrong] instead of
# minimizing log(1-D(G(z)))[minizing likelihood of discriminator being correct]
# From paper [https://arxiv.org/pdf/1406.2661.pdf]
g_loss = self.loss(outputs, real_labels)
# Optimize generator
self.D.zero_grad()
self.G.zero_grad()
g_loss.backward()
self.g_optimizer.step()
generator_iter += 1
if ((i + 1) % 100) == 0:
print("Epoch: [%2d] [%4d/%4d] D_loss: %.8f, G_loss: %.8f" %
((epoch + 1),
(i + 1), train_loader.dataset.__len__() //
self.batch_size, d_loss.data[0], g_loss.data[0]))
z = Variable(
torch.randn(self.batch_size,
100).cuda(self.cuda_index))
# ============ TensorBoard logging ============#
# (1) Log the scalar values
info = {'d_loss': d_loss.data[0], 'g_loss': g_loss.data[0]}
for tag, value in info.items():
self.logger.scalar_summary(tag, value, i + 1)
# (2) Log values and gradients of the parameters (histogram)
for tag, value in self.D.named_parameters():
tag = tag.replace('.', '/')
self.logger.histo_summary(tag, self.to_np(value),
i + 1)
self.logger.histo_summary(tag + '/grad',
self.to_np(value.grad),
i + 1)
# (3) Log the images
info = {
'real_images':
self.to_np(
images.view(-1, 32, 32)[:self.number_of_images]),
'generated_images':
self.generate_img(z, self.number_of_images)
}
for tag, images in info.items():
self.logger.image_summary(tag, images, i + 1)
if generator_iter % 1000 == 0:
print('Generator iter-{}'.format(generator_iter))
self.save_model()
if not os.path.exists('training_result_images/'):
os.makedirs('training_result_images/')
# Denormalize images and save them in grid 8x8
z = Variable(torch.randn(self.batch_size,
100)).cuda(self.cuda_index)
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()
grid = utils.make_grid(samples)
utils.save_image(
grid,
'training_result_images/gan_image_iter_{}.png'.format(
str(generator_iter).zfill(3)))
self.t_end = time.time()
print('Time of training-{}'.format((self.t_end - self.t_begin)))
# Save the trained parameters
self.save_model()
def evaluate(self, test_loader, D_model_path, G_model_path):
self.load_model(D_model_path, G_model_path)
z = Variable(torch.randn(self.batch_size, 100)).cuda(self.cuda_index)
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()
grid = utils.make_grid(samples)
print("Grid of 8x8 images saved to 'gan_model_image.png'.")
utils.save_image(grid, 'gan_model_image.png')
def generate_img(self, z, number_of_images):
samples = self.G(z).data.cpu().numpy()[:number_of_images]
generated_images = []
for sample in samples:
generated_images.append(sample.reshape(32, 32))
return generated_images
def to_np(self, x):
return x.data.cpu().numpy()
def save_model(self):
torch.save(self.G.state_dict(), './generator.pkl')
torch.save(self.D.state_dict(), './discriminator.pkl')
print('Models save to ./generator.pkl & ./discriminator.pkl ')
def load_model(self, D_model_filename, G_model_filename):
D_model_path = os.path.join(os.getcwd(), D_model_filename)
G_model_path = os.path.join(os.getcwd(), G_model_filename)
self.D.load_state_dict(torch.load(D_model_path))
self.G.load_state_dict(torch.load(G_model_path))
print('Generator model loaded from {}.'.format(G_model_path))
print('Discriminator model loaded from {}-'.format(D_model_path))
class WGAN_GP(object):
def __init__(self, args):
print("WGAN_GradientPenalty init model.")
self.G = Generator(args.channels)
self.D = Discriminator(args.channels)
self.C = args.channels
# Check if cuda is available
self.check_cuda(args.cuda)
# WGAN values from paper
self.learning_rate = 1e-4
self.b1 = 0.5
self.b2 = 0.999
self.batch_size = args.batch_size
# WGAN_gradient penalty uses ADAM
self.d_optimizer = optim.Adam(self.D.parameters(), lr=self.learning_rate, betas=(self.b1, self.b2))
self.g_optimizer = optim.Adam(self.G.parameters(), lr=self.learning_rate, betas=(self.b1, self.b2))
# Set the logger
self.logger = Logger('./logs')
self.logger.writer.flush()
self.number_of_images = 10
self.generator_iters = args.generator_iters
self.critic_iter = 5
self.lambda_term = 10
def check_cuda(self, cuda_flag=False):
if cuda_flag:
self.cuda_index = 0
self.cuda = True
self.D.cuda(self.cuda_index)
self.G.cuda(self.cuda_index)
print("Cuda enabled flag: {}".format(self.cuda))
def train(self, train_loader):
self.t_begin = t.time()
self.file = open(self.args.resultdir + "inception_score_graph.txt", "w")
# Now batches are callable self.data.next()
self.data = self.get_infinite_batches(train_loader)
one = torch.tensor(1.0)
mone = torch.tensor(-1.0)
if self.cuda:
one = one.cuda(self.cuda_index)
mone = mone.cuda(self.cuda_index)
for g_iter in range(self.generator_iters):
print("[%d/%d]"%(g_iter, self.generator_iters))
# Requires grad, Generator requires_grad = False
for p in self.D.parameters():
p.requires_grad = True
d_loss_real = 0
d_loss_fake = 0
Wasserstein_D = 0
# Train Dicriminator forward-loss-backward-update self.critic_iter times while 1 Generator forward-loss-backward-update
for d_iter in range(self.critic_iter):
self.D.zero_grad()
images = self.data.__next__()
# Check for batch to have full batch_size
if (images.size()[0] != self.batch_size):
continue
z = torch.rand((self.batch_size, 100, 1, 1))
if self.cuda:
images, z = Variable(images.cuda(self.cuda_index)), Variable(z.cuda(self.cuda_index))
else:
images, z = Variable(images), Variable(z)
# Train discriminator
# WGAN - Training discriminator more iterations than generator
# Train with real images
d_loss_real = self.D(images)
d_loss_real = d_loss_real.mean()
d_loss_real.backward(mone)
# Train with fake images
if self.cuda:
z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index)
else:
z = Variable(torch.randn(self.batch_size, 100, 1, 1))
fake_images = self.G(z)
d_loss_fake = self.D(fake_images)
d_loss_fake = d_loss_fake.mean()
d_loss_fake.backward(one)
# Train with gradient penalty
gradient_penalty = self.calculate_gradient_penalty(images.data, fake_images.data)
gradient_penalty.backward(retain_graph=True)
d_loss = d_loss_fake - d_loss_real + gradient_penalty
Wasserstein_D = d_loss_real - d_loss_fake
self.d_optimizer.step()
# Generator update
for p in self.D.parameters():
p.requires_grad = False # to avoid computation
self.G.zero_grad()
# train generator
# compute loss with fake images
z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index)
fake_images = self.G(z)
g_loss = self.D(fake_images)
g_loss = g_loss.mean()
g_loss.backward(mone)
g_cost = -g_loss
self.g_optimizer.step()
with torch.no_grad():
# Saving model and sampling images every 1000th generator iterations
if (g_iter) % 1000 == 0:
self.save_model()
# # Workaround because graphic card memory can't store more than 830 examples in memory for generating image
# # Therefore doing loop and generating 800 examples and stacking into list of samples to get 8000 generated images
# # This way Inception score is more correct since there are different generated examples from every class of Inception model
sample_list = []
for i in range(10):
#samples = self.data.__next__()
z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index)
samples = self.G(z)
sample_list.append(samples.data.cpu().numpy())
#
# # Flattening list of list into one list
new_sample_list = list(chain.from_iterable(sample_list))
print("Calculating Inception Score over 8k generated images")
# # Feeding list of numpy arrays
inception_score = get_inception_score(new_sample_list, cuda=True, batch_size=32,
resize=True, splits=10)
if not os.path.exists(self.args.resultdir + 'training_result_images/'):
os.makedirs(self.args.resultdir + 'training_result_images/')
# Denormalize images and save them in grid 8x8
z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index)
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()[:64]
grid = utils.make_grid(samples)
utils.save_image(grid, self.args.resultdir + 'training_result_images/img_generatori_iter_{}.png'.format(str(g_iter).zfill(3)))
# Testing
time = t.time() - self.t_begin
#print("Real Inception score: {}".format(inception_score))
print("Generator iter: {}".format(g_iter))
print("Time {}".format(time))
# Write to file inception_score, gen_iters, time
output = str(g_iter) + " " + str(time) + " " + str(inception_score[0]) + "\n"
self.file.write(output)
# ============ TensorBoard logging ============#
# (1) Log the scalar valuesself.args.resultdir +
info = {
'Wasserstein distance': Wasserstein_D.item(),
'Loss D': d_loss.item(),
'Loss G': g_cost.item(),
'Loss D Real': d_loss_real.item(),
'Loss D Fake': d_loss_fake.item()
}
for tag, value in info.items():
self.logger.scalar_summary(tag, value, g_iter + 1)
# (3) Log the images
info = {
'real_images': self.real_images(images, self.number_of_images),
'generated_images': self.generate_img(z, self.number_of_images)
}
for tag, images in info.items():
self.logger.image_summary(tag, images, g_iter + 1)
self.t_end = t.time()
print('Time of training-{}'.format((self.t_end - self.t_begin)))
#self.file.close()
# Save the trained parameters
self.save_model()
def evaluate(self, test_loader, D_model_path, G_model_path):
self.load_model(D_model_path, G_model_path)
z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index)
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()
grid = utils.make_grid(samples)
print("Grid of 8x8 images saved to 'dgan_model_image.png'.")
utils.save_image(grid, self.args.resultdir + 'dgan_model_image.png')
def calculate_gradient_penalty(self, real_images, fake_images):
eta = torch.FloatTensor(self.batch_size,1,1,1).uniform_(0,1)
eta = eta.expand(self.batch_size, real_images.size(1), real_images.size(2), real_images.size(3))
if self.cuda:
eta = eta.cuda(self.cuda_index)
else:
eta = eta
interpolated = eta * real_images + ((1 - eta) * fake_images)
if self.cuda:
interpolated = interpolated.cuda(self.cuda_index)
else:
interpolated = interpolated
# define it to calculate gradient
interpolated = Variable(interpolated, requires_grad=True)
# calculate probability of interpolated examples
prob_interpolated = self.D(interpolated)
# calculate gradients of probabilities with respect to examples
gradients = autograd.grad(outputs=prob_interpolated, inputs=interpolated,
grad_outputs=torch.ones(
prob_interpolated.size()).cuda(self.cuda_index) if self.cuda else torch.ones(
prob_interpolated.size()),
create_graph=True, retain_graph=True)[0]
grad_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * self.lambda_term
return grad_penalty
def real_images(self, images, number_of_images):
if (self.C == 3):
return self.to_np(images.view(-1, self.C, 32, 32)[:self.number_of_images])
else:
return self.to_np(images.view(-1, 32, 32)[:self.number_of_images])
def generate_img(self, z, number_of_images):
samples = self.G(z).data.cpu().numpy()[:number_of_images]
generated_images = []
for sample in samples:
if self.C == 3:
generated_images.append(sample.reshape(self.C, 32, 32))
else:
generated_images.append(sample.reshape(32, 32))
return generated_images
def to_np(self, x):
return x.data.cpu().numpy()
def save_model(self):
torch.save(self.G.state_dict(), './' + self.args.resultdir + 'generator.pkl')
torch.save(self.D.state_dict(), './' + self.args.resultdir + 'discriminator.pkl')
print('Models save to ./generator.pkl & ./discriminator.pkl ')
def load_model(self, D_model_filename, G_model_filename):
D_model_path = os.path.join(os.getcwd(), D_model_filename)
G_model_path = os.path.join(os.getcwd(), G_model_filename)
self.D.load_state_dict(torch.load(D_model_path))
self.G.load_state_dict(torch.load(G_model_path))
print('Generator model loaded from {}.'.format(G_model_path))
print('Discriminator model loaded from {}-'.format(D_model_path))
def get_infinite_batches(self, data_loader):
while True:
for i, (images, _) in enumerate(data_loader):
yield images
def generate_latent_walk(self, number):
if not os.path.exists(self.args.resultdir + 'interpolated_images/'):
os.makedirs(self.args.resultdir + 'interpolated_images/')
number_int = 10
# interpolate between twe noise(z1, z2).
z_intp = torch.FloatTensor(1, 100, 1, 1)
z1 = torch.randn(1, 100, 1, 1)
z2 = torch.randn(1, 100, 1, 1)
if self.cuda:
z_intp = z_intp.cuda()
z1 = z1.cuda()
z2 = z2.cuda()
z_intp = Variable(z_intp)
images = []
alpha = 1.0 / float(number_int + 1)
print(alpha)
for i in range(1, number_int + 1):
z_intp.data = z1*alpha + z2*(1.0 - alpha)
alpha += alpha
fake_im = self.G(z_intp)
fake_im = fake_im.mul(0.5).add(0.5) #denormalize
images.append(fake_im.view(self.C,32,32).data.cpu())
grid = utils.make_grid(images, nrow=number_int )
utils.save_image(grid, self.args.resultdir + 'interpolated_images/interpolated_{}.png'.format(str(number).zfill(3)))
print("Saved interpolated images.")