{"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 VAE can be trained in PyTorch to generate images. VAE are easy to train, but have low generation quality compared to other gen AI algorithms.","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\n\nfrom sklearn.model_selection import train_test_split","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-02-25T18:31:12.807858Z","iopub.execute_input":"2025-02-25T18:31:12.808169Z","iopub.status.idle":"2025-02-25T18:31:12.812545Z","shell.execute_reply.started":"2025-02-25T18:31:12.808145Z","shell.execute_reply":"2025-02-25T18:31:12.811584Z"}},"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-02-25T18:31:12.867851Z","iopub.execute_input":"2025-02-25T18:31:12.868157Z","iopub.status.idle":"2025-02-25T18:31:14.620184Z","shell.execute_reply.started":"2025-02-25T18:31:12.868130Z","shell.execute_reply":"2025-02-25T18:31:14.619290Z"}},"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-02-25T18:31:14.621315Z","iopub.execute_input":"2025-02-25T18:31:14.621631Z","iopub.status.idle":"2025-02-25T18:31:14.703870Z","shell.execute_reply.started":"2025-02-25T18:31:14.621598Z","shell.execute_reply":"2025-02-25T18:31:14.703188Z"}},"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-02-25T18:31:14.705528Z","iopub.execute_input":"2025-02-25T18:31:14.705863Z","iopub.status.idle":"2025-02-25T18:31:15.118066Z","shell.execute_reply.started":"2025-02-25T18:31:14.705829Z","shell.execute_reply":"2025-02-25T18:31:15.116900Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model definition","metadata":{}},{"cell_type":"code","source":"class VAE_encoder(nn.Module):\n    def __init__(self):\n        super(VAE_encoder, self).__init__()\n        self.conv1 = nn.Conv2d(1,32,kernel_size=3,padding='same')\n        self.conv2 = nn.Conv2d(32,32,kernel_size=3,padding='same')\n        self.conv3 = nn.Conv2d(32,64,kernel_size=3,padding='same')\n        self.conv4 = nn.Conv2d(64,64,kernel_size=3,padding='same')\n        self.conv5 = nn.Conv2d(64,128,kernel_size=3,padding='same')\n        self.conv6 = nn.Conv2d(128,128,kernel_size=3,padding='same')\n\n        self.linear1 = nn.Linear(7*7*128, 64)\n        self.linear2 = nn.Linear(64, 2)\n        self.linear3 = nn.Linear(64, 2)\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        mean = self.linear2(X)\n        log_var = self.linear3(X)\n\n        return mean, log_var\n\n\nclass VAE_decoder(nn.Module):\n    def __init__(self):\n        super(VAE_decoder, self).__init__()\n        self.linear1 = nn.Linear(2, 49*32)\n\n        self.linear2 = nn.Linear(2, 64)\n        self.linear3 = nn.Linear(64, 28*28)\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.linear2(X))\n        #X = F.leaky_relu(self.linear3(X))\n        #return torch.reshape(X, shape=(-1, 1, 28, 28))\n\n        \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\n\nclass VAE(nn.Module):\n    def __init__(self, device):\n        super(VAE, self).__init__()\n        self.encoder = VAE_encoder()\n        self.decoder = VAE_decoder()\n        self.device = device\n\n    def forward(self, X):\n        mean, log_var = self.encoder(X)\n        z = torch.randn(size=(2,), device=self.device)*torch.exp(log_var/2) + mean\n        Y = self.decoder(z)\n\n        return mean, log_var, Y","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-25T18:31:15.119392Z","iopub.execute_input":"2025-02-25T18:31:15.119638Z","iopub.status.idle":"2025-02-25T18:31:15.132623Z","shell.execute_reply.started":"2025-02-25T18:31:15.119607Z","shell.execute_reply":"2025-02-25T18:31:15.131948Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Loss definition","metadata":{}},{"cell_type":"code","source":"def reconstruction_loss(y_true, y_pred):\n    return torch.sum((y_true - y_pred)**2)\n\ndef KL_loss(mean, log_var):\n    return - torch.mean(1 + log_var - torch.square(mean) - torch.exp(log_var)) / 2\n\ndef complete_loss(y_true, y_pred, mean, log_var):\n    return reconstruction_loss(y_true, y_pred) + KL_loss(mean, log_var)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-25T18:31:15.133492Z","iopub.execute_input":"2025-02-25T18:31:15.133807Z","iopub.status.idle":"2025-02-25T18:31:15.153247Z","shell.execute_reply.started":"2025-02-25T18:31:15.133777Z","shell.execute_reply":"2025-02-25T18:31:15.152511Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Training loop","metadata":{}},{"cell_type":"code","source":"def fit(model, train_data_loader, optimizer, epochs, loss_fn, device):\n    history = {}\n    train_size = len(train_data_loader.dataset)\n    for epoch in range(epochs):\n        \n        model.eval()\n        for (X,) in train_loader:\n            X = X.to(device)\n            mean, std, Y = model(X[1:2])\n            plt.imshow(X[1:2].cpu().permute(0, 2, 3, 1).detach().numpy().reshape(28,28,1))\n            plt.show()\n            plt.imshow(Y.cpu().permute(0, 2, 3, 1).detach().numpy().reshape(28,28,1))\n            plt.show()\n            break\n            \n        model.train()\n        print(\"epoch\", epoch+1,\"/\", epochs)\n        model.train()\n        train_loss_mean = 0\n        for batch_idx, (X,) in enumerate(train_loader):\n            X = X.to(device)\n            \n            mean, log_var, y_pred = model(X)\n            loss = loss_fn(X, y_pred, mean, log_var)\n\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n\n            train_loss_mean += loss.item()\n\n        print(\"loss :\", train_loss_mean/train_size)\n\n\n    return history","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-25T18:31:15.153994Z","iopub.execute_input":"2025-02-25T18:31:15.154215Z","iopub.status.idle":"2025-02-25T18:31:15.169794Z","shell.execute_reply.started":"2025-02-25T18:31:15.154187Z","shell.execute_reply":"2025-02-25T18:31:15.168980Z"}},"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-02-25T18:31:15.170387Z","iopub.execute_input":"2025-02-25T18:31:15.170618Z","iopub.status.idle":"2025-02-25T18:31:15.187861Z","shell.execute_reply.started":"2025-02-25T18:31:15.170598Z","shell.execute_reply":"2025-02-25T18:31:15.187053Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = VAE(device)\n\noptimizer = torch.optim.Adam(model.parameters(), lr=3e-4)\nloss_fn = complete_loss\n\nepochs = 20","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-25T18:31:15.189529Z","iopub.execute_input":"2025-02-25T18:31:15.189797Z","iopub.status.idle":"2025-02-25T18:31:15.210234Z","shell.execute_reply.started":"2025-02-25T18:31:15.189765Z","shell.execute_reply":"2025-02-25T18:31:15.209446Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.to(device)\n\nsummary(model, (1, 28, 28))\n\nhistory = fit(model, train_loader, optimizer, epochs, loss_fn, device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-25T18:31:15.211095Z","iopub.execute_input":"2025-02-25T18:31:15.211398Z","iopub.status.idle":"2025-02-25T18:34:12.922417Z","shell.execute_reply.started":"2025-02-25T18:31:15.211368Z","shell.execute_reply":"2025-02-25T18:34:12.921630Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Results interpretation","metadata":{}},{"cell_type":"markdown","source":"We observe that the model learns quickly to reconstruct samples of decent quality, albeit quite blurred. Apart from this expected result, I learned something of value : if instead of using a leaky relu activation function for the last layer of the decoder, I use a sigmoid or a relu, the training is veery slow and almost always gets stuck generatinf black images. \\\nTo me, this is because of the following reason : The images we try to generate are almost completely black, apart of a few white pixels. Initially, the model seems to learn to generate completely black images from the black background. However, relu and sigmoid are saturating function, i.e. the gradient is very small (0 for relu) when their output is $0$. Thus, the network is stuck here because the gradient is more or less $0$ when backpropagating, and therefore it cannot learn the useful features to generate the white pixels. The leaky relu does not have this problem, as its gradient is non zero when its ouput $0$ or less. \\\nThus, I retain that when dealing with a problem where the output is very biaised toward one value, I must make sure that the model will be able to learn even if it outputs only this biased value. In my case, this value is $0$. To make sure the model can still learn, one must check that the gradient is not near $0$ when the model only output the biased value.\n\n\n\n","metadata":{}}]}