chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,285 @@
|
||||
import os
|
||||
|
||||
import matplotlib.animation as animation
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.parallel
|
||||
import torch.utils.data
|
||||
import torchvision.datasets as dset
|
||||
import torchvision.transforms as transforms
|
||||
import torchvision.utils as vutils
|
||||
from scipy.stats import entropy
|
||||
from torch.autograd import Variable
|
||||
from torch.nn import functional as F
|
||||
|
||||
import ray
|
||||
|
||||
# Training parameters
|
||||
workers = 2
|
||||
batch_size = 64
|
||||
image_size = 32
|
||||
|
||||
# Number of channels in the training images. For color images this is 3
|
||||
nc = 1
|
||||
|
||||
# Size of z latent vector (i.e. size of generator input)
|
||||
nz = 100
|
||||
|
||||
# Size of feature maps in generator
|
||||
ngf = 32
|
||||
|
||||
# Size of feature maps in discriminator
|
||||
ndf = 32
|
||||
|
||||
# Beta1 hyperparam for Adam optimizers
|
||||
beta1 = 0.5
|
||||
|
||||
# iterations of actual training in each Trainable _train
|
||||
train_iterations_per_step = 5
|
||||
|
||||
MODEL_PATH = os.path.expanduser("~/.ray/models/mnist_cnn.pt")
|
||||
|
||||
|
||||
def get_data_loader(data_dir="~/data"):
|
||||
dataset = dset.MNIST(
|
||||
root=data_dir,
|
||||
download=True,
|
||||
transform=transforms.Compose(
|
||||
[
|
||||
transforms.Resize(image_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.13066,), (0.30131,)),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
# Create the dataloader
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
dataset, batch_size=batch_size, shuffle=True, num_workers=workers
|
||||
)
|
||||
|
||||
return dataloader
|
||||
|
||||
|
||||
# __GANmodel_begin__
|
||||
# custom weights initialization called on netG and netD
|
||||
def weights_init(m):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find("Conv") != -1:
|
||||
nn.init.normal_(m.weight.data, 0.0, 0.02)
|
||||
elif classname.find("BatchNorm") != -1:
|
||||
nn.init.normal_(m.weight.data, 1.0, 0.02)
|
||||
nn.init.constant_(m.bias.data, 0)
|
||||
|
||||
|
||||
# Generator Code
|
||||
class Generator(nn.Module):
|
||||
def __init__(self):
|
||||
super(Generator, self).__init__()
|
||||
self.main = nn.Sequential(
|
||||
# input is Z, going into a convolution
|
||||
nn.ConvTranspose2d(nz, ngf * 4, 4, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(ngf * 4),
|
||||
nn.ReLU(True),
|
||||
nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
|
||||
nn.BatchNorm2d(ngf * 2),
|
||||
nn.ReLU(True),
|
||||
nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
|
||||
nn.BatchNorm2d(ngf),
|
||||
nn.ReLU(True),
|
||||
nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
return self.main(input)
|
||||
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
def __init__(self):
|
||||
super(Discriminator, self).__init__()
|
||||
self.main = nn.Sequential(
|
||||
nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
|
||||
nn.BatchNorm2d(ndf * 2),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
|
||||
nn.BatchNorm2d(ndf * 4),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
return self.main(input)
|
||||
|
||||
|
||||
# __GANmodel_end__
|
||||
|
||||
|
||||
# __INCEPTION_SCORE_begin__
|
||||
class Net(nn.Module):
|
||||
"""
|
||||
LeNet for MNist classification, used for inception_score
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(Net, self).__init__()
|
||||
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
|
||||
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
|
||||
self.conv2_drop = nn.Dropout2d()
|
||||
self.fc1 = nn.Linear(320, 50)
|
||||
self.fc2 = nn.Linear(50, 10)
|
||||
|
||||
def forward(self, x):
|
||||
x = F.relu(F.max_pool2d(self.conv1(x), 2))
|
||||
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
|
||||
x = x.view(-1, 320)
|
||||
x = F.relu(self.fc1(x))
|
||||
x = F.dropout(x, training=self.training)
|
||||
x = self.fc2(x)
|
||||
return F.log_softmax(x, dim=1)
|
||||
|
||||
|
||||
def inception_score(imgs, mnist_model_ref, batch_size=32, splits=1):
|
||||
N = len(imgs)
|
||||
dtype = torch.FloatTensor
|
||||
dataloader = torch.utils.data.DataLoader(imgs, batch_size=batch_size)
|
||||
cm = ray.get(mnist_model_ref) # Get the mnist model from Ray object store.
|
||||
up = nn.Upsample(size=(28, 28), mode="bilinear").type(dtype)
|
||||
|
||||
def get_pred(x):
|
||||
x = up(x)
|
||||
x = cm(x)
|
||||
return F.softmax(x).data.cpu().numpy()
|
||||
|
||||
preds = np.zeros((N, 10))
|
||||
for i, batch in enumerate(dataloader, 0):
|
||||
batch = batch.type(dtype)
|
||||
batchv = Variable(batch)
|
||||
batch_size_i = batch.size()[0]
|
||||
preds[i * batch_size : i * batch_size + batch_size_i] = get_pred(batchv)
|
||||
|
||||
# Now compute the mean kl-div
|
||||
split_scores = []
|
||||
for k in range(splits):
|
||||
part = preds[k * (N // splits) : (k + 1) * (N // splits), :]
|
||||
py = np.mean(part, axis=0)
|
||||
scores = []
|
||||
for i in range(part.shape[0]):
|
||||
pyx = part[i, :]
|
||||
scores.append(entropy(pyx, py))
|
||||
split_scores.append(np.exp(np.mean(scores)))
|
||||
|
||||
return np.mean(split_scores), np.std(split_scores)
|
||||
|
||||
|
||||
# __INCEPTION_SCORE_end__
|
||||
|
||||
|
||||
def train_func(
|
||||
netD,
|
||||
netG,
|
||||
optimG,
|
||||
optimD,
|
||||
criterion,
|
||||
dataloader,
|
||||
iteration,
|
||||
device,
|
||||
mnist_model_ref,
|
||||
):
|
||||
real_label = 1
|
||||
fake_label = 0
|
||||
|
||||
for i, data in enumerate(dataloader, 0):
|
||||
if i >= train_iterations_per_step:
|
||||
break
|
||||
|
||||
netD.zero_grad()
|
||||
real_cpu = data[0].to(device)
|
||||
b_size = real_cpu.size(0)
|
||||
label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
|
||||
output = netD(real_cpu).view(-1)
|
||||
errD_real = criterion(output, label)
|
||||
errD_real.backward()
|
||||
D_x = output.mean().item()
|
||||
|
||||
noise = torch.randn(b_size, nz, 1, 1, device=device)
|
||||
fake = netG(noise)
|
||||
label.fill_(fake_label)
|
||||
output = netD(fake.detach()).view(-1)
|
||||
errD_fake = criterion(output, label)
|
||||
errD_fake.backward()
|
||||
D_G_z1 = output.mean().item()
|
||||
errD = errD_real + errD_fake
|
||||
optimD.step()
|
||||
|
||||
netG.zero_grad()
|
||||
label.fill_(real_label)
|
||||
output = netD(fake).view(-1)
|
||||
errG = criterion(output, label)
|
||||
errG.backward()
|
||||
D_G_z2 = output.mean().item()
|
||||
optimG.step()
|
||||
|
||||
is_score, is_std = inception_score(fake, mnist_model_ref)
|
||||
|
||||
# Output training stats
|
||||
if iteration % 10 == 0:
|
||||
print(
|
||||
"[%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z))"
|
||||
": %.4f / %.4f \tInception score: %.4f"
|
||||
% (
|
||||
iteration,
|
||||
len(dataloader),
|
||||
errD.item(),
|
||||
errG.item(),
|
||||
D_x,
|
||||
D_G_z1,
|
||||
D_G_z2,
|
||||
is_score,
|
||||
)
|
||||
)
|
||||
|
||||
return errG.item(), errD.item(), is_score
|
||||
|
||||
|
||||
def plot_images(dataloader):
|
||||
# Plot some training images
|
||||
real_batch = next(iter(dataloader))
|
||||
plt.figure(figsize=(8, 8))
|
||||
plt.axis("off")
|
||||
plt.title("Original Images")
|
||||
plt.imshow(
|
||||
np.transpose(
|
||||
vutils.make_grid(real_batch[0][:64], padding=2, normalize=True).cpu(),
|
||||
(1, 2, 0),
|
||||
)
|
||||
)
|
||||
|
||||
plt.show()
|
||||
|
||||
|
||||
def demo_gan(checkpoint_paths):
|
||||
img_list = []
|
||||
fixed_noise = torch.randn(64, nz, 1, 1)
|
||||
for path in checkpoint_paths:
|
||||
checkpoint_dict = torch.load(os.path.join(path, "checkpoint.pt"))
|
||||
|
||||
loadedG = Generator()
|
||||
loadedG.load_state_dict(checkpoint_dict["netGmodel"])
|
||||
with torch.no_grad():
|
||||
fake = loadedG(fixed_noise).detach().cpu()
|
||||
img_list.append(vutils.make_grid(fake, padding=2, normalize=True))
|
||||
|
||||
fig = plt.figure(figsize=(8, 8))
|
||||
plt.axis("off")
|
||||
ims = [[plt.imshow(np.transpose(i, (1, 2, 0)), animated=True)] for i in img_list]
|
||||
ani = animation.ArtistAnimation(
|
||||
fig, ims, interval=1000, repeat_delay=1000, blit=True
|
||||
)
|
||||
ani.save("./generated.gif", writer="imagemagick", dpi=72)
|
||||
plt.show()
|
||||
Reference in New Issue
Block a user