<jupyter_start><jupyter_text>Setup  Imports<jupyter_code>from vaiutils import path_consts, channels_last, plot_images, smoothen
from tqdm import tqdm_notebook
from os.path import exists
import torch
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torchvision.datasets import MNIST
from torch.utils.data.dataloader import DataLoader<jupyter_output><empty_output><jupyter_text>Load variables and data<jupyter_code>for k, v in path_consts('WGAN', 'MNIST'):
    exec(k + '=v')
batch_size = 64
transform = transforms.Compose(
    [transforms.Scale(32, 1), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
data = DataLoader(MNIST(DIR_DATA, transform=transform),
                  batch_size=batch_size, shuffle=True, drop_last=True)
plot_images(channels_last(iter(data).next()[0].numpy()))<jupyter_output><empty_output><jupyter_text>Create Model<jupyter_code>z_dim = 100
n_critic = 5
critic_clip = 0.01
def init_weights(model):
    for m in model.modules():
        if type(m) == nn.Linear or type(m) == nn.Conv2d or type(m) == nn.ConvTranspose2d:
            nn.init.normal(m.weight.data, std=0.02)
        elif type(m) == nn.BatchNorm2d:
            nn.init.normal(m.weight.data, mean=1, std=0.02)
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.conv1 = nn.ConvTranspose2d(
            z_dim, 256, 4)
        self.bn1 = nn.BatchNorm2d(256)
        self.conv2 = nn.ConvTranspose2d(
            256, 128, 4, stride=2, padding=1)
        self.bn2 = nn.BatchNorm2d(128)
        self.conv3 = nn.ConvTranspose2d(
            128, 64, 4, stride=2, padding=1)
        self.bn3 = nn.BatchNorm2d(64)
        self.conv4 = nn.ConvTranspose2d(
            64, 1, 4, stride=2, padding=1)
        init_weights(self)

    def relu(self, x):
        return F.relu(x, inplace=True)

    def forward(self, x):
        x = x.view(-1, z_dim, 1, 1)
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = F.tanh(self.conv4(x))
        return x
class Critic(nn.Module):
    def __init__(self):
        super(Critic, self).__init__()
        self.conv1 = nn.Conv2d(1, 64, 4, stride=2, padding=1, bias=False)
        self.conv2 = nn.Conv2d(64, 128, 4, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(128)
        self.conv3 = nn.Conv2d(128, 256, 4, stride=2, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(256)
        self.conv4 = nn.Conv2d(256, 1, 4, bias=False)
        init_weights(self)

    def lrelu(self, x):
        return F.leaky_relu(x,0.2, inplace=True)

    def forward(self, x):
        x = self.lrelu(self.conv1(x))
        x = self.lrelu(self.conv2(x))
        x = self.lrelu(self.conv3(x))
        x = self.conv4(x).squeeze().mean()
        return x<jupyter_output><empty_output><jupyter_text>Train Model<jupyter_code>generator = Generator().cuda()
critic = Critic().cuda()

optimizer_g = optim.RMSprop(generator.parameters(), lr=5e-4)
optimizer_c = optim.RMSprop(critic.parameters(), lr=5e-4)
def save_model():
    %mkdir -p $DIR_OUTPUT
    torch.save(generator.state_dict(), DIR_OUTPUT + '/generator.dat')
    torch.save(critic.state_dict(), DIR_OUTPUT + '/critic.dat')


def load_model():
    if not exists(DIR_OUTPUT + '/generator.dat'):
        return
    generator.load_state_dict(torch.load(DIR_OUTPUT + '/generator.dat'))
    critic.load_state_dict(torch.load(DIR_OUTPUT + '/critic.dat'))
history = {'loss': []}


def optimize(epochs=1):
    load_model()
    for epoch in tqdm_notebook(range(epochs)):
        for batch, datum in enumerate(tqdm_notebook(data)):
            x_batch = Variable(datum[0].cuda())
            z_batch = Variable(torch.randn(batch_size, z_dim).cuda())

            optimizer_c.zero_grad()
            loss = critic(x_batch) - critic(generator(z_batch))
            (-loss).backward()
            optimizer_c.step()

            #for p in critic.parameters():
                #p.data.clamp_(-critic_clip, critic_clip)

            if (batch + 1) % n_critic == 0:
                z_batch = Variable(torch.randn(batch_size, z_dim).cuda())

                optimizer_g.zero_grad()
                loss_g = -critic(generator(z_batch))
                loss_g.backward()
                optimizer_g.step()

                if ((batch + 1) / n_critic) % int(100 / n_critic) == 0:
                    history['loss'].append(loss.cpu().data.numpy()[0])

            if batch % 500 == 0 or batch==len(data)-1:
                z_batch = Variable(torch.randn(batch_size, z_dim).cuda())
                x_gen = generator(z_batch)
                plot_images(channels_last(x_gen[:5].cpu().data.numpy()))
                save_model()
optimize(10)
plot(smoothen(history['loss'],0.1)), np.array(history['loss'][-10:]).mean()<jupyter_output><empty_output><jupyter_text>Generate Images<jupyter_code>z_batch = Variable(torch.randn(64, z_dim).cuda())
x_gen = generator(z_batch)
plot_images(channels_last(x_gen.cpu().data.numpy()))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14379717437584671, 'fraction_numerical': 0.03309463905554481, 'mean_word_length': 2.9723289777094544, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '270012', 'license': 'MIT', 'n_tokens_mistral': 1969, 'n_tokens_neox': 1855, 'n_words': 330, 'path': 'Research-Papers/GANs/WGAN/Notebooks/WGAN.ipynb', 'repo_name': 'svaisakh/aiprojects', 'size': '5167'}