<jupyter_start><jupyter_text>Wasserstein GAN (WGAN)Based on paper [Wasserstein GAN](https://arxiv.org/abs/1701.07875)  Outline- Introduction- Prerequest- Datasets- Build Models    - Generator Models    - Discriminator Models- Models Settings- Training- Result  Introduction Abstract :We introduce a new algorithm named WGAN, an alternative to traditional GAN training. In this new model, we show that we can improve the stability of learning, get rid of problems like mode collapse, and provide meaningful learning curves useful for debugging and hyperparameter searches. Furthermore, we show that the corresponding optimization problem is sound, and provide extensive theoretical work highlighting the deep connections to other distances between distributions. ![WGAN](https://miro.medium.com/max/1600/1*Yfa9bZL0d4NHaU1mHbGzjw.jpeg)source: https://medium.com/@jonathan_hui/gan-wasserstein-gan-wgan-gp-6a1a2aa1b490  Prerequest<jupyter_code>from google.colab import drive
drive.mount('/content/drive')
# import All prerequisites
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
from torchvision.utils import save_image
import numpy as np
import os

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ROOT = "/content/drive/My Drive/Colab Notebooks/DSC_UI_GAN/Batch1/W2/WGAN"
sample_dir = os.path.join(ROOT, 'sample')
# Make dir if no exist
if not os.path.exists(sample_dir):
    os.makedirs(sample_dir)<jupyter_output><empty_output><jupyter_text>Dataset<jupyter_code>batch_size = 100

# MNIST Dataset
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])])

train_dataset = datasets.MNIST(root='./mnist_data/', train=True, transform=transform, download=True)

# Data Loader (Input Pipeline)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)

examples = enumerate(train_loader)
batch_idx, (example_data, example_targets) = next(examples)
## Print example
import matplotlib.pyplot as plt

fig = plt.figure()
for i in range(6):
  plt.subplot(2,3,i+1)
  plt.tight_layout()
  plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
  plt.title("Ground Truth: {}".format(example_targets[i]))
  plt.xticks([])
  plt.yticks([])
fig<jupyter_output><empty_output><jupyter_text>Build Models  Generator Model<jupyter_code>class Generator(nn.Module):
    def __init__(self, g_input_dim, g_output_dim):
        super(Generator, self).__init__()

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Linear(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *block(g_input_dim, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Linear(1024, g_output_dim), 
            nn.Tanh())

    def forward(self, z):
        image = self.model(z)
        image = image.view(image.size(0), -1)
        return image<jupyter_output><empty_output><jupyter_text>Discriminator Model<jupyter_code>class Discriminator(nn.Module):
    def __init__(self, d_input_dim):
        super(Discriminator, self).__init__()

        self.model = nn.Sequential(
            nn.Linear(d_input_dim, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            nn.Sigmoid(),
        )

    def forward(self, image):
        img_flat = image.view(image.size(0), -1)
        validity = self.model(img_flat)
        return validity<jupyter_output><empty_output><jupyter_text>Build network<jupyter_code># build network
z_dim = 100
mnist_dim = train_dataset.train_data.size(1) * train_dataset.train_data.size(2)

G = Generator(g_input_dim = z_dim, g_output_dim = mnist_dim).to(device)
D = Discriminator(mnist_dim).to(device)
print(G, D)<jupyter_output><empty_output><jupyter_text>Train Process![WGAN Algorithm](https://github.com/DSC-UI-SRIN/Introduction-to-GAN/raw/master/2%20-%20%20Wasserstein%20GANs/images/wgan_algorithm.png)<jupyter_code># optimizer
lr = 0.00005
n_critic =  5
clip_value = 0.01
G_optimizer = torch.optim.RMSprop(params=[p for p in G.parameters() if p.requires_grad], lr=lr)
D_optimizer = torch.optim.RMSprop(params=[p for p in D.parameters() if p.requires_grad], lr=lr)
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
epochs = 200
batches_done = 0
list_loss_D = []
list_loss_G = []
for epoch in range(epochs):

    for i, (imgs, _) in enumerate(train_loader):

        # Configure input
        real_imgs = Variable(imgs.type(Tensor))

        # ---------------------
        #  Train Discriminator
        # ---------------------

        D_optimizer.zero_grad()

        # Sample noise as generator input
        z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], z_dim))))

        # Generate a batch of images
        fake_imgs = G(z).detach()
        # Adversarial loss
        d_loss = -torch.mean(D(real_imgs)) + torch.mean(D(fake_imgs))

        d_loss.backward()
        D_optimizer.step()

        # Clip weights of discriminator
        for p in D.parameters():
            p.data.clamp_(-clip_value, clip_value)

        # Train the generator every n_critic iterations
        if i % n_critic == 0:

            # -----------------
            #  Train Generator
            # -----------------

            G_optimizer.zero_grad()

            # Generate a batch of images
            gen_imgs = G(z)
            # Adversarial loss
            g_loss = -torch.mean(D(gen_imgs))


            g_loss.backward()
            G_optimizer.step()

            # add loss to list
            list_loss_D.append(d_loss.item())
            list_loss_G.append(g_loss.item())
        
        if i % 300 == 0:
            print(
              "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
              % (epoch, epochs, i, len(train_loader), d_loss.item(), g_loss.item()))

    if epoch % 5 == 0:
        save_image(gen_imgs.view(gen_imgs.size(0), 1, 28, 28), os.path.join(sample_dir, "%d.png" % epoch), nrow=5, normalize=True)

torch.save(G, os.path.join(ROOT, 'G.pt'))
torch.save(D, os.path.join(ROOT, 'D.pt'))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1300763358778626, 'fraction_numerical': 0.020763358778625954, 'mean_word_length': 3.125314861460957, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '414460', 'license': 'MIT', 'n_tokens_mistral': 2307, 'n_tokens_neox': 2151, 'n_words': 466, 'path': '2 - Wasserstein GANs/WGAN.ipynb', 'repo_name': 'drainvers/Introduction-to-GAN', 'size': '6550'}