import argparse
import os
import random
import warnings
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
from ignite.contrib.handlers import ProgressBar
from ignite.engine import Engine, Events
from ignite.handlers import ModelCheckpoint, Timer
from ignite.metrics import RunningAverage
try:
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
except ImportError:
raise ImportError(
"Please install torchvision to run this example, for example "
"via conda by running 'conda install -c pytorch torchvision'. "
)
PRINT_FREQ = 100
FAKE_IMG_FNAME = "fake_sample_epoch_{:04d}.png"
REAL_IMG_FNAME = "real_sample_epoch_{:04d}.png"
LOGS_FNAME = "logs.tsv"
PLOT_FNAME = "plot.svg"
SAMPLES_FNAME = "samples.svg"
CKPT_PREFIX = "networks"
class Net(nn.Module):
""" A base class for both generator and the discriminator.
Provides a common weight initialization scheme.
"""
def weights_init(self):
for m in self.modules():
classname = m.__class__.__name__
if "Conv" in classname:
m.weight.data.normal_(0.0, 0.02)
elif "BatchNorm" in classname:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
def forward(self, x):
return x
class Generator(Net):
""" Generator network.
Args:
nf (int): Number of filters in the second-to-last deconv layer
"""
def __init__(self, z_dim, nf, nc):
super(Generator, self).__init__()
self.net = nn.Sequential(
# input is Z, going into a convolution
nn.ConvTranspose2d(in_channels=z_dim, out_channels=nf * 8, kernel_size=4, stride=1, padding=0, bias=False),
nn.BatchNorm2d(nf * 8),
nn.ReLU(inplace=True),
# state size. (nf*8) x 4 x 4
nn.ConvTranspose2d(in_channels=nf * 8, out_channels=nf * 4, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(nf * 4),
nn.ReLU(inplace=True),
# state size. (nf*4) x 8 x 8
nn.ConvTranspose2d(in_channels=nf * 4, out_channels=nf * 2, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(nf * 2),
nn.ReLU(inplace=True),
# state size. (nf*2) x 16 x 16
nn.ConvTranspose2d(in_channels=nf * 2, out_channels=nf, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(nf),
nn.ReLU(inplace=True),
# state size. (nf) x 32 x 32
nn.ConvTranspose2d(in_channels=nf, out_channels=nc, kernel_size=4, stride=2, padding=1, bias=False),
nn.Tanh()
# state size. (nc) x 64 x 64
)
self.weights_init()
def forward(self, x):
return self.net(x)
class Discriminator(Net):
""" Discriminator network.
Args:
nf (int): Number of filters in the first conv layer.
"""
def __init__(self, nc, nf):
super(Discriminator, self).__init__()
self.net = nn.Sequential(
# input is (nc) x 64 x 64
nn.Conv2d(in_channels=nc, out_channels=nf, kernel_size=4, stride=2, padding=1, bias=False),
nn.LeakyReLU(0.2, inplace=True),
# state size. (nf) x 32 x 32
nn.Conv2d(in_channels=nf, out_channels=nf * 2, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(nf * 2),
nn.LeakyReLU(0.2, inplace=True),
# state size. (nf*2) x 16 x 16
nn.Conv2d(in_channels=nf * 2, out_channels=nf * 4, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(nf * 4),
nn.LeakyReLU(0.2, inplace=True),
# state size. (nf*4) x 8 x 8
nn.Conv2d(in_channels=nf * 4, out_channels=nf * 8, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(nf * 8),
nn.LeakyReLU(0.2, inplace=True),
# state size. (nf*8) x 4 x 4
nn.Conv2d(in_channels=nf * 8, out_channels=1, kernel_size=4, stride=1, padding=0, bias=False),
nn.Sigmoid(),
)
self.weights_init()
def forward(self, x):
output = self.net(x)
return output.view(-1, 1).squeeze(1)
def check_manual_seed(seed):
""" If manual seed is not specified, choose a random one and communicate it to the user.
"""
seed = seed or random.randint(1, 10000)
random.seed(seed)
torch.manual_seed(seed)
print("Using manual seed: {seed}".format(seed=seed))
def check_dataset(dataset, dataroot):
"""
Args:
dataset (str): Name of the dataset to use. See CLI help for details
dataroot (str): root directory where the dataset will be stored.
Returns:
dataset (data.Dataset): torchvision Dataset object
"""
resize = transforms.Resize(64)
crop = transforms.CenterCrop(64)
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
if dataset in {"imagenet", "folder", "lfw"}:
dataset = dset.ImageFolder(root=dataroot, transform=transforms.Compose([resize, crop, to_tensor, normalize]))
nc = 3
elif dataset == "lsun":
dataset = dset.LSUN(
root=dataroot, classes=["bedroom_train"], transform=transforms.Compose([resize, crop, to_tensor, normalize])
)
nc = 3
elif dataset == "cifar10":
dataset = dset.CIFAR10(
root=dataroot, download=True, transform=transforms.Compose([resize, to_tensor, normalize])
)
nc = 3
elif dataset == "mnist":
dataset = dset.MNIST(root=dataroot, download=True, transform=transforms.Compose([resize, to_tensor, normalize]))
nc = 1
elif dataset == "fake":
dataset = dset.FakeData(size=256, image_size=(3, 64, 64), transform=to_tensor)
nc = 3
else:
raise RuntimeError("Invalid dataset name: {}".format(dataset))
return dataset, nc
def main(
dataset,
dataroot,
z_dim,
g_filters,
d_filters,
batch_size,
epochs,
learning_rate,
beta_1,
saved_G,
saved_D,
seed,
n_workers,
device,
alpha,
output_dir,
):
# seed
check_manual_seed(seed)
# data
dataset, num_channels = check_dataset(dataset, dataroot)
loader = data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=n_workers, drop_last=True)
# netowrks
netG = Generator(z_dim, g_filters, num_channels).to(device)
netD = Discriminator(num_channels, d_filters).to(device)
# criterion
bce = nn.BCELoss()
# optimizers
optimizerG = optim.Adam(netG.parameters(), lr=learning_rate, betas=(beta_1, 0.999))
optimizerD = optim.Adam(netD.parameters(), lr=learning_rate, betas=(beta_1, 0.999))
# load pre-trained models
if saved_G:
netG.load_state_dict(torch.load(saved_G))
if saved_D:
netD.load_state_dict(torch.load(saved_D))
# misc
real_labels = torch.ones(batch_size, device=device)
fake_labels = torch.zeros(batch_size, device=device)
fixed_noise = torch.randn(batch_size, z_dim, 1, 1, device=device)
def get_noise():
return torch.randn(batch_size, z_dim, 1, 1, device=device)
# The main function, processing a batch of examples
def step(engine, batch):
# unpack the batch. It comes from a dataset, so we have <images, labels> pairs. Discard labels.
real, _ = batch
real = real.to(device)
# -----------------------------------------------------------
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
netD.zero_grad()
# train with real
output = netD(real)
errD_real = bce(output, real_labels)
D_x = output.mean().item()
errD_real.backward()
# get fake image from generator
noise = get_noise()
fake = netG(noise)
# train with fake
output = netD(fake.detach())
errD_fake = bce(output, fake_labels)
D_G_z1 = output.mean().item()
errD_fake.backward()
# gradient update
errD = errD_real + errD_fake
optimizerD.step()
# -----------------------------------------------------------
# (2) Update G network: maximize log(D(G(z)))
netG.zero_grad()
# Update generator. We want to make a step that will make it more likely that discriminator outputs "real"
output = netD(fake)
errG = bce(output, real_labels)
D_G_z2 = output.mean().item()
errG.backward()
# gradient update
optimizerG.step()
return {"errD": errD.item(), "errG": errG.item(), "D_x": D_x, "D_G_z1": D_G_z1, "D_G_z2": D_G_z2}
# ignite objects
trainer = Engine(step)
checkpoint_handler = ModelCheckpoint(output_dir, CKPT_PREFIX, n_saved=10, require_empty=False)
timer = Timer(average=True)
# attach running average metrics
monitoring_metrics = ["errD", "errG", "D_x", "D_G_z1", "D_G_z2"]
RunningAverage(alpha=alpha, output_transform=lambda x: x["errD"]).attach(trainer, "errD")
RunningAverage(alpha=alpha, output_transform=lambda x: x["errG"]).attach(trainer, "errG")
RunningAverage(alpha=alpha, output_transform=lambda x: x["D_x"]).attach(trainer, "D_x")
RunningAverage(alpha=alpha, output_transform=lambda x: x["D_G_z1"]).attach(trainer, "D_G_z1")
RunningAverage(alpha=alpha, output_transform=lambda x: x["D_G_z2"]).attach(trainer, "D_G_z2")
# attach progress bar
pbar = ProgressBar()
pbar.attach(trainer, metric_names=monitoring_metrics)
@trainer.on(Events.ITERATION_COMPLETED(every=PRINT_FREQ))
def print_logs(engine):
fname = os.path.join(output_dir, LOGS_FNAME)
columns = ["iteration",] + list(engine.state.metrics.keys())
values = [str(engine.state.iteration),] + [str(round(value, 5)) for value in engine.state.metrics.values()]
with open(fname, "a") as f:
if f.tell() == 0:
print("\t".join(columns), file=f)
print("\t".join(values), file=f)
message = "[{epoch}/{max_epoch}][{i}/{max_i}]".format(
epoch=engine.state.epoch, max_epoch=epochs, i=(engine.state.iteration % len(loader)), max_i=len(loader)
)
for name, value in zip(columns, values):
message += " | {name}: {value}".format(name=name, value=value)
pbar.log_message(message)
# adding handlers using `trainer.on` decorator API
@trainer.on(Events.EPOCH_COMPLETED)
def save_fake_example(engine):
fake = netG(fixed_noise)
path = os.path.join(output_dir, FAKE_IMG_FNAME.format(engine.state.epoch))
vutils.save_image(fake.detach(), path, normalize=True)
# adding handlers using `trainer.on` decorator API
@trainer.on(Events.EPOCH_COMPLETED)
def save_real_example(engine):
img, y = engine.state.batch
path = os.path.join(output_dir, REAL_IMG_FNAME.format(engine.state.epoch))
vutils.save_image(img, path, normalize=True)
# adding handlers using `trainer.add_event_handler` method API
trainer.add_event_handler(
event_name=Events.EPOCH_COMPLETED, handler=checkpoint_handler, to_save={"netG": netG, "netD": netD}
)
# automatically adding handlers via a special `attach` method of `Timer` handler
timer.attach(
trainer,
start=Events.EPOCH_STARTED,
resume=Events.ITERATION_STARTED,
pause=Events.ITERATION_COMPLETED,
step=Events.ITERATION_COMPLETED,
)
# adding handlers using `trainer.on` decorator API
@trainer.on(Events.EPOCH_COMPLETED)
def print_times(engine):
pbar.log_message("Epoch {} done. Time per batch: {:.3f}[s]".format(engine.state.epoch, timer.value()))
timer.reset()
# adding handlers using `trainer.on` decorator API
@trainer.on(Events.EPOCH_COMPLETED)
def create_plots(engine):
try:
import matplotlib as mpl
mpl.use("agg")
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
except ImportError:
warnings.warn("Loss plots will not be generated -- pandas or matplotlib not found")
else:
df = pd.read_csv(os.path.join(output_dir, LOGS_FNAME), delimiter="\t", index_col="iteration")
_ = df.plot(subplots=True, figsize=(20, 20))
_ = plt.xlabel("Iteration number")
fig = plt.gcf()
path = os.path.join(output_dir, PLOT_FNAME)
fig.savefig(path)
# adding handlers using `trainer.on` decorator API
@trainer.on(Events.EXCEPTION_RAISED)
def handle_exception(engine, e):
if isinstance(e, KeyboardInterrupt) and (engine.state.iteration > 1):
engine.terminate()
warnings.warn("KeyboardInterrupt caught. Exiting gracefully.")
create_plots(engine)
checkpoint_handler(engine, {"netG_exception": netG, "netD_exception": netD})
else:
raise e
# Setup is done. Now let's run the training
trainer.run(loader, epochs)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--dataset",
required=True,
choices={"cifar10", "lsun", "imagenet", "folder", "lfw", "fake", "mnist"},
help="Type of the dataset to be used.",
)
parser.add_argument("--dataroot", required=True, help="path to dataset")
parser.add_argument("--workers", type=int, default=2, help="number of data loading workers")
parser.add_argument("--batch-size", type=int, default=64, help="input batch size")
parser.add_argument("--z-dim", type=int, default=100, help="size of the latent z vector")
parser.add_argument(
"--g-filters", type=int, default=64, help="Number of filters in the second-to-last generator deconv layer"
)
parser.add_argument("--d-filters", type=int, default=64, help="Number of filters in first discriminator conv layer")
parser.add_argument("--epochs", type=int, default=25, help="number of epochs to train for")
parser.add_argument("--lr", type=float, default=0.0002, help="learning rate")
parser.add_argument("--beta-1", type=float, default=0.5, help="beta_1 for adam")
parser.add_argument("--no-cuda", action="store_true", help="disables cuda")
parser.add_argument("--saved-G", default="", help="path to pickled generator (to continue training)")
parser.add_argument("--saved-D", default="", help="path to pickled discriminator (to continue training)")
parser.add_argument("--output-dir", default=".", help="directory to output images and model checkpoints")
parser.add_argument("--seed", type=int, help="manual seed")
parser.add_argument("--alpha", type=float, default=0.98, help="smoothing constant for exponential moving averages")
args = parser.parse_args()
dev = "cpu" if (not torch.cuda.is_available() or args.no_cuda) else "cuda:0"
try:
os.makedirs(args.output_dir)
except FileExistsError:
if (not os.path.isdir(args.output_dir)) or (len(os.listdir(args.output_dir)) > 0):
raise FileExistsError("Please provide a path to a non-existing or empty directory.")
main(
dataset=args.dataset,
dataroot=args.dataroot,
z_dim=args.z_dim,
g_filters=args.g_filters,
d_filters=args.d_filters,
batch_size=args.batch_size,
epochs=args.epochs,
learning_rate=args.lr,
beta_1=args.beta_1,
saved_D=args.saved_D,
saved_G=args.saved_G,
seed=args.seed,
device=dev,
n_workers=args.workers,
alpha=args.alpha,
output_dir=args.output_dir,
)
