<jupyter_start><jupyter_text><jupyter_code># when running in colab notebooks, first install library
!pip install import-ai
# upload respective dataset manually from examples directory of the library or download as below
!apt install subversion
!svn checkout https://github.com/srirambandi/ai/trunk/examples/MNIST
from google.colab import drive
drive.mount('/content/drive')
import ai
import numpy as np
z_dim = 100
gf_dim = 64
df_dim = 64
ai.manual_seed(2357)
def data_generator(m):
    train_dict = np.load('MNIST/train.npy', allow_pickle=True)
    test_dict = np.load('MNIST/test.npy', allow_pickle=True)
    data = np.concatenate([train_dict.item()['data'], test_dict.item()['data']])
    data = data.transpose(1, 2, 0)   # making data batch-last
    data = data.reshape(1, *data.shape) / 127.5 - 1.   # adding channel dimension and normalizing data
    epoch = 0
    
    while True:
        epoch += 1
        for batch in range(int(data.shape[-1] / m)):
            yield data[...,batch * m:(batch + 1) * m], epoch
class Generator(ai.Module):
    def __init__(self):
        self.g_fc = ai.Linear(z_dim, 8*gf_dim * 2 * 2)
        self.g_bn1 = ai.BatchNorm((8*gf_dim, 2, 2))
        self.g_deconv1 = ai.ConvTranspose2d(8*gf_dim, 4*gf_dim, kernel_size=5, stride=2, padding=2, a=1)
        self.g_bn2 = ai.BatchNorm((4*gf_dim, 4, 4))
        self.g_deconv2 = ai.ConvTranspose2d(4*gf_dim, 2*gf_dim, kernel_size=5, stride=2, padding=2, a=0)
        self.g_bn3 = ai.BatchNorm((2*gf_dim, 7, 7))
        self.g_deconv3 = ai.ConvTranspose2d(2*gf_dim, gf_dim, kernel_size=5, stride=2, padding=2, a=1)
        self.g_bn4 = ai.BatchNorm((gf_dim, 14, 14))
        self.g_deconv4 = ai.ConvTranspose2d(gf_dim, 1, kernel_size=5, stride=2, padding=2, a=1)
        
    def forward(self, z):
        o1 = ai.G.reshape(self.g_fc(z), (8*gf_dim, 2, 2))
        o2 = ai.G.relu(self.g_bn1(o1))
        o3 = ai.G.relu(self.g_bn2(self.g_deconv1(o2)))
        o4 = ai.G.relu(self.g_bn3(self.g_deconv2(o3)))
        o5 = ai.G.relu(self.g_bn4(self.g_deconv3(o4)))
        fake_image = ai.G.tanh(self.g_deconv4(o5))
        return fake_image
class Critic(ai.Module):
    def __init__(self):
        self.d_conv1 = ai.Conv2d(1, 64, kernel_size=5, stride=2, padding=2)
        self.d_conv2 = ai.Conv2d(64, 2*64, kernel_size=5, stride=2, padding=2)
        self.d_bn1 = ai.BatchNorm((2*64, 7, 7))
        self.d_conv3 = ai.Conv2d(2*64, 3*64, kernel_size=5, stride=2, padding=2)
        self.d_bn2 = ai.BatchNorm((3*64, 4, 4))
        self.d_conv4 = ai.Conv2d(3*64, 4*64, kernel_size=5, stride=2, padding=2)
        self.d_bn3 = ai.BatchNorm((4*64, 2, 2))
        self.d_fc = ai.Linear(1024, 1)
        
    def forward(self, image):
        o1 = ai.G.lrelu(self.d_conv1(image))
        o2 = ai.G.lrelu(self.d_bn1(self.d_conv2(o1)))
        o3 = ai.G.lrelu(self.d_bn2(self.d_conv3(o2)))
        o4 = ai.G.lrelu(self.d_bn3(self.d_conv4(o3)))
        o5 = self.d_fc(o4)
        return o5
generator = Generator()
critic = Critic()
print(generator)
print(critic)
lr = 0.00005
c = 0.01
g_optim = ai.Optimizer(generator.parameters(), optim_fn='RMSProp', lr=lr)
c_optim = ai.Optimizer(critic.parameters(), optim_fn='RMSProp', lr=lr)
it, epoch = 0, 0
m = 64   # batch size
n_critic = 5   # number of critic updates per generator update
c = 0.01  # clip value for critic weigthts after each update

batch_size = ai.Parameter((1, 1), init_ones=True, constant=float(m), eval_grad=False)   # batch size 'm' as a Parameter constant
neg_ones = ai.Parameter((1, 1), init_ones=True, constant=-1.0, eval_grad=False)

# real images data generator
data = data_generator(m)

sample_z = np.random.uniform(-1, 1, (z_dim, m))
sampled_images = None
def sampler(sampled_images):
    ai.G.grad_mode = False

    # generate images like real data
    fake_images = generator.forward(sample_z).data
    fake_images = (fake_images + 1.) / 2.

    if sampled_images is not None:
        sampled_images = np.concatenate([sampled_images, fake_images], axis=-1)
    else:
        sampled_images = fake_images
    
    ai.G.grad_mode = True

    return sampled_images
for it in range(10000):

    # freeze generator before optimizing critic
    for p in generator.parameters():
        p.eval_grad = False

    critic_iter = 100 if it < 25 or it % 500 == 0 else n_critic

    # training critic to identify real/fake data
    for _ in range(critic_iter):

        real_images, epoch = data.__next__()
        if (real_images.shape[-1] != m):
            continue

        c_loss_real = critic.forward(real_images)

        z = np.random.randn(z_dim, m)
        fake_images = generator.forward(z)

        c_loss_fake = critic.forward(fake_images)

        # add the above two losses and reduce their mean
        c_loss = c_loss_fake - c_loss_real
        c_loss = ai.G.sum(c_loss)
        c_loss = c_loss / batch_size
        c_loss.grad = np.ones(c_loss.shape)

        c_loss.backward()
        c_optim.step()
        c_optim.zero_grad()

        ai.clip_grad_value(critic.parameters(), c)

    # unfreeze generator
    for p in generator.parameters():
        p.eval_grad = True

    # training generator to fool critic with fake data
    z = np.random.uniform(-1, 1, (z_dim, m))
    fake_images = generator.forward(z)

    g_loss = critic.forward(fake_images)
    g_loss = ai.G.sum(g_loss)
    g_loss = g_loss / batch_size
    g_loss = neg_ones * g_loss
    g_loss.grad = np.ones(g_loss.shape)

    g_loss.backward()
    g_optim.step()
    g_optim.zero_grad()
    c_optim.zero_grad()

    if it%10 == 0:
        print('Iter: {}, Epoch: {}, c_loss: {}, g_loss: {}'.format(it, epoch, c_loss.data[0, 0], g_loss.data[0, 0]))
    
    if it%100 == 0:
        sampled_images=sampler(sampled_images)
        np.save('/content/drive/My Drive/WGAN/sampled_images.npy', sampled_images)
        generator.save('/content/drive/My Drive/WGAN/Generator.npy')
        critic.save('/content/drive/My Drive/WGAN/Critic.npy')
def save_images(sampled_images):
    
    import imageio

    grids = []
    h, w = int(np.sqrt(m)), int(np.sqrt(m))
    for i in range(int(sampled_images.shape[-1]/m)):
        grid = []
        batch = sampled_images[...,i*m:(i+1)*m]

        for r in range(h):
            row = np.squeeze(batch[...,r*w])

            for c in range(1, w):
                row = np.concatenate([row, np.squeeze(batch[...,r*w + c])], axis=1)

            grid.append(row)

        grid = np.concatenate(grid, axis=0)
        grids.append(grid)

    imageio.mimsave('/content/drive/My Drive/WGAN/wgan.gif', grids)

sampled_images = np.load('/content/drive/My Drive/WGAN/sampled_images.npy')
save_images(sampled_images)
from IPython.display import Image

with open('/content/drive/My Drive/WGAN/wgan.gif','rb') as f:
    display(Image(data=f.read(), format='png'))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14015817223198593, 'fraction_numerical': 0.03456356180433509, 'mean_word_length': 2.9657375145180023, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 1, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '588753', 'license': 'MIT', 'n_tokens_mistral': 2723, 'n_tokens_neox': 2566, 'n_words': 528, 'path': 'wgan.ipynb', 'repo_name': 'srirambandi/GAN', 'size': '6828'}