{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":3004,"databundleVersionId":861823,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Purpose of this notebook","metadata":{}},{"cell_type":"markdown","source":"The goal of this notebook is to show how a GAN can be trained in PyTorch to generate images. GAN are a little bit cumbersome to train, as the adversarial nature of the training makes the training process less stable than VAE or diffusion models for example","metadata":{}},{"cell_type":"markdown","source":"## Importing useful libraries","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nimport numpy as np\nimport pandas as pd\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torchsummary import summary","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:15.900435Z","iopub.execute_input":"2025-03-07T20:31:15.900762Z","iopub.status.idle":"2025-03-07T20:31:15.904988Z","shell.execute_reply.started":"2025-03-07T20:31:15.900733Z","shell.execute_reply":"2025-03-07T20:31:15.903906Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load dataset","metadata":{}},{"cell_type":"code","source":"df_train = pd.read_csv('/kaggle/input/digit-recognizer/train.csv')\n\nprint(df_train.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:15.906438Z","iopub.execute_input":"2025-03-07T20:31:15.906755Z","iopub.status.idle":"2025-03-07T20:31:17.563483Z","shell.execute_reply.started":"2025-03-07T20:31:15.906731Z","shell.execute_reply":"2025-03-07T20:31:17.562664Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X = df_train.drop(labels='label', axis=1).values\n# we do not need labels, we will generate digits randomly","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:17.565363Z","iopub.execute_input":"2025-03-07T20:31:17.565600Z","iopub.status.idle":"2025-03-07T20:31:17.647777Z","shell.execute_reply.started":"2025-03-07T20:31:17.565580Z","shell.execute_reply":"2025-03-07T20:31:17.647085Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BATCH_SIZE = 256\n\nplt.imshow(np.reshape((X/255)[1], (28, 28, 1)))\n\ntorch_X = torch.from_numpy(X/255)\n\ntorch_X = torch_X.view(-1, 1,28,28).float()\n\ntrain = torch.utils.data.TensorDataset(torch_X)\n\ntrain_loader = torch.utils.data.DataLoader(train,\n                                           batch_size = BATCH_SIZE,\n                                           shuffle = False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:17.648810Z","iopub.execute_input":"2025-03-07T20:31:17.649011Z","iopub.status.idle":"2025-03-07T20:31:18.068410Z","shell.execute_reply.started":"2025-03-07T20:31:17.648993Z","shell.execute_reply":"2025-03-07T20:31:18.067679Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model definition","metadata":{}},{"cell_type":"code","source":"class discriminator(nn.Module):\n    def __init__(self):\n        super(discriminator, self).__init__()\n        self.conv1 = nn.Conv2d(1,8,kernel_size=3,padding='same')\n        self.conv2 = nn.Conv2d(8,8,kernel_size=3,padding='same')\n        self.conv3 = nn.Conv2d(8,16,kernel_size=3,padding='same')\n        self.conv4 = nn.Conv2d(16,16,kernel_size=3,padding='same')\n        self.conv5 = nn.Conv2d(16,32,kernel_size=3,padding='same')\n        self.conv6 = nn.Conv2d(32,32,kernel_size=3,padding='same')\n\n        self.linear1 = nn.Linear(7*7*32, 64)\n        self.linear2 = nn.Linear(64, 1)\n\n    def forward(self, X):\n        X = F.leaky_relu(self.conv1(X))\n        X = X + F.leaky_relu(self.conv2(X))\n        X = nn.AvgPool2d(kernel_size=2)(X)\n        X = F.leaky_relu(self.conv3(X))\n        X = X + F.leaky_relu(self.conv4(X))\n        X = nn.AvgPool2d(kernel_size=2)(X)\n        X = F.leaky_relu(self.conv5(X))\n        X = X + F.leaky_relu(self.conv6(X))\n        X = nn.Flatten()(X)\n\n        X = F.leaky_relu(self.linear1(X))\n        output = nn.Sigmoid()(self.linear2(X))\n        \n        return output\n\n\nclass generator(nn.Module):\n    def __init__(self):\n        super(generator, self).__init__()\n        self.linear1 = nn.Linear(128, 49*32)\n        \n        self.conv1 = nn.Conv2d(32,32,kernel_size=3,padding='same')\n        self.conv2 = nn.Conv2d(32,16,kernel_size=3,padding='same')\n        self.conv3 = nn.Conv2d(16,16,kernel_size=3,padding='same')\n        self.conv4 = nn.Conv2d(16,8,kernel_size=3,padding='same')\n        self.conv5 = nn.Conv2d(8,8,kernel_size=3,padding='same')\n        self.conv6 = nn.Conv2d(8,1,kernel_size=3,padding='same')\n\n    def forward(self, X):        \n        X = F.leaky_relu(self.linear1(X))\n        X = torch.reshape(X, shape=(-1, 32, 7, 7))\n\n        X = X + F.leaky_relu(self.conv1(X))\n        X = F.leaky_relu(self.conv2(X))\n        X = nn.Upsample(scale_factor=2)(X)\n        X = X + F.leaky_relu(self.conv3(X))\n        X = F.leaky_relu(self.conv4(X))\n        X = nn.Upsample(scale_factor=2)(X)\n        X = X + F.leaky_relu(self.conv5(X))\n        X = F.leaky_relu(self.conv6(X))\n\n        return X","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:18.069084Z","iopub.execute_input":"2025-03-07T20:31:18.069308Z","iopub.status.idle":"2025-03-07T20:31:18.080900Z","shell.execute_reply.started":"2025-03-07T20:31:18.069290Z","shell.execute_reply":"2025-03-07T20:31:18.080164Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Loss definition","metadata":{}},{"cell_type":"code","source":"def discriminator_loss_fake(fake_images_rating):\n    return -torch.mean(torch.log(1 - fake_images_rating + 1e-8))\n\ndef discriminator_loss_true(true_images_rating):\n    return -torch.mean(torch.log(true_images_rating + 1e-8))\n    \ndef generator_loss(fake_images_rating):\n    return -torch.mean(torch.log(fake_images_rating + 1e-8))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:18.081730Z","iopub.execute_input":"2025-03-07T20:31:18.082010Z","iopub.status.idle":"2025-03-07T20:31:18.098746Z","shell.execute_reply.started":"2025-03-07T20:31:18.081981Z","shell.execute_reply":"2025-03-07T20:31:18.097996Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Training loop","metadata":{}},{"cell_type":"code","source":"def fit(generator, discriminator, train_data_loader, generator_optimizer, discriminator_optimizer, epochs, loss_fn1, loss_fn2, loss_fn3, device):\n    train_size = len(train_data_loader.dataset)\n    for epoch in range(epochs):\n        \n        generator.eval()\n        random_vector = torch.randn(size=(1, 128), device=device)\n        generated_image = generator(random_vector)\n        plt.imshow(generated_image.cpu().permute(0, 2, 3, 1).detach().numpy().reshape(28,28,1))\n        plt.show()\n            \n        print(\"epoch\", epoch+1,\"/\", epochs)\n        \n        generator.train()\n        for batch_idx, (X,) in enumerate(train_loader):\n            X = X.to(device)\n\n            # training the generator\n\n            for _ in range(1):\n                random_vector = torch.randn(size=(X.shape[0], 128), device=device)\n                generated_images = generator(random_vector)\n                fake_images_rating = discriminator(generated_images)\n    \n                generator_loss = loss_fn3(fake_images_rating)\n    \n                generator_optimizer.zero_grad()\n                generator_loss.backward()\n                generator_optimizer.step()\n\n            # training the discriminator\n\n            random_vector = torch.randn(size=(X.shape[0], 128), device=device)\n            generated_images = generator(random_vector)\n            fake_images_rating = discriminator(generated_images)\n\n            true_images_rating = discriminator(X)\n\n            discriminator_loss = loss_fn1(fake_images_rating) + loss_fn2(true_images_rating)\n\n            discriminator_optimizer.zero_grad()\n            discriminator_loss.backward()\n            discriminator_optimizer.step()\n\n            if batch_idx%25 == 0:\n                print(\"generator loss :    \", generator_loss.item())\n                print(\"discriminator loss :\", discriminator_loss.item())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:18.099555Z","iopub.execute_input":"2025-03-07T20:31:18.099818Z","iopub.status.idle":"2025-03-07T20:31:18.116522Z","shell.execute_reply.started":"2025-03-07T20:31:18.099798Z","shell.execute_reply":"2025-03-07T20:31:18.115714Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model compilation","metadata":{}},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"Using {device} device\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:18.117249Z","iopub.execute_input":"2025-03-07T20:31:18.117500Z","iopub.status.idle":"2025-03-07T20:31:18.135018Z","shell.execute_reply.started":"2025-03-07T20:31:18.117480Z","shell.execute_reply":"2025-03-07T20:31:18.134368Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model_generator = generator()\nmodel_discriminator = discriminator()\n\ngenerator_optimizer = torch.optim.RMSprop(model_generator.parameters(), lr=1e-4)\ndiscriminator_optimizer = torch.optim.RMSprop(model_discriminator.parameters(), lr=1e-4)\nloss_fn1 = discriminator_loss_fake\nloss_fn2 = discriminator_loss_true\nloss_fn3 = generator_loss\n\nepochs = 100","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:18.136562Z","iopub.execute_input":"2025-03-07T20:31:18.136753Z","iopub.status.idle":"2025-03-07T20:31:18.155100Z","shell.execute_reply.started":"2025-03-07T20:31:18.136735Z","shell.execute_reply":"2025-03-07T20:31:18.154467Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model_generator.to(device)\nmodel_discriminator.to(device)\n\nsummary(model_generator, (128,))\nsummary(model_discriminator, (1, 28, 28))\n\nfit(model_generator, model_discriminator, train_loader,\n    generator_optimizer, discriminator_optimizer, epochs,\n    loss_fn1, loss_fn2, loss_fn3, device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-07T20:31:18.156211Z","iopub.execute_input":"2025-03-07T20:31:18.156480Z","iopub.status.idle":"2025-03-07T20:34:39.795292Z","shell.execute_reply.started":"2025-03-07T20:31:18.156453Z","shell.execute_reply":"2025-03-07T20:34:39.794615Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Results interpretation","metadata":{}},{"cell_type":"markdown","source":"The model learns quite fast to generate samples of decent quality. We observe that the improvment in generation quality seems uncorrelated to the loss values during the training process. In order to evaluate different GAN quality, this is a problem as comparing the loss is the easiest and fastest way to make comparisons. \\\nIn order to really compare two models, one would have to define a metric that can evaluate the intrisic quality of generated samples, like FID. \\\nFor some reason, the Adam optimizer does not work in this setting. My guess is that due to the non stationary target distribution during training, the first momentum of Adam is a problem. This could explain why RMSProp, which only uses the second moment of Adam, works. This idea is supported by this paper : https://sgo-workshop.github.io/CameraReady2019/11.pdf\nThis is just a guess for now.\nGAN have been a very used algorithm to generate samples from a distribution. However, they have been replaced by diffusion models. These type of models are much more complex mathematically, but their training is not adversarial which makes it much more stable. We will train in a future notebook a diffusion model.","metadata":{}}]}