TP_GAN_cat_image_generator.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5wczghxT56FL"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "L'objectif de ce TP est de coder un GAN pour générer des petits chats. Il abordera les notions principales telles que: les CNN, convolutions transposées, batchnorm, binary crossentropy, "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jyemBO3ISvqi"
      },
      "source": [
        "## Les imports"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zOPaFwyAIQJl"
      },
      "source": [
        "\n",
        "On va d'abord importer les librairies nécessaires dont les classiques numpy, matplotlib ainsi que Pytorch (quand même....) et certains de ses modules et classes utilisés utilisés souvent (Dataset, nn)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dvBgYvzTNL8g"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torchvision.io import read_image\n",
        "from torch import nn\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython.display import clear_output\n",
        "import gdown\n",
        "import os\n",
        "from tqdm import tqdm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4flw12gySJn6"
      },
      "source": [
        "## Téléchargement du dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lz7SLVnrIQJp"
      },
      "source": [
        "\n",
        "Importons maintenant le dataset. On va télécharget le dataset de chats depuis ce [drive](https://drive.google.com/uc?id=1F9I7iDmQ_I9Qsrav1UXlD4OiIBVSU5sl) avec la commande gdown (ou à la mano si vous préférez).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MLZNgqQ7Obze",
        "outputId": "42773885-4c2b-4d15-9c21-bcc2c0a4a6f9"
      },
      "outputs": [],
      "source": [
        "# Téléchargement du dataset\n",
        "url = 'https://drive.google.com/uc?id=1F9I7iDmQ_I9Qsrav1UXlD4OiIBVSU5sl'\n",
        "output = 'dataset.tgz'\n",
        "if not os.path.exists(output):\n",
        "    gdown.download(url, output, quiet=False)\n",
        "\n",
        "# Dézippage du dataset\n",
        "\n",
        "def unzip(zip_file, dest_dir):\n",
        "    import zipfile\n",
        "    with zipfile.ZipFile(zip_file, 'r') as zip_ref:\n",
        "        zip_ref.extractall(dest_dir)\n",
        "\n",
        "unzip('dataset.tgz', './')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zqylv3scUY76"
      },
      "source": [
        "## Quelques paramètres généraux"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a9kEINnMIQJq"
      },
      "source": [
        "\n",
        "Les paramètres peuvent être modifiés pour tester un peu (sauf peut-être la taille de l'image pour ce dataset).\n",
        "\n",
        "Pour ce dataset, on peut prendre une batch size pas trop grande sinon Colab crash (ou votre PC). De même, la dimension de l'espace latent peut être ajustée, ici vu que les chats c'est pas si simple, environ 100 c'est bien.\n",
        "\n",
        "Pour charger le dataset, on parcourt notre dossier dataset en chargeant les images qui s'y trouvent. Les images ont leurs pixels entre 0 et 255 qu'on va renormaliser dès la création du dataset entre -1 et 1, ce qui est plus adapté pour les réseaux de neurones et bien pour des GAN car on a une moyenne nulle.\n",
        "\n",
        "On va aussi les afficher parce que c'est bien de savoir sur quoi on travaille :\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DFtB5YCWO2Pm"
      },
      "outputs": [],
      "source": [
        "BATCH_SIZE = 64\n",
        "LATENT_DIM = 100\n",
        "IMG_SHAPE = (64,64,3)\n",
        "\n",
        "x_train = []\n",
        "batch = []\n",
        "i = 0\n",
        "\n",
        "for files in os.listdir(\"./dataset\"):\n",
        "  image = read_image(f\"./dataset/{files}\")\n",
        "\n",
        "  # On normalise les données entre -1 et 1\n",
        "  image = image/255\n",
        "  image = image*2 -1\n",
        "\n",
        "  if i < BATCH_SIZE:\n",
        "    batch.append(image)\n",
        "    i += 1\n",
        "  else:\n",
        "    x_train.append(torch.stack(batch))\n",
        "    batch = []\n",
        "    i = 0\n",
        "\n",
        "\n",
        "x_train = torch.stack(x_train)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mosCTmbq4ixg",
        "outputId": "bef60ed6-b27b-4e4f-87c0-b1cbbecce02b"
      },
      "outputs": [],
      "source": [
        "print(x_train.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 961
        },
        "id": "ZF6nSfKQP-IM",
        "outputId": "d0732b42-911d-40a1-ac20-f916f36eca59"
      },
      "outputs": [],
      "source": [
        "test_batch = next(iter(x_train))\n",
        "test_batch = test_batch.permute(0,2,3,1)\n",
        "fig = plt.figure(figsize=(12,12))\n",
        "for i in range(25):\n",
        "  plt.subplot(5,5,i+1)\n",
        "  plt.axis('off')\n",
        "  # On oublie pas de faire image * 0.5 + 0.5 pour revenir dans [0,1]\n",
        "  plt.imshow(test_batch[i]*0.5+0.5)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cWFiDuRM6ocP"
      },
      "source": [
        "## Pytorch et les GPUs...\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lzToi7YP3iX7"
      },
      "source": [
        "\n",
        "Pytorch a besoin qu'on lui précise avec quel appareil on travaille (CPU ou GPU par exemple).\n",
        "\n",
        "Du coup, ici on choisit le meilleur appareil disponible (généralement GPU, sinon CPU)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VYHvZDsK7UPY",
        "outputId": "03a48a2e-1988-4c68-8f88-65d09db1c0a0"
      },
      "outputs": [],
      "source": [
        "device = (\n",
        "    \"cuda\"\n",
        "    if torch.cuda.is_available()\n",
        "    else \"mps\"\n",
        "    if torch.backends.mps.is_available()\n",
        "    else \"cpu\"\n",
        ")\n",
        "print(f\"Using {device} device\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IZZBKWamPNxO"
      },
      "source": [
        "## Le discriminateur"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J7ZondMqIQJt"
      },
      "source": [
        "\n",
        "On va d'abord faire le discriminateur.\n",
        "\n",
        "Pour cela on va définir le modèle en utilisant `nn.Sequential()`, comme dans les autres TP.\n",
        "\n",
        "En entrée, on va avoir une image de dimension (3,64,64) (c'est en couleur pour rappel donc il y a bien 3 canaux).\n",
        "\n",
        "Pour vous aider, un petit rappel des blocs à mettre:\n",
        "- `Conv2D(in_features,out_features,kernel_size,strides,padding)`\n",
        "\n",
        "  Typiquement, on utilise une taille de kernel de 3 et des filtres qui sont des puissances de 2 (32,64,128...).\n",
        "  \n",
        "  Pour rappel, le padding `'same'` équivaut à rajouter des zéros pour garder la même taille d'image en sortie (si les strides sont de 1) et avec `'valid'` on en rajoute pas. `'same'` est donc conseillé ici pour éviter les surprises (et de se torturer la tête sur la taille en sortie même si vous savez évidemment que c'est $\\frac{n+2p-f}{s}+1$). A savoir que sur Pytorch cela ne marche pas forcément, et que vous allez peut être devoir la préciser à la main.\n",
        "   \n",
        "   On ne met pas l'activation tout de suite car on ajoute d'abord de la batch normalization pour renormaliser le batch en 0 et ainsi profiter au maximum de la non-linéarité de la ReLU en ce point.\n",
        "\n",
        "  Pour rappel, on augmente le nombre de filtres au fur et à mesure de l'architecture dans le discriminateur.\n",
        "- `BatchNorm2D(in_features)`\n",
        "- `LeakyReLU(alpha=0.2)` (pente de la leaky relu dans le domaine $]-\\infty,0]$)\n",
        "\n",
        "Mettez 3-4 blocs comme ça et faites des essais.\n",
        "\n",
        "Ici pas de `Pooling` car on réduit la taille des images directement avec du stride (typiquement 2 à chaque Convolution en comptant bien la taille qu'on obtient à la fin).\n",
        "\n",
        "En sortie on veut une dimension (1) grâce à une `Flatten` puis des `Dense` en oubliant pas la sigmoïde à la fin.\n",
        "\n",
        "Si vous voulez pas trop vous embêter pour votre `forward`, vous pouvez utiliser un nn.Sequential dans votre modèle.\n",
        "\n",
        "*Alternative* :\n",
        "- C'est aussi possible de bien calculer la taille de l'image pour terminer par une `Convolution` avec une sortie de dimension (1,1,1) puis une `Flatten` (un peu mieux, dit \"Full Convolutional\").\n",
        "\n",
        "  **Exemple** :\n",
        "  - Image de (64,64,3) -> 4 Blocs de convolution avec du stride de 2 : $64/2^4$ -> Feature map de taille (4,4,nombre_de_filtres).\n",
        "\n",
        "    On peut ensuite finir avec une `Conv2D` avec une taille de filtre de 4 et du padding `'valid'` pour juste faire une combinaison de tous les pixels restants, ce qui nous donne bien notre unique pixel de dimension (1,1,1) sans oublier la sigmoïde et on peut ensuite `Flatten` tout ça).\n",
        "\n",
        "\n",
        "Juste ici, petite antisèche de l'architecture si vous séchez. Attention l'écriture est en pseudo-code, il faut l'adapter à Pytorch.\n",
        "<details>\n",
        "<summary>Antisèche</summary>\n",
        "Conv2D(32 filtres,kernel taille 3,stride 2,padding 1) -> BN -> LR <br />  \n",
        "-> Conv2D(64,kernel 3,stride 2,padding 1) -> BN -> LR <br />  \n",
        "-> Conv2D(128,3,2,1) -> BN -> LR <br />  \n",
        "-> Conv2D(256,3,2,1) -> BN -> LR <br />  \n",
        "-> Conv2D(1,4,1,1) -> Sigmoïde -> Flatten ou Flatten -> Linear(1) -> Sigmoïde\n",
        "</details>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TasqSOiOk24X"
      },
      "outputs": [],
      "source": [
        "discriminator = nn.Sequential(\n",
        "    #dim : (3, 64, 64)\n",
        "    nn.Conv2d(3, 64, 3, 2, 1), #Un stride de 2 permet de diviser par 2 la dimension de l'image en sortie, donc pas besoin de MaxPooling\n",
        "                               #Le nombre de filtre à chaque étape est in_features*out_features. Ici, on applique un fiftre de profondeur 3 (les 3\n",
        "                               #filtres étant différents) (3 au début car RVB),\n",
        "                               #puis on somme les 3 valeurs pour obtenir un pixel de la première couche (des 64 couches qu'on doit créer).\n",
        "                               #On a donc in_features * out_features filtre différents à chaque couche.\n",
        "    nn.BatchNorm2d(64),\n",
        "    nn.LeakyReLU(0.2),\n",
        "    #dim : (64, 32, 32)\n",
        "\n",
        "    nn.Conv2d(64, 128, 3, 2, 1),\n",
        "    nn.BatchNorm2d(128),\n",
        "    nn.LeakyReLU(0.2),\n",
        "    #dim : (128, 16, 16)\n",
        "\n",
        "    nn.Conv2d(128, 256, 3, 2, 1),\n",
        "    nn.BatchNorm2d(256),\n",
        "    nn.LeakyReLU(0.2)  ,\n",
        "    #dim : (256, 8, 8)\n",
        "\n",
        "    nn.Flatten(),\n",
        "    nn.Linear(256 * 8 * 8, 1), #Revient à choisir une bonne taille de conv pour se ramener à 1x1x1 = nn.Conv2d(128, 1, 3, 2, 1)\n",
        "    nn.Sigmoid()\n",
        "\n",
        ").to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZaDZbie7Qw2D"
      },
      "source": [
        "## Le générateur"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b0X7PhKJIQJu"
      },
      "source": [
        "\n",
        "On va faire le générateur.\n",
        "En entrée, il va prendre un vecteur de l'espace latent de dimension `latent_dim`.\n",
        "\n",
        "On augmente la taille de se vecteur grâce à une couche dense.\n",
        "\n",
        "`DensLinear(4 x 4 x 1024,input_shape=(latent_dim, ))`\n",
        "\n",
        "Pour redimensionner notre vecteur en une \"image\" de taille 4 x 4 x 1024, on va pouvoir utilser la couche `Unflatten()`\n",
        "\n",
        "\n",
        "On peut aussi directement redimensionner le vecteur de l'espace latent sans la couche Dense.\n",
        "\n",
        "\n",
        "Pour vous aider, un petit rappel des blocs à mettre ensuite:\n",
        "- `Conv2DTranspose(in_features,out_features,kernel_size,strides,padding,output_padding)`\n",
        "\n",
        "  Pour rappel, on diminue le nombre de filtres au fur et à mesure de l'architecture dans le générateur (512,256,...)\n",
        "\n",
        "- `BatchNorm2D(in_features)`\n",
        "- `ReLU()`\n",
        "\n",
        "Mettez 2-3 blocs comme ça et faites des essais aussi. **Faites bien attention** à la taille de vos images tout au long de l'architecture pour bien avoir la taille d'image finale voulue.\n",
        "\n",
        "**Exemple**:\n",
        "\n",
        "- On a transformé notre vecteur latent en image de dimension (1024,4,4). Pour avoir du (3,64,64), il faut donc 4 blocs de `Conv2DTranspose` (avec padding) pour avoir $4*2^4=64$. Libre à vous de changer l'entrée pour mettre le nombre de blocs que vous voulez.\n",
        "\n",
        "\n",
        "Ici aussi pas de `UpSampling2D` car on augmente la taille avec du stride (2 aussi souvent) à chaque convolution transposée.\n",
        "\n",
        "En sortie on veut une dimension (64,64,3) grâce à une `Conv2D` avec 3 filtres, on oublie pas de prendre une activation en tangente hyperbolique `'tanh'` pour avoir des pixels dans [-1,1].\n",
        "\n",
        "<details>\n",
        "<summary>Antisèche</summary>\n",
        "  (Dans le doute on peut flatten ici)\n",
        "  Dense(100,4*4*1024 neurones) -> Unflatten(en (1024,4,4))\n",
        "  -> ConvTransposée(256 filtres,kernel taille 3,stride 2,padding 1) -> BN -> ReLU <br />  \n",
        "  -> ConvTransposée(128,kernel 3,stride 2,padding 1) -> BN -> ReLU <br />  \n",
        "  -> ConvTransposée(64,3,2,1) -> BN -> ReLU <br />  \n",
        "  -> ConvTransposée(32,3,2,1) -> BN -> ReLU <br />  \n",
        "  -> ConvTransposée(3,3,1,1) -> Tanh\n",
        "</details>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eHCUgA6lKuwj"
      },
      "outputs": [],
      "source": [
        "generator = nn.Sequential(\n",
        "#On a au début un vecteur de bruit. Il faut le ramener à une image de profondeur qui va diminuer. C'est l'inverse d'un réseau de convolution\n",
        "#ce réseau. On fait donc un unflatten. Conv2DTranspose fait une opération d'agrandissement de l'image (stride qui encapsule les valeurs dans des 0).\n",
        "#et diminution de la profondeur ET une opération de convolution (il faut bien des filtres pr l'entrainement) qui ne change pas la taille.\n",
        "#Ainsi la taille de l'image augmente et la profondeur diminue.\n",
        "\n",
        "   #dim : (latent_dim)\n",
        "   nn.Linear(LATENT_DIM, 1024*4*4), #on a un vecteur de bruit de taille 100, on le trainsforme en taille 1024*4*4 pr le flatten.\n",
        "   nn.Unflatten(1, (1024, 4, 4)),\n",
        "   #dim : (4, 4, 1024)\n",
        "   nn.ConvTranspose2d(1024, 512, 3, 2, 1, 1), #pour trouver le padding et le outputpadding, on regarde sur la doc Pytorch de convTranspose2D\n",
        "   nn.BatchNorm2d(512),\n",
        "   nn.ReLU(),\n",
        "   #dim : (8, 8, 512)\n",
        "\n",
        "   nn.ConvTranspose2d(512, 256, 3, 2, 1, 1),\n",
        "   nn.BatchNorm2d(256),\n",
        "   nn.ReLU(),\n",
        "   #dim : (16, 16, 256)\n",
        "\n",
        "   nn.ConvTranspose2d(256, 128, 3, 2, 1, 1),\n",
        "   nn.BatchNorm2d(128),\n",
        "   nn.ReLU(),\n",
        "   #dim : (32, 32, 128)\n",
        "\n",
        "   nn.ConvTranspose2d(128, 3, 3, 2, 1, 1),\n",
        "   nn.BatchNorm2d(3),\n",
        "   nn.Tanh(),\n",
        "   #dim : (64, 64, 3)\n",
        "\n",
        ").to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u0Fbk2BnWJOO"
      },
      "source": [
        "## Le train step"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SrS02DU5IQJw"
      },
      "source": [
        "\n",
        "Cette partie est assez importante car elle permet de comprendre comment on entraîne vraiment un GAN, c'est-à-dire à quoi on compare les sorties pour entraîner correctement le discriminateur et le générateur. Ici, pas de `model.fit` malheureusement.\n",
        "On va définir un `train_step`, c'est à dire ce qu'on va faire comme opérations à chaque batch :\n",
        "\n",
        "Entraîner le discriminateur :\n",
        "\n",
        "- Générer des images fausses à partir de bruit gaussien et en prédire les labels : on a besoin ici d'un vecteur latent de dimension `(batch_size,latent_dim,1,1)`. Ensuite, on fait passer ce bruit dans le générateur pour obtenir des fausses images. Enfin, on récupère la sortie du discriminateur sur celles-ci.\n",
        "\n",
        "\n",
        "```\n",
        "noise = torch.randn(shape).to(device) (merci pytorch...)\n",
        "fake_images = generator(...)\n",
        "fake_predictions = discriminator(...)\n",
        "```\n",
        "\n",
        "\n",
        "- Prendre des images vraies du dataset et en prédire aussi les labels. Donc la sortie du discriminateur sur les vraies images.\n",
        "\n",
        "\n",
        "```\n",
        "real_predictions = ...\n",
        "```\n",
        "\n",
        "\n",
        "- Calculer la loss en comparant les prédictions sur les fausses avec des 0 et les prédictions sur les vraies avec des 1.\n",
        "\n",
        "\n",
        "\n",
        "La binary crossentropy prend en argument les labels visés puis ceux prédits.\n",
        "Pour avoir des 1 ou des 0 : `torch.ones_like(tensor)` ou `torch.zeros_like(tensor)` avec `tensor` le tenseur dont on veut imiter la taille (par exemple les labels issus des prédictions, vu que c'est à ça que l'on va vouloir les comparer). Dans l'exemple suivant, cela corresponds aux `true_labels` et les prédictions précédentes aux `predictions`.\n",
        "```\n",
        "discriminator_loss_on_real = loss(real_labels,real_predictions)\n",
        "discriminator_loss_on_fake = loss(fake_labels,fake_predictions)\n",
        "discriminator_loss = discriminator_loss_on_real + discriminator_loss_on_fake\n",
        "```\n",
        "\n",
        "On peut éventuellement diviser la loss du discriminateur par 2.\n",
        "\n",
        "- Calculer les gradients en fonction de la loss calculée et les différents paramètres du modèle\n",
        "- On applique les gradients calculés avec l'optimisateur choisi\n",
        "\n",
        "Entraîner le générateur (presque la même chose):\n",
        "- Générer des images fausses à partir de bruit gaussien et en prédire les labels.\n",
        "- Calculer la loss en comparant les prédictions sur les fausses avec des 1 (on veut tromper le discriminateur).\n",
        "- Calculer les gradients en fonction de la loss calculée et les différents paramètres du modèle\n",
        "- On applique les gradients calculés avec l'optimisateur choisi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "moZPLFqPSnz9"
      },
      "outputs": [],
      "source": [
        "def train_step(real_images, generator, discriminator, loss, g_opt, d_opt):\n",
        "  global LATENT_DIM\n",
        "  batch_size = real_images.size(dim=0)\n",
        "\n",
        "\n",
        "  noise = torch.randn(size = (BATCH_SIZE,LATENT_DIM)).to(device)\n",
        "  fake_images = generator(noise)\n",
        "\n",
        "  ### On entraine le discriminateur\n",
        "\n",
        "  discriminator.zero_grad()\n",
        "\n",
        "  fake_images_predictions = discriminator(fake_images)\n",
        "  real_images_predictions = discriminator(real_images)\n",
        "\n",
        "  fake_labels = torch.zeros_like(fake_images_predictions) #on veut autant de 0 que d'images dans le batch\n",
        "  fake_labels = fake_labels.type_as(real_images)\n",
        "  real_labels = torch.ones_like(real_images_predictions)\n",
        "  real_labels = real_labels.type_as(real_images)\n",
        "\n",
        "  ### On calcule la loss pour les vraies et les fausses images\n",
        "\n",
        "  fake_loss = loss(fake_images_predictions, fake_labels)\n",
        "  real_loss = loss(real_images_predictions, real_labels)\n",
        "\n",
        "  disc_loss = (fake_loss + real_loss)/2\n",
        "\n",
        "  ### Backpropagation\n",
        "\n",
        "  disc_loss.backward()\n",
        "  d_opt.step()\n",
        "  # d_opt.zero_grad()\n",
        "\n",
        "\n",
        "\n",
        "  ### On entraine le générateur\n",
        "\n",
        "  generator.zero_grad()\n",
        "\n",
        "  fake_images = generator(noise)\n",
        "  fake_images_predictions = discriminator(fake_images)\n",
        "\n",
        "  real_labels = torch.ones_like(fake_images_predictions)\n",
        "  real_labels = real_labels.type_as(real_images)\n",
        "\n",
        "\n",
        "  #Au début le discriminateur ne retourne que des 0 pour les images (moches) du générateur\n",
        "  #On entraine le générateur pour que le discriminateur retourne que des 1 (des trucs proches des vrais images de chat)\n",
        "  #Donc on entraine QUE le générateur sur la loss du discriminateur par rapport à 1 (les vrais beaux chats)\n",
        "\n",
        "\n",
        "  ### On calcule la loss pour le générateur\n",
        "\n",
        "  gen_loss = loss(fake_images_predictions, real_labels)\n",
        "\n",
        "  ### Backpropagation\n",
        "\n",
        "  gen_loss.backward()\n",
        "  g_opt.step()\n",
        "  # g_opt.zero_grad()\n",
        "  return gen_loss,disc_loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wPUvdewb7Ak5"
      },
      "source": [
        "## Le train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92gcoGyJIQJx"
      },
      "source": [
        "On code la fonction d'entraînement principale.\n",
        "Il reste à compléter la loss et les optimisateurs à utiliser, typiquement ici la `BCELoss` et `Adam(parameters,lr=2e-4,betas = (0.5, 0.5)`.\n",
        "\n",
        "Dans un premier temps, on peut prendre les mêmes optimiseurs pour les deux quitte à adapter pour tester après (changer le learning rate par exemple pour rééquilibrer un peu l'entraînement)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KOXC7WKmTI_Z"
      },
      "outputs": [],
      "source": [
        "def train(dataset,generator,discriminator,epochs,fixed_noise):\n",
        "\n",
        "  ###A compléter###\n",
        "  loss = nn.BCELoss() #on définit enfin la loss\n",
        "  g_opt = torch.optim.Adam(generator.parameters(), lr=2e-4, betas = (0.5, 0.5))\n",
        "  d_opt = torch.optim.Adam(discriminator.parameters(), lr=2e-4, betas = (0.5, 0.5))\n",
        "\n",
        "  Lgen_loss = []\n",
        "  Ldisc_loss = []\n",
        "  X = []\n",
        "  j = 0\n",
        "\n",
        "  for epoch in range(epochs):\n",
        "    progress_bar = tqdm(dataset)\n",
        "    ## Vu que c'est un dataloader, on ne peut itérer directement dessus avec son indice. On va juste prendre à chaque fois batch par batch.\n",
        "    for _,image_batch in enumerate(progress_bar):\n",
        "        j += 1\n",
        "        real_images = image_batch.to(device)\n",
        "        gen_loss, disc_loss = train_step(real_images, generator, discriminator, loss, g_opt, d_opt)\n",
        "\n",
        "        X.append(j)\n",
        "        Lgen_loss.append(gen_loss.item())\n",
        "        Ldisc_loss.append(disc_loss.item())\n",
        "\n",
        "        progress_bar.set_description(f\"Epoch {epoch+1}/{epochs} | Gen Loss: {gen_loss} | Disc Loss: {disc_loss}\")\n",
        "\n",
        "    clear_output(wait=False)\n",
        "    generate_and_save_plots(X, Lgen_loss, Ldisc_loss) # Définie après, pour générer les courbes des loss\n",
        "    summarize_performance(generator,fixed_noise) # Définie après, pour afficher les images générées"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4iSDZildZALv"
      },
      "source": [
        "## L'affichage à chaque epoch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpNuUbbgIQJy"
      },
      "source": [
        "\n",
        "Petite fonction qui affiche les images obtenues à chaque epoch. On va **afficher** 25 images avec la même seed (toujours du même vecteur latent) pour voir l'amélioration progressive de l'image. Ce n'est pas de l'overfitting sur un seul vecteur car on **entraîne** bien à partir de vecteurs différents à chaque fois avant.\n",
        "\n",
        "On affiche aussi après chaque epoch les courbes des loss du générateur et du discriminateur pour suivre l'entraînement et l'arrêter si on constate un souci."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2D5j4ZlLXaOh"
      },
      "outputs": [],
      "source": [
        "def summarize_performance(generator,fixed_noise):\n",
        "  fake_images = generator(fixed_noise).detach().cpu()\n",
        "  fake_images = torch.permute(fake_images, (0, 2, 3, 1))\n",
        "  fig = plt.figure(figsize=(12,12))\n",
        "  for i in range(25):\n",
        "    plt.subplot(5,5,i+1)\n",
        "    plt.axis('off')\n",
        "    plt.imshow(fake_images[i]*0.5+0.5)\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SDbwKnC1pLFN"
      },
      "outputs": [],
      "source": [
        "def generate_and_save_plots(X, Lgen_loss, Ldisc_loss):\n",
        "    fig = plt.figure(figsize=(4,4))\n",
        "    plt.plot(X,Lgen_loss, label = 'gen_loss')\n",
        "    plt.plot(X,Ldisc_loss, label = 'disc_loss')\n",
        "    plt.legend()\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dAx8pu5P7sOR"
      },
      "source": [
        "## Ici vous lancez tout!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tNStSQXmIQJ0"
      },
      "source": [
        "Choisissez le nombre d'epochs que vous voulez. Ici une vingtaine d'epochs peut suffire, on va pas attendre 3h quand même."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "vKElP8cy8dil",
        "outputId": "7492cfb3-1f4a-41e4-a2d3-e28c9d8593cc"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 20\n",
        "fixed_noise = torch.randn(size=(BATCH_SIZE,LATENT_DIM)).to(device)\n",
        "train(x_train,generator,discriminator,EPOCHS,fixed_noise)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ypSz-8HzIQJ0"
      },
      "source": [
        "## Sauvegarde du modèle et inférence\n",
        "\n",
        "Maintenant que votre modèle s'est entraîné pendant 10 minutes voir plus, on ne veut pas le réentraîner à chaque fois.\n",
        "On va donc voir comment le sauvegarder sur votre drive et ensuite l'inférer, c'est-à-dire générer les images sans toucher aux poids."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SnoImcTtIQJ1",
        "outputId": "7360c092-4045-4c9f-aedb-796c3f177c36"
      },
      "outputs": [],
      "source": [
        "# Mount son drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "# Google va vous demander alors de vous logger avec votre compte google et accepter les conditions blabla..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ECsgo2zaIQJ1"
      },
      "source": [
        "On peut sauvegarder uniquement les poids du modèle à chaque fois pour ne pas avoir à stocker d'autres informations inutiles."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wsFfGkajIQJ1"
      },
      "outputs": [],
      "source": [
        "# On save les poids du générateur dans un fichier .h5 avec la méthode .save_weights\n",
        "torch.save(generator,'/content/drive/My Drive/generator.h5')\n",
        "\n",
        "# Faire pareil pour le discriminateur\n",
        "torch.save(discriminator,'/content/drive/My Drive/discriminator.h5')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XwrekpydIQJ2"
      },
      "source": [
        "Ensuite il ne reste plus qu'à load les poids des deux modèles, générer des vecteurs latents et inférer le modèle dessus.\n",
        "Pour inférer le modèle, on a déjà fait ça dans la boucle d'entraînement, il suffit de le considérer comme une simple fonction et faire ```model(entree)```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 961
        },
        "id": "sUrBuwOIIQJ2",
        "outputId": "f61c72b3-e54b-46a9-8283-7e997965220c"
      },
      "outputs": [],
      "source": [
        "# On peut alors charger les poids d'un modèle déjà entraîné avec .load_weights\n",
        "generator = torch.load('/content/drive/My Drive/generator.h5')\n",
        "\n",
        "# Faire pareil pour le discriminateur\n",
        "discriminator = torch.load('/content/drive/My Drive/discriminator.h5')\n",
        "\n",
        "# On peut alors générer des images avec le générateur\n",
        "latent_vector = torch.randn(size=(25, LATENT_DIM)).to(device) #Création de 25 vecteurs latents aléatoires pour générer 25 images\n",
        "fake_images = generator(latent_vector).detach().cpu()               #Génération des images avec le générateur\n",
        "fake_images = torch.permute(fake_images, (0, 2, 3, 1))              #On déplace les dimensions de notre tenseur pour les adapter à un format compatible avec plt\n",
        "\n",
        "fig = plt.figure(figsize=(12,12))\n",
        "for i in range(25):\n",
        "    plt.subplot(5,5,i+1)\n",
        "    plt.axis('off')\n",
        "    plt.imshow(fake_images[i]*0.5+0.5)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "prn-zEwC7jhy"
      },
      "source": [
        "# À développer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oq_5UD18IQJ3"
      },
      "source": [
        "## Autres datasets\n",
        "\n",
        "Vous pouvez essayer de générer des images sur d'autres datasets comme Fashion MNIST, CIFAR, CelebA ou n'importe quel type d'images qui vous font plaisir. S'il faut le télécharger et l'importer sur Colab, vous pouvez directement monter votre Drive et uploader votre dataset sur ce dernier.\n",
        "\n",
        "## Autres architectures\n",
        "\n",
        "Vous pouvez rajouter des blocs dans le générateur ou le discriminateur, essayer d'ajouter du Dropout, enlever les biais dans les couches, modifier la dimension de l'espace latent, initialiser les poids d'une certaine façon... Si vous êtes déjà à jour sur les CNN avancés (ResNET, MobileNet, EfficientNet), vous pouvez essayer de faire un GAN sur ces bases, ce qui permettra de résoudre certains problèmes comme les vanishing gradient.\n",
        "\n",
        "## Tips d'entraînement possibles\n",
        "\n",
        "- Utiliser une loss différente, du genre la Wasserstein loss.\n",
        "\n",
        "- Changer le label visé pour les images vraies de 1 en 0.9\n",
        "\n",
        "- Rajouter du bruit sur les images\n",
        "\n",
        "- Ajouter des labels à l'entrée du générateur et du discriminateur (vous pourrez alors même choisir les classes des images générées)\n",
        "\n",
        "- Entraîner le discrminateur plus que le générateur (typiquement entre 3 à 5 boucles d'entraînement à chaque fois que le générateur en fait une).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pM6h4yvial7i"
      },
      "source": [
        "# Ce qu'il faut retenir de ce TP\n",
        "- Architecture du générateur et du discriminateur :\n",
        "\n",
        "  `Conv2D(Transpose) -> BatchNorm -> ReLU ou LeakyReLU`\n",
        "\n",
        "  + Spécifités en entrée et sortie selon ce que l'on veut\n",
        "\n",
        "  Eviter les fully connected/denses si possible"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.5"
    },
    "vscode": {
      "interpreter": {
        "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}