{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9cfe3442",
   "metadata": {},
   "source": [
    "# TP – GANs avancés : Conditional GAN, Progressive Growing & Rejection Sampling\n",
    "\n",
    "Dans ce notebook, nous allons implémenter et explorer plusieurs variantes de GANs sur **FashionMNIST** :\n",
    "\n",
    "1. Un **GAN convolutionnel standard** sur des images 32×32.  \n",
    "2. Un **Conditional GAN (cGAN)** qui génère par classe à l’aide d’**embeddings**.  \n",
    "3. Un **Progressive Growing GAN (PG-GAN)** jouet (16×16 → 32×32).  \n",
    "4. Une **évaluation détaillée des discriminateurs** (accuracies, par classe).  \n",
    "5. Une **calibration** simple du discriminateur (décalage de logit).  \n",
    "6. Du **Discriminator Rejection Sampling (DRS)** pour filtrer / repondérer les échantillons.\n",
    "\n",
    "Nous allons alterner :\n",
    "- des cellules **Markdown** expliquant les concepts,  \n",
    "- des cellules de **code** où vous devrez compléter des parties marquées par :\n",
    "\n",
    "<font color='blue'>**TODO:**</font>\n",
    "\n",
    "et des commentaires `# TODO:` dans le code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "295a96b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms, utils\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# ------------------------------------------------------------------\n",
    "# Device (CPU / GPU)\n",
    "# ------------------------------------------------------------------\n",
    "if torch.cuda.is_available():\n",
    "    print(\"CUDA is available. Using GPU.\")\n",
    "    device = torch.device(\"cuda\")\n",
    "elif torch.backends.mps.is_available():\n",
    "    print(\"MPS is available. Using Apple Silicon GPU.\")\n",
    "    device = torch.device(\"mps\")\n",
    "else:\n",
    "    print(\"Using CPU.\")\n",
    "    device = torch.device(\"cpu\")\n",
    "\n",
    "print(\"Device:\", device)\n",
    "\n",
    "# ------------------------------------------------------------------\n",
    "# Hyperparameters (utilisés partout dans le notebook)\n",
    "# ------------------------------------------------------------------\n",
    "batch_size = 128\n",
    "z_dim = 128  # dimension du bruit (latent space)\n",
    "image_size = 32  # taille des images (32x32)\n",
    "nc = 1  # nombre de canaux (FashionMNIST = niveaux de gris)\n",
    "num_classes = 10  # FashionMNIST: 10 classes\n",
    "\n",
    "# ------------------------------------------------------------------\n",
    "# Transforms & Datasets (train / test)\n",
    "# ------------------------------------------------------------------\n",
    "transform_32 = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(image_size),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5,), (0.5,)),  # ramener les pixels dans [-1,1]\n",
    "    ]\n",
    ")\n",
    "\n",
    "train_dataset = datasets.FashionMNIST(\n",
    "    root=\"./data\", train=True, download=True, transform=transform_32\n",
    ")\n",
    "test_dataset = datasets.FashionMNIST(\n",
    "    root=\"./data\", train=False, download=True, transform=transform_32\n",
    ")\n",
    "\n",
    "train_loader = DataLoader(\n",
    "    train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True\n",
    ")\n",
    "test_loader = DataLoader(\n",
    "    test_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51219abd",
   "metadata": {},
   "source": [
    "## 1. Unconditional Conv-GAN sur FashionMNIST (32×32)\n",
    "\n",
    "Nous commençons par implémenter un **DCGAN** standard :\n",
    "\n",
    "- Un **générateur convolutionnel** \\( G(z) \\) qui prend un bruit `z` de dimension `z_dim` et produit une image 32×32.  \n",
    "- Un **discriminateur convolutionnel** \\( D(x) \\) qui prend une image 32×32 et renvoie une probabilité que l’image soit réelle.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Complétes les classes ConvGenerator et ConvDiscriminator pour y intégrer les activations finales appropriées (Tanh pour G, Sigmoid pour D)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f0c1e47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------\n",
    "# 1) Conv GAN standard\n",
    "# -------------------------------\n",
    "class ConvGenerator(nn.Module):\n",
    "    def __init__(self, z_dim=128, ngf=64, nc=1):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            # input Z: (z_dim) -> (ngf*4) x 4 x 4\n",
    "            nn.ConvTranspose2d(z_dim, ngf * 4, 4, 1, 0, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 4),\n",
    "            nn.ReLU(True),\n",
    "            # (ngf*4) x 4 x 4 -> (ngf*2) x 8 x 8\n",
    "            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 2),\n",
    "            nn.ReLU(True),\n",
    "            # (ngf*2) x 8 x 8 -> (ngf) x 16 x 16\n",
    "            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf),\n",
    "            nn.ReLU(True),\n",
    "            # (ngf) x 16 x 16 -> (nc) x 32 x 32\n",
    "            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),\n",
    "        )\n",
    "        self.activation = # TODO\n",
    "\n",
    "    def forward(self, z):\n",
    "        # z: (B, z_dim)\n",
    "        z = z.view(z.size(0), z.size(1), 1, 1)\n",
    "        z = self.net(z)\n",
    "         # TODO\n",
    "        return z\n",
    "\n",
    "\n",
    "class ConvDiscriminator(nn.Module):\n",
    "    def __init__(self, nc=1, ndf=64):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            # input image: (nc) x 32 x 32\n",
    "            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),  # 16x16\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),  # 8x8\n",
    "            nn.BatchNorm2d(ndf * 2),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),  # 4x4\n",
    "            nn.BatchNorm2d(ndf * 4),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),  # 1x1\n",
    "        )\n",
    "        self.activation =  # TODO\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.net(x)\n",
    "        # TODO\n",
    "        return out.view(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e318f696",
   "metadata": {},
   "source": [
    "### 1.1. Boucle d'entraînement du GAN\n",
    "\n",
    "Nous allons maintenant écrire une fonction `train_gan` qui :\n",
    "\n",
    "- prend en entrée un générateur `G`, un discriminateur `D` et un `dataloader`,  \n",
    "- effectue plusieurs epochs d'entraînement,  \n",
    "- affiche quelques échantillons générés à la fin de chaque epoch.\n",
    "\n",
    "<font color='blue'>**TODO:**</font>  Identifiez dans le code ci-dessous l'étape de mise à jour de `D` (avec réelles + fausses), et complétez mise à jour de `G`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af880961",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_gan(\n",
    "    generator: nn.Module,\n",
    "    discriminator: nn.Module,\n",
    "    dataloader: DataLoader,\n",
    "    num_epochs: int = 10,\n",
    "    lr: float = 2e-4,\n",
    "    beta1: float = 0.5,\n",
    "    z_dim: int = 128,\n",
    "    name: str = \"DCGAN_32x32\",\n",
    "):\n",
    "    \"\"\"\n",
    "    Train an unconditional GAN (DCGAN-style) on FashionMNIST.\n",
    "\n",
    "    Args:\n",
    "        generator: ConvGenerator instance G.\n",
    "        discriminator: ConvDiscriminator instance D.\n",
    "        dataloader: DataLoader over training images.\n",
    "        num_epochs: number of training epochs.\n",
    "        lr: learning rate for Adam optimizers.\n",
    "        beta1: beta1 parameter for Adam.\n",
    "        z_dim: latent dimension.\n",
    "        name: string used for plots / saving images.\n",
    "    \"\"\"\n",
    "    G = generator.to(device)\n",
    "    D = discriminator.to(device)\n",
    "\n",
    "    criterion = nn.BCELoss()\n",
    "    optim_D = optim.Adam(D.parameters(), lr=lr, betas=(beta1, 0.999))\n",
    "    optim_G = optim.Adam(G.parameters(), lr=lr, betas=(beta1, 0.999))\n",
    "\n",
    "    G.train()\n",
    "    D.train()\n",
    "\n",
    "    fixed_noise = torch.randn(64, z_dim, device=device)\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        for x_real, _ in dataloader:\n",
    "            x_real = x_real.to(device)\n",
    "\n",
    "            real_labels = torch.ones(x_real.size(0), device=device)\n",
    "            fake_labels = torch.zeros(x_real.size(0), device=device)\n",
    "\n",
    "            # --------- Train D ---------\n",
    "            # 1) Calculer D(x_real) et la loss sur les vraies images.\n",
    "            # 2) Générer un batch de fausses images x_fake = G(z).\n",
    "            # 3) Calculer D(x_fake) et la loss sur les fausses images.\n",
    "            # 4) Additionner les deux pour obtenir loss_D, puis backward + step.\n",
    "\n",
    "            D.zero_grad()\n",
    "            out_real = D(x_real)\n",
    "            loss_D_real = criterion(out_real, real_labels)\n",
    "\n",
    "            z = torch.randn(x_real.size(0), z_dim, device=device)\n",
    "            x_fake = G(z).detach()\n",
    "            out_fake = D(x_fake)\n",
    "            loss_D_fake = criterion(out_fake, fake_labels)\n",
    "\n",
    "            loss_D = loss_D_real + loss_D_fake\n",
    "            loss_D.backward()\n",
    "            optim_D.step()\n",
    "\n",
    "            # --------- Train G ---------\n",
    "            # <font color='blue'>**TODO:**</font> Compléter la partie suivante :\n",
    "            # 1) Regénérer un batch de z.\n",
    "            # 2) Générer x_fake = G(z).\n",
    "            # 3) Calculer D(x_fake) (on veut que D(x_fake) ≈ 1).\n",
    "            # 4) Calculer loss_G et mettre à jour G.\n",
    "\n",
    "            # TODO\n",
    "\n",
    "            # TODO\n",
    "\n",
    "        # --------- Visualisation à la fin de l'epoch ---------\n",
    "        with torch.no_grad():\n",
    "            fake_images = G(fixed_noise).cpu()\n",
    "        grid = utils.make_grid(fake_images, nrow=8, normalize=True, value_range=(-1, 1))\n",
    "        plt.figure(figsize=(6, 6))\n",
    "        plt.imshow(grid.permute(1, 2, 0))\n",
    "        plt.axis(\"off\")\n",
    "        plt.title(f\"{name.replace('_', ' ')} - Epoch {epoch+1}\")\n",
    "        if epoch == num_epochs - 1:\n",
    "            plt.savefig(f\"{name}_samples_epoch{epoch+1}.png\")\n",
    "        plt.show()\n",
    "\n",
    "        print(\n",
    "            f\"[Epoch {epoch+1}/{num_epochs}] \"\n",
    "            f\"Loss_D: {loss_D.item():.4f}  Loss_G: {loss_G.item():.4f}\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "718e31d1",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Vérifiez que votre GAN fonctionne correctement en l'entraînant sur FashionMNIST sur 5 epochs et en affichant les échantillons générés à chaque epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2af26e6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Entraînement du GAN convolutionnel (unconditional) sur FashionMNIST 32x32\n",
    "\n",
    "G = ConvGenerator(z_dim=z_dim, nc=nc)\n",
    "D = ConvDiscriminator(nc=nc)\n",
    "train_gan(G, D, train_loader, num_epochs=5, z_dim=z_dim, name=\"DCGAN_32x32\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a0c876d",
   "metadata": {},
   "source": [
    "## 2. Conditional GAN (cGAN) avec embeddings de classe\n",
    "\n",
    "Nous allons maintenant rendre le GAN **conditionnel** sur la classe `y` de FashionMNIST.\n",
    "\n",
    "Idée :\n",
    "- Apprendre un **embedding de classe** `e(y)` (vector de dimension `emb_dim`).\n",
    "- Générateur : `G(z, y)` utilise à la fois le bruit `z` et l'embedding `e(y)`.\n",
    "- Discriminateur : `D(x, y)` reçoit l'image **et** une carte dérivée de `e(y)`.\n",
    "\n",
    "Le générateur est généralement composé d'un bloc linéaire suivi de couches convolutionnelles transposées (ConvTranspose2d) pour augmenter la résolution. On concatène l'embedding de classe au vecteur de bruit avant de le passer dans le bloc linéaire.\n",
    "\n",
    "Pour le discriminateur, on peut créer une carte de caractéristiques à partir de l'embedding de classe et la concaténer aux canaux de l'image d'entrée. \n",
    "\n",
    "<font color='blue'>**TODO:**</font> A partir des fonctions forward ci-dessous, complétez les parties manquantes des fonctions `__init__` des classes `CondConvGenerator` et `CondConvDiscriminator` afin d'intégrer les embeddings de classe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22c8681a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------\n",
    "# 2) Conditional Conv GAN\n",
    "# -------------------------------\n",
    "\n",
    "\n",
    "class CondConvGenerator(nn.Module):\n",
    "    \"\"\"\n",
    "    Conditional convolutional generator G(z, y).\n",
    "\n",
    "    - y is a class in {0,...,9}\n",
    "    - y is first embedded into a dense vector, then concatenated to z.\n",
    "    - The rest of the architecture is similar to the unconditional ConvGenerator.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, z_dim=128, num_classes=10, ngf=64, nc=1, emb_dim=50):\n",
    "        super().__init__()\n",
    "        self.label_emb = nn.Embedding(num_classes, emb_dim)\n",
    "\n",
    "        self.fc =  # TODO\n",
    "\n",
    "        in_channels = z_dim + emb_dim\n",
    "        self.net = nn.Sequential(\n",
    "            nn.ConvTranspose2d(in_channels, ngf * 4, 4, 1, 0, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 4),\n",
    "            nn.ReLU(True),\n",
    "            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 2),\n",
    "            nn.ReLU(True),\n",
    "            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf),\n",
    "            nn.ReLU(True),\n",
    "            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "\n",
    "    def forward(self, z, y):\n",
    "        # z: (B, z_dim), y: (B,)\n",
    "        y_emb = self.label_emb(y)  # (B, emb_dim)\n",
    "        h = torch.cat([z, y_emb], dim=1)  # (B, z_dim+emb_dim)\n",
    "        h = self.fc(h)\n",
    "        h = h.view(h.size(0), h.size(1), 1, 1)\n",
    "        x_fake = self.net(h)\n",
    "        return x_fake\n",
    "\n",
    "\n",
    "class CondConvDiscriminator(nn.Module):\n",
    "    \"\"\"\n",
    "    Conditional convolutional discriminator D(x, y).\n",
    "\n",
    "    - y is embedded, then expanded into a 32x32 map.\n",
    "    - This map is concatenated as an extra channel to x.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, num_classes=10, nc=1, ndf=64, emb_dim=50):\n",
    "        super().__init__()\n",
    "        self.label_emb = nn.Embedding(num_classes, emb_dim)\n",
    "\n",
    "        # On transforme l'embedding en carte 32x32 et on concatène comme canal\n",
    "        self.fc_label =  # TODO\n",
    "\n",
    "        in_channels = nc + 1  # 1 canal image + 1 canal label\n",
    "\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, ndf, 4, 2, 1, bias=False),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ndf * 2),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ndf * 4),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        # x: (B,1,32,32), y: (B,)\n",
    "        y_emb = self.label_emb(y)  # (B, emb_dim)\n",
    "        y_map = self.fc_label(y_emb)  # (B, 32*32)\n",
    "        y_map = y_map.view(x.size(0), 1, 32, 32)  # (B,1,32,32)\n",
    "        x_cat = torch.cat([x, y_map], dim=1)  # (B,2,32,32)\n",
    "        out = self.net(x_cat)\n",
    "        return out.view(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bdc4cef",
   "metadata": {},
   "source": [
    "### 2.2. Boucle d'entraînement du Conditional GAN\n",
    "\n",
    "Nous allons maintenant entraîner le **cGAN** défini ci-dessus :\n",
    "\n",
    "- Générateur : `G(z, y)`, conditionné par la classe `y`.  \n",
    "- Discriminateur : `D(x, y)`, qui doit renvoyer une probabilité que `x` soit **réelle**.\n",
    "\n",
    "Sur chaque batch, nous allons :\n",
    "\n",
    "1. **Mettre à jour le discriminateur** `D` :\n",
    "   - Calculer la loss sur les **images réelles** `x_{\\text{real}}` avec label cible 1.\n",
    "   - Générer des **images fausses** `x_{\\text{fake}} = G(z, y)` avec label cible 0.\n",
    "   - Additionner les deux pertes et faire un `backward()` + `step()` sur `D`.\n",
    "\n",
    "2. **Mettre à jour le générateur** `G` :\n",
    "   - Générer un nouveau batch `x_{\\text{fake}} = G(z, y)`.\n",
    "   - Évaluer `D(x_{\\text{fake}}, y)` et **tromper** le discriminateur (labels cibles = 1).\n",
    "   - Calculer la loss de `G` et faire un `backward()` + `step()` sur `G`.\n",
    "\n",
    "Nous garderons aussi un **bruit fixe** `fixed_z` et des labels `fixed_labels` pour visualiser,\n",
    "à la fin de chaque epoch, **10×10 images générées (une grille par classe)**.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Dans la fonction `train_cgan` ci-dessous :\n",
    "- Identifiez les blocs où l’on met à jour le **discriminateur** et le **générateur**.  \n",
    "- Reprenez la logique de `train_gan` (GAN non conditionnel) en l’adaptant au cas **conditionnel**.  \n",
    "- Les lignes marquées par `# TODO` sont celles que vous devrez re-coder vous-mêmes dans la version étudiante."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a347ced3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_cgan(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    dataloader,\n",
    "    num_epochs: int = 10,\n",
    "    lr: float = 2e-4,\n",
    "    beta1: float = 0.5,\n",
    "    z_dim: int = 128,\n",
    "    name: str = \"cDCGAN_32x32\",\n",
    "):\n",
    "    \"\"\"\n",
    "    Train a conditional GAN (cGAN) on FashionMNIST.\n",
    "\n",
    "    - The generator G takes as input (z, y) and outputs x_fake = G(z, y).\n",
    "    - The discriminator D takes as input (x, y) and outputs p_real = D(x, y).\n",
    "\n",
    "    Training loop (per batch):\n",
    "        1) Update D:\n",
    "            - maximize log D(x_real, y) + log(1 - D(x_fake, y))\n",
    "        2) Update G:\n",
    "            - maximize log D(x_fake, y)\n",
    "\n",
    "    Args:\n",
    "        generator: CondConvGenerator (G).\n",
    "        discriminator: CondConvDiscriminator (D).\n",
    "        dataloader: DataLoader over (x, y) pairs.\n",
    "        num_epochs: number of epochs.\n",
    "        lr: learning rate for Adam.\n",
    "        beta1: beta1 for Adam.\n",
    "        z_dim: latent dimension.\n",
    "        name: string used in plot titles / filenames.\n",
    "    \"\"\"\n",
    "    G = generator.to(device)\n",
    "    D = discriminator.to(device)\n",
    "\n",
    "    criterion = nn.BCELoss()\n",
    "    optim_D = optim.Adam(D.parameters(), lr=lr, betas=(beta1, 0.999))\n",
    "    optim_G = optim.Adam(G.parameters(), lr=lr, betas=(beta1, 0.999))\n",
    "\n",
    "    G.train()\n",
    "    D.train()\n",
    "\n",
    "    # Pour visualiser par classe : 10 classes × 10 samples = 100 images\n",
    "    fixed_z = torch.randn(100, z_dim, device=device)\n",
    "    fixed_labels = torch.arange(0, 10, device=device).repeat_interleave(10)\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        for i, (x, y) in enumerate(dataloader):\n",
    "            x, y = x.to(device), y.to(device)\n",
    "\n",
    "        # TODO\n",
    "\n",
    "        # TODO\n",
    "\n",
    "        # --------- Visualisation par classe à la fin de l'epoch ---------\n",
    "        with torch.no_grad():\n",
    "            fake_imgs = G(fixed_z, fixed_labels).cpu()\n",
    "        grid = utils.make_grid(fake_imgs, nrow=10, normalize=True, value_range=(-1, 1))\n",
    "        plt.figure(figsize=(6, 6))\n",
    "        plt.imshow(grid.permute(1, 2, 0))\n",
    "        plt.axis(\"off\")\n",
    "        plt.title(f\"{name.replace('_', ' ')} - Epoch {epoch+1}\")\n",
    "        if epoch == num_epochs - 1:\n",
    "            plt.savefig(f\"{name}_final.png\")\n",
    "        plt.show()\n",
    "\n",
    "        print(\n",
    "            f\"[Epoch {epoch+1}/{num_epochs}] \"\n",
    "            f\"Loss_D: {loss_D.item():.4f}  Loss_G: {loss_G.item():.4f}\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c52dc53",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Entraînez votre cGAN sur FashionMNIST pendant 10 epochs et affichez les échantillons générés à chaque epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5bead20",
   "metadata": {},
   "outputs": [],
   "source": [
    "Gc = CondConvGenerator(z_dim=z_dim, num_classes=num_classes, nc=nc)\n",
    "Dc = CondConvDiscriminator(num_classes=num_classes, nc=nc)\n",
    "train_cgan(Gc, Dc, train_loader, num_epochs=5, z_dim=z_dim, name=\"cDCGAN_32x32\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5709c0a1",
   "metadata": {},
   "source": [
    "## 3. Progressive Growing : de 16×16 à 32×32\n",
    "\n",
    "Nous allons maintenant implémenter un **mini Progressive Growing GAN** (PG-GAN) :\n",
    "\n",
    "1. On commence par entraîner un GAN standard en **16×16** (mêmes idées que le DCGAN 32×32).  \n",
    "2. Puis on ajoute un **bloc supplémentaire** pour passer à 32×32 et on fait un **fade-in** progressif entre :\n",
    "   - l’ancien chemin (image 16×16 upsamplée),\n",
    "   - le nouveau chemin (image générée en 32×32).\n",
    "\n",
    "L’objectif **pédagogique** n’est pas de battre les SOTA, mais de :\n",
    "- voir que le modèle 16×16 est très proche de ce que nous avons déjà fait,  \n",
    "- comprendre comment le paramètre  $\\alpha$ permet de **mélanger progressivement** deux résolutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eee78a8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform_16 = transforms.Compose(\n",
    "    [transforms.Resize(16), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]\n",
    ")\n",
    "\n",
    "train_dataset_16 = datasets.FashionMNIST(\n",
    "    root=\"./data\", train=True, download=True, transform=transform_16\n",
    ")\n",
    "\n",
    "dataloader_16 = DataLoader(\n",
    "    train_dataset_16,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce48074f",
   "metadata": {},
   "source": [
    "### 3.1. Générateur et discriminateur en 16×16\n",
    "\n",
    "Nous définissons maintenant :\n",
    "\n",
    "- un **générateur 16×16** `Gen16`, très proche du `ConvGenerator` 32×32 mais qui s’arrête à la résolution 16×16 ;  \n",
    "- un **discriminateur 16×16** `Disc16`, très proche du `ConvDiscriminator` 32×32 mais adapté aux tailles 16×16.\n",
    "\n",
    "La différence principale est que le module `Gen16` propose au choix une méthode avec ou sans **tanh** en sortie (pour faciliter le fade-in plus tard).\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Observez les classes `Gen16` et `Disc16` ci-dessous et lancez leur entraênement sur FashionMNIST en 16×16.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a82086e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Gen16(nn.Module):\n",
    "    \"\"\"\n",
    "    Generator producing 16x16 images (unconditional).\n",
    "\n",
    "    This is structurally very close to the ConvGenerator(32x32),\n",
    "    but we stop one upsampling earlier so the output resolution is 16x16.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, z_dim=128, ngf=64, nc=1):\n",
    "        super().__init__()\n",
    "        # On part d'un vecteur latent et on fabrique un feature map 4x4\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(z_dim, ngf * 4 * 4 * 4),  # (ngf*4, 4, 4)\n",
    "            nn.BatchNorm1d(ngf * 4 * 4 * 4),\n",
    "            nn.ReLU(True),\n",
    "        )\n",
    "\n",
    "        # Blocs upsample + conv : 4x4 -> 8x8 -> 16x16\n",
    "        self.block1 = nn.Sequential(\n",
    "            nn.Upsample(scale_factor=2, mode=\"nearest\"),  # 4x4 -> 8x8\n",
    "            nn.Conv2d(ngf * 4, ngf * 2, 3, 1, 1),\n",
    "            nn.BatchNorm2d(ngf * 2),\n",
    "            nn.ReLU(True),\n",
    "        )\n",
    "\n",
    "        self.block2 = nn.Sequential(\n",
    "            nn.Upsample(scale_factor=2, mode=\"nearest\"),  # 8x8 -> 16x16\n",
    "            nn.Conv2d(ngf * 2, ngf, 3, 1, 1),\n",
    "            nn.BatchNorm2d(ngf),\n",
    "            nn.ReLU(True),\n",
    "        )\n",
    "\n",
    "        # Tête image 16x16\n",
    "        self.head16 = nn.Conv2d(ngf, nc, 3, 1, 1)\n",
    "\n",
    "    def forward(self, z, output_features=False):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            z: tensor of shape (B, z_dim)\n",
    "            output_features: if True, also return the feature map h16\n",
    "                             (used later by the 32x32 fade-in generator).\n",
    "\n",
    "        Returns:\n",
    "            x16: (B, 1, 16, 16) if output_features=False\n",
    "            (x16, h16) if output_features=True\n",
    "        \"\"\"\n",
    "        # z : (B, z_dim)\n",
    "        h = self.fc(z)  # (B, ngf*4*4*4)\n",
    "        h = h.view(z.size(0), -1, 4, 4)  # (B, ngf*4, 4, 4)\n",
    "\n",
    "        h = self.block1(h)  # (B, ngf*2, 8, 8)\n",
    "        h16 = self.block2(h)  # (B, ngf, 16, 16)\n",
    "\n",
    "        x16 = torch.tanh(self.head16(h16))  # (B, 1, 16, 16)\n",
    "\n",
    "        if output_features:\n",
    "            return x16, h16\n",
    "        else:\n",
    "            return x16\n",
    "\n",
    "\n",
    "class Disc16(nn.Module):\n",
    "    \"\"\"\n",
    "    16x16 discriminator (unconditional).\n",
    "\n",
    "    Very similar to ConvDiscriminator, but the input resolution is 16x16.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, nc=1, ndf=64):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),  # 8x8\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),  # 4x4\n",
    "            nn.BatchNorm2d(ndf * 2),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf * 2, 1, 4, 1, 0, bias=False),  # 1x1\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.net(x)\n",
    "        return out.view(-1)\n",
    "\n",
    "\n",
    "class Disc32_FadeIn(nn.Module):\n",
    "    \"\"\"\n",
    "    32x32 discriminator with fade-in on top of a 16x16 discriminator.\n",
    "\n",
    "    - For alpha = 0.0, it mostly uses the old downsampled path.\n",
    "    - For alpha = 1.0, it mostly uses the new 32->16 conv block.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, disc16: Disc16, nc=1, ndf=64):\n",
    "        super().__init__()\n",
    "        self.disc16 = disc16\n",
    "        # new block: 32x32 -> 16x16\n",
    "        self.new_block = nn.Sequential(\n",
    "            nn.Conv2d(nc, ndf, 3, 1, 1),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf, nc, 3, 2, 1),  # 32x32 -> 16x16\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "        )\n",
    "        self.downsample_old = nn.AvgPool2d(2)\n",
    "\n",
    "    def forward(self, x, alpha: float = 1.0):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: 32x32 image (B,1,32,32)\n",
    "            alpha: fade-in coefficient in [0,1]\n",
    "\n",
    "        Returns:\n",
    "            p_real: probability x is real (B,)\n",
    "        \"\"\"\n",
    "        # ancien chemin : simple downsampling\n",
    "        x_old = self.downsample_old(x)  # (B,1,16,16)\n",
    "        # nouveau chemin : bloc conv 32->16\n",
    "        x_new = self.new_block(x)  # (B,1,16,16)\n",
    "\n",
    "        # TODO: mélanger x_old et x_new avec la même formule qu'au niveau du générateur\n",
    "        x_mix = (1 - alpha) * x_old + alpha * x_new\n",
    "\n",
    "        out = self.disc16(x_mix)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92d394fa",
   "metadata": {},
   "source": [
    "### 3.2. Entraîner le GAN 16×16\n",
    "\n",
    "Avant de passer à la phase 32×32, nous devons d’abord entraîner un GAN sur des images 16×16 :\n",
    "\n",
    "- Générateur : `Gen16`\n",
    "- Discriminateur : `Disc16`\n",
    "\n",
    "Nous allons réutiliser **exactement la même boucle d’entraînement** que pour le DCGAN 32×32 (`train_gan`), simplement en changeant :\n",
    "- le modèle (`Gen16`, `Disc16`),\n",
    "- le dataloader (`dataloader_16`).\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Comparez les résultats visuels des 16×16 aux 32×32 :\n",
    "- Les vêtements sont-ils reconnaissables ?\n",
    "- Les artefacts sont-ils plus ou moins visibles ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e542800d",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen16 = Gen16(z_dim=z_dim, nc=nc).to(device)\n",
    "disc16 = Disc16(nc=nc).to(device)\n",
    "\n",
    "train_gan(gen16, disc16, dataloader_16, num_epochs=5, z_dim=z_dim, name=\"DCGAN_16x16\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9401ed4a",
   "metadata": {},
   "source": [
    "### 3.3. Générateur 32×32 avec fade-in\n",
    "\n",
    "Nous construisons maintenant un générateur 32×32 à partir de `Gen16` :\n",
    "\n",
    "- `Gen16` produit une image 16×16 et des features 16×16.\n",
    "- Nous :\n",
    "  1. **Upsamplons** l’ancienne image 16×16 en 32×32 (`x16_up`),\n",
    "  2. Appliquons un **nouveau bloc** sur les features pour produire une image 32×32 (`x32_new`),\n",
    "  3. Faisons un **blend progressif** entre les deux :\n",
    "\n",
    "$$\n",
    "x_{\\text{out}} = (1 - \\alpha) \\cdot x_{16\\uparrow} + \\alpha \\cdot x_{32,\\text{new}},\n",
    "$$\n",
    "\n",
    "avec $ \\alpha \\in [0,1] $ qui augmente au cours de l’entraînement.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Compléter le `forward` du générateur 32×32 en utilisant cette formule de `fade-in`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "090f76cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Gen32_FadeIn(nn.Module):\n",
    "    \"\"\"\n",
    "    32x32 generator built on top of a 16x16 generator (Gen16) with fade-in.\n",
    "\n",
    "    - For alpha = 0.0, we mostly use the upsampled 16x16 image.\n",
    "    - For alpha = 1.0, we mostly use the new 32x32 head.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, gen16: Gen16, nc=1, ngf=64):\n",
    "        super().__init__()\n",
    "        self.gen16 = gen16\n",
    "        self.upsample = nn.Upsample(scale_factor=2, mode=\"nearest\")\n",
    "\n",
    "        # new 32x32 \"head\" on top of 16x16 features\n",
    "        self.new_head = nn.Sequential(\n",
    "            nn.Upsample(scale_factor=2, mode=\"nearest\"),  # 16x16 -> 32x32 (features)\n",
    "            nn.Conv2d(ngf, ngf, 3, 1, 1),\n",
    "            nn.BatchNorm2d(ngf),\n",
    "            nn.ReLU(True),\n",
    "            nn.Conv2d(ngf, nc, 3, 1, 1),\n",
    "        )\n",
    "\n",
    "    def forward(self, z, alpha: float = 1.0):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            z: latent vector (B, z_dim)\n",
    "            alpha: fade-in coefficient in [0,1]\n",
    "\n",
    "        Returns:\n",
    "            x_out: blended 32x32 image (B, 1, 32, 32)\n",
    "        \"\"\"\n",
    "        # 1) Get original 16x16 output and features\n",
    "        x16, h16 = self.gen16(z, output_features=True)  # ancien output + features\n",
    "        x16_up = self.upsample(x16)  # (B,1,32,32)\n",
    "\n",
    "        # 2) New 32x32 branch\n",
    "        x32_new = self.new_head(h16)  # (B,1,32,32), not yet tanh\n",
    "        x32_new = torch.tanh(x32_new)\n",
    "\n",
    "        # 3) Fade-in between old path (upsampled 16x16) and new 32x32 branch\n",
    "        x_out =  # TODO\n",
    "\n",
    "        return x_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca6d632",
   "metadata": {},
   "source": [
    "### 3.4. Discriminateur 16×16 et discriminateur 32×32 avec fade-in\n",
    "\n",
    "Côté discriminateur, l’idée est symétrique :\n",
    "\n",
    "- `Disc16` discrimine sur des images 16×16.  \n",
    "- `Disc32_FadeIn` :\n",
    "  - possède un **nouveau bloc 32→16** qui ramène les images 32×32 en features 16×16,\n",
    "  - mélange l’**ancien chemin downsamplé** (`x_old`) et le **nouveau chemin** (`x_new`) :\n",
    "    $$\n",
    "    x_{\\text{mix}} = (1 - \\alpha) \\cdot x_{\\text{old}} + \\alpha \\cdot x_{\\text{new}}.\n",
    "    $$\n",
    "  - donne ce `x_mix` au `Disc16`.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Compléter le `forward` de `Disc32_FadeIn` en utilisant le même schéma de fade-in que pour le générateur."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adf8e55d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Disc32_FadeIn(nn.Module):\n",
    "    \"\"\"\n",
    "    32x32 discriminator with fade-in on top of a 16x16 discriminator.\n",
    "\n",
    "    - For alpha = 0.0, it mostly uses the old downsampled path.\n",
    "    - For alpha = 1.0, it mostly uses the new 32->16 conv block.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, disc16: Disc16, nc=1, ndf=64):\n",
    "        super().__init__()\n",
    "        self.disc16 = disc16\n",
    "        # new block: 32x32 -> 16x16\n",
    "        self.new_block = nn.Sequential(\n",
    "            nn.Conv2d(nc, ndf, 3, 1, 1),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(ndf, nc, 3, 2, 1),  # 32x32 -> 16x16\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "        )\n",
    "        self.downsample_old = nn.AvgPool2d(2)\n",
    "\n",
    "    def forward(self, x, alpha: float = 1.0):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: 32x32 image (B,1,32,32)\n",
    "            alpha: fade-in coefficient in [0,1]\n",
    "\n",
    "        Returns:\n",
    "            p_real: probability x is real (B,)\n",
    "        \"\"\"\n",
    "        # ancien chemin : simple downsampling\n",
    "        x_old = self.downsample_old(x)  # (B,1,16,16)\n",
    "        # nouveau chemin : bloc conv 32->16\n",
    "        x_new = self.new_block(x)  # (B,1,16,16)\n",
    "\n",
    "        #  mélanger x_old et x_new avec la même formule qu'au niveau du générateur\n",
    "        x_mix =  # TODO\n",
    "\n",
    "        out = self.disc16(x_mix)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a944844",
   "metadata": {},
   "source": [
    "### 3.5. Boucle d'entraînement progressive (16×16 → 32×32)\n",
    "\n",
    "Nous allons maintenant écrire une boucle `train_progressive` qui :\n",
    "\n",
    "1. Prenant en entrée un **Gen16** et un **Disc16** déjà entraînés.  \n",
    "2. Prenant en entrée un `Gen32_FadeIn` et `Disc32_FadeIn` à entraîner.  \n",
    "3. Parcourt le dataloader en 32×32 (`dataloader_32`) et, à chaque batch :\n",
    "\n",
    "   - Fait évoluer un coefficient $\\alpha \\in [0,1]$ en fonction du **nombre de pas de gradient**.\n",
    "   - Utilise ce $\\alpha$ pour :\n",
    "     - `G(z, alpha)` (blender les chemins 16→32),\n",
    "     - `D(x, alpha)` (blender les chemins 32→16).\n",
    "\n",
    "Une manière simple de faire évoluer $\\alpha$ est :\n",
    "\n",
    "$$\n",
    "\\alpha = \\min\\left(1, \\frac{\\text{step}}{\\text{total\\_steps}}\\right),\n",
    "$$\n",
    "\n",
    "où `step` est le nombre global de batchs vus depuis le début de l'entraînement\n",
    "et `total_steps = num_epochs * len(dataloader_32)`.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Compléter la mise à jour de $\\alpha$ dans la fonction `train_progressive` et vérifier son effet visuel sur les images générées. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b5e2fb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_progressive(\n",
    "    G,\n",
    "    D,\n",
    "    dataloader_32,\n",
    "    num_epochs: int = 5,\n",
    "    z_dim: int = 128,\n",
    "    name: str = \"Progressive_GAN_32x32\",\n",
    "):\n",
    "    \"\"\"\n",
    "    Train a simple progressive GAN from 16x16 to 32x32.\n",
    "\n",
    "    Args:\n",
    "        G: Generator FadeIn (32x32 generator on top of pretrained Gen16)\n",
    "        D: Discriminator FadeIn (32x32 discriminator on top of pretrained Disc16)\n",
    "        dataloader_32: DataLoader over 32x32 images (same dataset, higher resolution)\n",
    "        num_epochs: number of epochs for the progressive phase\n",
    "        z_dim: latent dimension\n",
    "    \"\"\"\n",
    "\n",
    "    criterion = nn.BCELoss()\n",
    "    optim_D = optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))\n",
    "    optim_G = optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))\n",
    "\n",
    "    fixed_noise = torch.randn(64, z_dim, device=device)\n",
    "\n",
    "    step = 0\n",
    "    total_steps =  # TODO\n",
    "    alpha = 0\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        for x, _ in dataloader_32:\n",
    "            x = x.to(device)\n",
    "            b = x.size(0)\n",
    "\n",
    "            # TODO\n",
    "\n",
    "            # TODO\n",
    "\n",
    "        # --------- Visualisation à la fin de l'epoch ---------\n",
    "        with torch.no_grad():\n",
    "            fake_imgs = G(fixed_noise, alpha=alpha).cpu()\n",
    "        grid = utils.make_grid(fake_imgs, nrow=8, normalize=True, value_range=(-1, 1))\n",
    "        plt.figure(figsize=(6, 6))\n",
    "        plt.imshow(grid.permute(1, 2, 0))\n",
    "        plt.axis(\"off\")\n",
    "        plt.title(f\"{name} - Epoch {epoch+1}, alpha={alpha:.2f}\")\n",
    "        if epoch == num_epochs - 1:\n",
    "            plt.savefig(f\"{name}_samples_epoch{epoch+1}.png\")\n",
    "        plt.show()\n",
    "\n",
    "        print(\n",
    "            f\"[Epoch {epoch+1}/{num_epochs}] \"\n",
    "            f\"alpha={alpha:.2f} \"\n",
    "            f\"Loss_D: {loss_D.item():.4f}  Loss_G: {loss_G.item():.4f}\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adfddcc3",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Entraînez votre PG-GAN sur FashionMNIST en 32×32 pendant 10 epochs et affichez les échantillons générés à chaque epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b5bf218",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- 2) Progressive Growing vers 32×32 ---\n",
    "Gpg = Gen32_FadeIn(gen16, nc=nc).to(device)\n",
    "Dpg = Disc32_FadeIn(disc16, nc=nc).to(device)\n",
    "train_progressive(\n",
    "    Gpg,\n",
    "    Dpg,\n",
    "    train_loader,\n",
    "    num_epochs=5,\n",
    "    z_dim=z_dim,\n",
    "    name=\"PGGAN_32x32\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0130579d",
   "metadata": {},
   "source": [
    "## 5. Calibration globale du discriminateur\n",
    "\n",
    "Dans la section précédente, nous avons regardé **l'accuracy par classe** du discriminateur\n",
    "(grâce à `eval_discriminator`).  \n",
    "\n",
    "Nous allons maintenant nous intéresser à une **calibration globale** (toutes classes confondues),\n",
    "en vue d’utiliser le discriminateur pour du **rejection sampling**. On va commencer par analyser la distribution des scores par classes afin de se faire une idée du niveau de confiance du discriminateur. La fonction suivante permet d'afficher les différents scores du discriminateur sur les images réelles et générées, par classe.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84829103",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_per_class_acc(acc_real, acc_fake, title=\"\"):\n",
    "    if title:\n",
    "        print(f\"\\n===== {title} =====\")\n",
    "    for c in range(10):\n",
    "        print(\n",
    "            f\"Classe {c}: \"\n",
    "            f\"real_acc = {acc_real[c]*100:.1f}%   \"\n",
    "            f\"fake_acc = {acc_fake[c]*100:.1f}%\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c8d25a7",
   "metadata": {},
   "source": [
    "Par défaut, on considère que le discriminateur n'est pas calibré, c'est-à-dire que les scores qu'il produit ne correspondent pas à des probabilités fiables. Plus spécifiquement, si on observe des performances très différentes selon que les images sont réelles ou générées, on peut s'attendre à ce que les scores soient biaisés. On considère que le modèle classifie une image comme réelle si son score est supérieure à $0.5$.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Complétez la fonction suivante afin de calculer l'accuracy du discriminateur par classe sur un dataset donné, puis utilisez-la pour évaluer les discriminateur des GANs entraîné précédemment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd6c8eed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_discriminator(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    test_loader,\n",
    "    z_dim=128,\n",
    "    device=device,\n",
    "    conditional=False,\n",
    "    use_pg=False,\n",
    "    alpha=1.0,\n",
    "    n_classes=10,\n",
    "    max_batches=None,\n",
    "):\n",
    "    \"\"\"\n",
    "    Évalue un discriminateur sur un test set (FashionMNIST) + samples générés.\n",
    "\n",
    "    Paramètres\n",
    "    ----------\n",
    "    generator : nn.Module\n",
    "        Générateur G.\n",
    "    discriminator : nn.Module\n",
    "        Discriminateur D.\n",
    "    test_loader : DataLoader\n",
    "        DataLoader du test set (images 32x32 idéalement).\n",
    "    z_dim : int\n",
    "        Dimension du bruit d'entrée du générateur.\n",
    "    conditional : bool\n",
    "        - False : G(z), D(x)\n",
    "        - True  : G(z, y), D(x, y)\n",
    "    use_pg : bool\n",
    "        - False : D(x) ou D(x, y)\n",
    "        - True  : D(x, alpha=alpha) ou D(x, y, alpha=alpha)\n",
    "    alpha : float\n",
    "        Coefficient de fade-in pour PG-GAN (typiquement 1.0 après training).\n",
    "    n_classes : int\n",
    "        Nombre de classes (10 pour FashionMNIST).\n",
    "    max_batches : int ou None\n",
    "        Si non None, limite le nombre de batchs de test (pratique pour aller vite).\n",
    "\n",
    "    Retourne\n",
    "    --------\n",
    "    results : dict\n",
    "        {\n",
    "            \"acc_real\": np.ndarray shape (n_classes,),\n",
    "            \"acc_fake\": np.ndarray shape (n_classes,),\n",
    "            \"scores\":   np.ndarray shape (N_total,),   # proba prédite d'être réel\n",
    "            \"labels\":   np.ndarray shape (N_total,),   # 1 = réel, 0 = fake\n",
    "            \"classes\":  np.ndarray shape (N_total,),   # classe true (0..9)\n",
    "        }\n",
    "    \"\"\"\n",
    "    G = generator.to(device)\n",
    "    D = discriminator.to(device)\n",
    "    G.eval()\n",
    "    D.eval()\n",
    "\n",
    "    real_correct = np.zeros(n_classes, dtype=np.int64)\n",
    "    real_total = np.zeros(n_classes, dtype=np.int64)\n",
    "    fake_correct = np.zeros(n_classes, dtype=np.int64)\n",
    "    fake_total = np.zeros(n_classes, dtype=np.int64)\n",
    "\n",
    "    all_scores = []  # p(real | x)\n",
    "    all_labels = []  # 1 pour réel, 0 pour fake\n",
    "    all_classes = []  # classe de l'image (FashionMNIST)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (x, y) in enumerate(test_loader):\n",
    "            if max_batches is not None and batch_idx >= max_batches:\n",
    "                break\n",
    "\n",
    "            x = x.to(device)\n",
    "            y = y.to(device)\n",
    "            B = x.size(0)\n",
    "\n",
    "            # ===================== RÉEL =====================\n",
    "            # forward du discriminateur\n",
    "            if conditional:\n",
    "                if use_pg:\n",
    "                    out_real = D(x, y, alpha=alpha)  # (B,)\n",
    "                else:\n",
    "                    out_real = D(x, y)  # (B,)\n",
    "            else:\n",
    "                if use_pg:\n",
    "                    out_real = D(x, alpha=alpha)  # (B,)\n",
    "                else:\n",
    "                    out_real = D(x)  # (B,)\n",
    "\n",
    "            # Assumer que out_real est déjà une proba (Sigmoid dans le réseau)\n",
    "            probs_real = out_real.view(-1)  # (B,)\n",
    "            preds_real =   # TODO\n",
    "            preds_real = preds_real.long()\n",
    "            # stats par classe\n",
    "            for c in range(n_classes):\n",
    "                mask = y == c\n",
    "                if mask.any():\n",
    "                    real_total[c] += mask.sum().item()\n",
    "                    real_correct[c] += (preds_real[mask] == 1).sum().item()\n",
    "\n",
    "            # logs pour calibration\n",
    "            all_scores.append(probs_real.detach().cpu())\n",
    "            all_labels.append(torch.ones(B, dtype=torch.long))  # 1 = réel\n",
    "            all_classes.append(y.detach().cpu())\n",
    "\n",
    "            # ===================== FAKE =====================\n",
    "            # même labels de classe y pour compter par classe\n",
    "            z = torch.randn(B, z_dim, device=device)\n",
    "\n",
    "            if conditional:\n",
    "                if use_pg:\n",
    "                    x_fake = G(z, y, alpha=alpha)\n",
    "                    out_fake = D(x_fake, y, alpha=alpha)\n",
    "                else:\n",
    "                    x_fake = G(z, y)\n",
    "                    out_fake = D(x_fake, y)\n",
    "            else:\n",
    "                if use_pg:\n",
    "                    x_fake = G(z, alpha=alpha)\n",
    "                    out_fake = D(x_fake, alpha=alpha)\n",
    "                else:\n",
    "                    x_fake = G(z)\n",
    "                    out_fake = D(x_fake)\n",
    "\n",
    "            probs_fake = out_fake.view(-1)\n",
    "            preds_fake = (probs_fake > 0.5).long()  # 0 attendu (fake)\n",
    "\n",
    "            for c in range(n_classes):\n",
    "                mask = y == c\n",
    "                if mask.any():\n",
    "                    fake_total[c] += mask.sum().item()\n",
    "                    fake_correct[c] += (preds_fake[mask] == 0).sum().item()\n",
    "\n",
    "            # logs pour calibration\n",
    "            all_scores.append(probs_fake.detach().cpu())\n",
    "            all_labels.append(torch.zeros(B, dtype=torch.long))  # 0 = fake\n",
    "            all_classes.append(y.detach().cpu())\n",
    "\n",
    "    acc_real = real_correct / np.maximum(real_total, 1)\n",
    "    acc_fake = fake_correct / np.maximum(fake_total, 1)\n",
    "\n",
    "    all_scores = torch.cat(all_scores, dim=0).numpy()\n",
    "    all_labels = torch.cat(all_labels, dim=0).numpy()\n",
    "    all_classes = torch.cat(all_classes, dim=0).numpy()\n",
    "\n",
    "    results = {\n",
    "        \"acc_real\": acc_real,\n",
    "        \"acc_fake\": acc_fake,\n",
    "        \"scores\": all_scores,\n",
    "        \"labels\": all_labels,\n",
    "        \"classes\": all_classes,\n",
    "    }\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99731eee",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_dcgan = eval_discriminator(\n",
    "    generator=G,\n",
    "    discriminator=D,\n",
    "    test_loader=test_loader,\n",
    "    z_dim=z_dim,\n",
    "    device=device,\n",
    "    conditional=False,\n",
    "    use_pg=False,\n",
    "    alpha=1.0,\n",
    ")\n",
    "\n",
    "print_per_class_acc(\n",
    "    res_dcgan[\"acc_real\"], res_dcgan[\"acc_fake\"], title=\"DCGAN non conditionnel\"\n",
    ")\n",
    "\n",
    "# Pour calibration :\n",
    "scores_dcgan = res_dcgan[\"scores\"]\n",
    "labels_dcgan = res_dcgan[\"labels\"]\n",
    "classes_dcgan = res_dcgan[\"classes\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96b41dec",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_cgan = eval_discriminator(\n",
    "    generator=Gc,\n",
    "    discriminator=Dc,\n",
    "    test_loader=test_loader,\n",
    "    z_dim=z_dim,\n",
    "    device=device,\n",
    "    conditional=True,  # << important\n",
    "    use_pg=False,\n",
    "    alpha=1.0,\n",
    ")\n",
    "\n",
    "print_per_class_acc(\n",
    "    res_cgan[\"acc_real\"], res_cgan[\"acc_fake\"], title=\"cGAN conditionnel\"\n",
    ")\n",
    "\n",
    "scores_cgan = res_cgan[\"scores\"]\n",
    "labels_cgan = res_cgan[\"labels\"]\n",
    "classes_cgan = res_cgan[\"classes\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8df711a",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_pg = eval_discriminator(\n",
    "    generator=Gpg,\n",
    "    discriminator=Dpg,\n",
    "    test_loader=test_loader,\n",
    "    z_dim=z_dim,\n",
    "    device=device,\n",
    "    conditional=False,\n",
    "    use_pg=True,  # << important\n",
    "    alpha=1.0,  # ou la valeur que tu utilises\n",
    ")\n",
    "\n",
    "print_per_class_acc(\n",
    "    res_pg[\"acc_real\"], res_pg[\"acc_fake\"], title=\"PG-GAN 32x32 (non conditionnel)\"\n",
    ")\n",
    "\n",
    "scores_pg = res_pg[\"scores\"]\n",
    "labels_pg = res_pg[\"labels\"]\n",
    "classes_pg = res_pg[\"classes\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b9df23d",
   "metadata": {},
   "source": [
    "On observe que le discriminateur a tendance à être **trop confiant** sur les images fausses au point d'avoir une accuracy proche de 100% sur ces dernières. Cela signifie que les scores produits par le discriminateur sont biaisés et qu'il est nécessaire de calibrer ce dernier avant de l'utiliser pour du rejection sampling. Une manière simple de calibrer le discriminateur est d'appliquer un **décalage de logit**.\n",
    "\n",
    "Idée :\n",
    "\n",
    "1. On mélange un gros batch de **vraies** images (label 1) et de **fausses** images générées (label 0).\n",
    "2. On collecte les scores du discriminateur $p = D(x)$, supposés être des probabilités “réel”.\n",
    "3. On passe en **logits** :\n",
    "\n",
    "$$\n",
    "\\ell = \\log\\frac{p}{1-p}.\n",
    "$$\n",
    "\n",
    "4. On choisit un **décalage** $( t_0 )$ pour recentrer la frontière de décision :\n",
    "   - $t_0 = \\mathrm{médiane}(\\ell) $\n",
    "   - On définit $ \\ell' = \\ell - t_0 $, puis $ p_\\mathrm{calib} = \\sigma(\\ell') $.\n",
    "\n",
    "Avec ce choix, on obtient (approximativement) **50 % des exemples** prédits “réel” après calibration\n",
    "(au seuil 0.5).  \n",
    "Cela nous donne une **version symétrique** du discriminateur, plus pratique pour le DRS.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Dans la fonction `eval_discriminator_global` :\n",
    "- Identifier où l’on :\n",
    "  - calcule les logits à partir des probabilités,\n",
    "  - calcule `t0` comme **médiane** des logits,\n",
    "  - recalcule les probabilités calibrées.  \n",
    "- Comprendre pourquoi ce choix de `t0` force la frontière à 50/50."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b50633",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_discriminator_global(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    test_loader,\n",
    "    z_dim=128,\n",
    "    device=device,\n",
    "    conditional=False,\n",
    "    use_pg=False,\n",
    "    alpha=1.0,\n",
    "    max_batches=None,\n",
    "):\n",
    "    \"\"\"\n",
    "    Évalue un discriminateur sur un mélange de vraies et de fausses images,\n",
    "    et calcule un décalage de logit pour obtenir 50% de prédictions 'réel'.\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "    generator : nn.Module\n",
    "        Générateur G.\n",
    "    discriminator : nn.Module\n",
    "        Discriminateur D (sorties après Sigmoid : p(real | x)).\n",
    "    test_loader : DataLoader\n",
    "        DataLoader du test set FashionMNIST (images 32x32 ou 16x16 selon modèle).\n",
    "    z_dim : int\n",
    "        Dimension du bruit d'entrée du générateur.\n",
    "    conditional : bool\n",
    "        - False : D(x) , G(z)\n",
    "        - True  : D(x, y) , G(z, y)\n",
    "    use_pg : bool\n",
    "        - False : discriminateur classique\n",
    "        - True  : D(..., alpha=alpha) , G(..., alpha=alpha) (PG-GAN)\n",
    "    alpha : float\n",
    "        Coefficient de fade-in utilisé dans le PG-GAN.\n",
    "    max_batches : int ou None\n",
    "        Si non None, limite le nombre de batchs test (pour aller plus vite).\n",
    "\n",
    "    Retour\n",
    "    ------\n",
    "    results : dict\n",
    "        {\n",
    "          \"acc_real\": float,           # accuracy sur les vraies (seuil 0.5, original)\n",
    "          \"acc_fake\": float,           # accuracy sur les fausses\n",
    "          \"t0\": float,                 # décalage de logit pour avoir 50% p_calib>0.5\n",
    "          \"acc_real_calib\": float,     # accuracy sur vraies après calibration\n",
    "          \"acc_fake_calib\": float,     # accuracy sur fausses après calibration\n",
    "          \"scores\": np.ndarray,        # p(real) original, concat (vraies + fausses)\n",
    "          \"labels\": np.ndarray,        # 1 pour réel, 0 pour fake\n",
    "          \"scores_calib\": np.ndarray,  # p_calib(real) après décalage\n",
    "        }\n",
    "    \"\"\"\n",
    "    G = generator.to(device)\n",
    "    D = discriminator.to(device)\n",
    "    G.eval()\n",
    "    D.eval()\n",
    "\n",
    "    n_real_correct = 0\n",
    "    n_real_total = 0\n",
    "    n_fake_correct = 0\n",
    "    n_fake_total = 0\n",
    "\n",
    "    all_scores = []  # p(real | x)\n",
    "    all_labels = []  # 1 = réel, 0 = fake\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (x, y) in enumerate(test_loader):\n",
    "            if max_batches is not None and batch_idx >= max_batches:\n",
    "                break\n",
    "\n",
    "            x = x.to(device)\n",
    "            y = y.to(device)\n",
    "            B = x.size(0)\n",
    "\n",
    "            # ===================== RÉEL =====================\n",
    "            if conditional:\n",
    "                if use_pg:\n",
    "                    out_real = D(x, y, alpha=alpha)\n",
    "                else:\n",
    "                    out_real = D(x, y)\n",
    "            else:\n",
    "                if use_pg:\n",
    "                    out_real = D(x, alpha=alpha)\n",
    "                else:\n",
    "                    out_real = D(x)\n",
    "\n",
    "            probs_real = out_real.view(-1)  # p(real)\n",
    "            preds_real = (probs_real > 0.5).long()  # 1 = prédit réel\n",
    "\n",
    "            n_real_total += B\n",
    "            n_real_correct += (preds_real == 1).sum().item()\n",
    "\n",
    "            all_scores.append(probs_real.cpu())\n",
    "            all_labels.append(torch.ones(B, dtype=torch.long))  # 1 = réel\n",
    "\n",
    "            # ===================== FAKE =====================\n",
    "            z = torch.randn(B, z_dim, device=device)\n",
    "\n",
    "            if conditional:\n",
    "                if use_pg:\n",
    "                    x_fake = G(z, y, alpha=alpha)\n",
    "                    out_fake = D(x_fake, y, alpha=alpha)\n",
    "                else:\n",
    "                    x_fake = G(z, y)\n",
    "                    out_fake = D(x_fake, y)\n",
    "            else:\n",
    "                if use_pg:\n",
    "                    x_fake = G(z, alpha=alpha)\n",
    "                    out_fake = D(x_fake, alpha=alpha)\n",
    "                else:\n",
    "                    x_fake = G(z)\n",
    "                    out_fake = D(x_fake)\n",
    "\n",
    "            probs_fake = out_fake.view(-1)\n",
    "            preds_fake = (probs_fake > 0.5).long()  # 0 attendu\n",
    "\n",
    "            n_fake_total += B\n",
    "            n_fake_correct += (preds_fake == 0).sum().item()\n",
    "\n",
    "            all_scores.append(probs_fake.cpu())\n",
    "            all_labels.append(torch.zeros(B, dtype=torch.long))  # 0 = fake\n",
    "\n",
    "    # -------------------------------\n",
    "    # 1) Accuracies originales\n",
    "    # -------------------------------\n",
    "    acc_real = n_real_correct / max(n_real_total, 1)\n",
    "    acc_fake = n_fake_correct / max(n_fake_total, 1)\n",
    "\n",
    "    scores = torch.cat(all_scores, dim=0).numpy()  # p(real)\n",
    "    labels = torch.cat(all_labels, dim=0).numpy()  # 1 réel / 0 fake\n",
    "\n",
    "    # -------------------------------\n",
    "    # 2) Passage en logit\n",
    "    # -------------------------------\n",
    "    eps = 1e-6\n",
    "    s_clamp = np.clip(scores, eps, 1.0 - eps)\n",
    "    logits = np.log(s_clamp / (1.0 - s_clamp))  # logit = log(p / (1-p))\n",
    "\n",
    "    # -------------------------------\n",
    "    # 3) Décalage de la frontière\n",
    "    #    On cherche t0 tel que 50% des logits soient > t0 et 50% < t0.\n",
    "    #    => t0 = médiane(logits)\n",
    "    # -------------------------------\n",
    "    t0 = np.median(logits)\n",
    "\n",
    "    # logits calibrés : l' -> l - t0\n",
    "    logits_calib = logits - t0\n",
    "    scores_calib = 1.0 / (1.0 + np.exp(-logits_calib))\n",
    "\n",
    "    # -------------------------------\n",
    "    # 4) Accuracies après calibration\n",
    "    #    (seuil 0.5 sur scores_calib)\n",
    "    # -------------------------------\n",
    "    preds_calib = (scores_calib > 0.5).astype(np.int64)\n",
    "\n",
    "    real_mask = labels == 1\n",
    "    fake_mask = labels == 0\n",
    "\n",
    "    acc_real_calib = (\n",
    "        float((preds_calib[real_mask] == 1).mean()) if real_mask.any() else 0.0\n",
    "    )\n",
    "    acc_fake_calib = (\n",
    "        float((preds_calib[fake_mask] == 0).mean()) if fake_mask.any() else 0.0\n",
    "    )\n",
    "\n",
    "    results = {\n",
    "        \"acc_real\": acc_real,\n",
    "        \"acc_fake\": acc_fake,\n",
    "        \"t0\": t0,\n",
    "        \"acc_real_calib\": acc_real_calib,\n",
    "        \"acc_fake_calib\": acc_fake_calib,\n",
    "        \"scores\": scores,\n",
    "        \"labels\": labels,\n",
    "        \"scores_calib\": scores_calib,\n",
    "    }\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90c678a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_dcgan = eval_discriminator_global(\n",
    "    generator=G,  # ConvGenerator entraîné\n",
    "    discriminator=D,  # ConvDiscriminator entraîné\n",
    "    test_loader=test_loader,\n",
    "    z_dim=z_dim,\n",
    "    device=device,\n",
    "    conditional=False,\n",
    "    use_pg=False,\n",
    "    alpha=1.0,\n",
    ")\n",
    "\n",
    "print(\"== DCGAN non cond ==\")\n",
    "print(f\"Acc real (original) : {res_dcgan['acc_real']*100:.2f}%\")\n",
    "print(f\"Acc fake (original) : {res_dcgan['acc_fake']*100:.2f}%\")\n",
    "print(f\"t0 (logit shift)    : {res_dcgan['t0']:.4f}\")\n",
    "print(f\"Acc real (calib)    : {res_dcgan['acc_real_calib']*100:.2f}%\")\n",
    "print(f\"Acc fake (calib)    : {res_dcgan['acc_fake_calib']*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6508c464",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_cgan = eval_discriminator_global(\n",
    "    generator=Gc,  # CondConvGenerator entraîné\n",
    "    discriminator=Dc,  # CondConvDiscriminator entraîné\n",
    "    test_loader=test_loader,\n",
    "    z_dim=z_dim,\n",
    "    device=device,\n",
    "    conditional=True,  # << important\n",
    "    use_pg=False,\n",
    "    alpha=1.0,\n",
    ")\n",
    "\n",
    "print(\"== cGAN conditionnel ==\")\n",
    "print(f\"Acc real (original) : {res_cgan['acc_real']*100:.2f}%\")\n",
    "print(f\"Acc fake (original) : {res_cgan['acc_fake']*100:.2f}%\")\n",
    "print(f\"t0 (logit shift)    : {res_cgan['t0']:.4f}\")\n",
    "print(f\"Acc real (calib)    : {res_cgan['acc_real_calib']*100:.2f}%\")\n",
    "print(f\"Acc fake (calib)    : {res_cgan['acc_fake_calib']*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c73cfab5",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_pg = eval_discriminator_global(\n",
    "    generator=Gpg,  # Gen32_FadeIn entraîné\n",
    "    discriminator=Dpg,  # Disc32_FadeIn (ou équivalent)\n",
    "    test_loader=test_loader,\n",
    "    z_dim=z_dim,\n",
    "    device=device,\n",
    "    conditional=False,\n",
    "    use_pg=True,  # << important\n",
    "    alpha=1.0,  # ou la valeur que tu utilises\n",
    ")\n",
    "\n",
    "print(\"== PG-GAN 32x32 ==\")\n",
    "print(f\"Acc real (original) : {res_pg['acc_real']*100:.2f}%\")\n",
    "print(f\"Acc fake (original) : {res_pg['acc_fake']*100:.2f}%\")\n",
    "print(f\"t0 (logit shift)    : {res_pg['t0']:.4f}\")\n",
    "print(f\"Acc real (calib)    : {res_pg['acc_real_calib']*100:.2f}%\")\n",
    "print(f\"Acc fake (calib)    : {res_pg['acc_fake_calib']*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c53600c8",
   "metadata": {},
   "source": [
    "## 6. Discriminator Rejection Sampling (DRS)\n",
    "\n",
    "Nous avons maintenant :\n",
    "\n",
    "- un discriminateur calibré globalement (via `t0`),\n",
    "- un estimateur de probabilité calibrée $ D_\\mathrm{calib}(x) $.\n",
    "\n",
    "On approxime le **rapport densité cible / densité du générateur** par :\n",
    "\n",
    "$$\n",
    "\\frac{p(x)}{q(x)} \\approx \\frac{D_\\mathrm{calib}(x)}{1 - D_\\mathrm{calib}(x)}.\n",
    "$$\n",
    "\n",
    "Nous allons :\n",
    "\n",
    "1. Estimer un **majorant** $ M \\approx \\max_x p(x)/q(x) $ sur un grand échantillon\n",
    "   de points générés (par exemple 1000).\n",
    "2. Utiliser ce $ M $ pour définir une **probabilité d’acceptation** :\n",
    "   $$\n",
    "   a(x) = \\min\\left(1, \\frac{w(x)}{M}\\right),\\quad\n",
    "   w(x) = \\frac{D_\\mathrm{calib}(x)}{1 - D_\\mathrm{calib}(x)}.\n",
    "   $$\n",
    "\n",
    "La fonction `estimate_M_p_over_q` ci-dessous implémente l’étape 1.\n",
    "\n",
    "<font color='blue'>**TODO:**</font>  Dans la fonction `estimate_M_p_over_q`, à partir des la valeur calibrée du discriminateur, calculez le ratio de densités \\( p(x)/q(x) \\) et estimez son maximum \\( M \\)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "358f3c96",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def estimate_M_p_over_q(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    z_dim,\n",
    "    t0,\n",
    "    device,\n",
    "    n_samples=1000,\n",
    "    batch_size=128,\n",
    "    conditional=False,\n",
    "    num_classes=10,\n",
    "    class_id=None,\n",
    "    use_pg=False,\n",
    "    alpha=1.0,\n",
    "):\n",
    "    \"\"\"\n",
    "    Estime M = max_x p(x)/q(x) ≈ max_x D_calib(x) / (1 - D_calib(x))\n",
    "    en utilisant n_samples images générées.\n",
    "\n",
    "    Paramètres\n",
    "    ----------\n",
    "    generator : nn.Module\n",
    "    discriminator : nn.Module\n",
    "    z_dim : int\n",
    "    t0 : float\n",
    "        Décalage de logit obtenu avec eval_discriminator_global.\n",
    "    device : torch.device\n",
    "    n_samples : int\n",
    "    batch_size : int\n",
    "    conditional : bool\n",
    "        Si True, on suppose G(z,y) et D(x,y).\n",
    "    num_classes : int\n",
    "    class_id : int or None\n",
    "        Si donné et conditional=True, on génère uniquement cette classe.\n",
    "        Sinon, on échantillonne des labels aléatoires.\n",
    "    use_pg : bool\n",
    "    alpha : float\n",
    "\n",
    "    Retour\n",
    "    ------\n",
    "    M : float\n",
    "        Estimation de max p/q (max des poids w(x)).\n",
    "    \"\"\"\n",
    "    G = generator.to(device)\n",
    "    D = discriminator.to(device)\n",
    "    G.eval()\n",
    "    D.eval()\n",
    "\n",
    "    eps = 1e-6\n",
    "    M = 0.0\n",
    "    n_done = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        while n_done < n_samples:\n",
    "            b = min(batch_size, n_samples - n_done)\n",
    "            z = torch.randn(b, z_dim, device=device)\n",
    "\n",
    "            # labels si conditional\n",
    "            if conditional:\n",
    "                if class_id is None:\n",
    "                    y = torch.randint(0, num_classes, (b,), device=device)\n",
    "                else:\n",
    "                    y = torch.full((b,), int(class_id), device=device)\n",
    "            else:\n",
    "                y = None\n",
    "\n",
    "            # Génération et passage dans D\n",
    "            if conditional:\n",
    "                if use_pg:\n",
    "                    x_fake = G(z, y, alpha=alpha)\n",
    "                    out = D(x_fake, y, alpha=alpha)\n",
    "                else:\n",
    "                    x_fake = G(z, y)\n",
    "                    out = D(x_fake, y)\n",
    "            else:\n",
    "                if use_pg:\n",
    "                    x_fake = G(z, alpha=alpha)\n",
    "                    out = D(x_fake, alpha=alpha)\n",
    "                else:\n",
    "                    x_fake = G(z)\n",
    "                    out = D(x_fake)\n",
    "\n",
    "            # out = p(real) brut (avant calibration)\n",
    "            p = out.view(-1).detach().cpu().numpy()\n",
    "            p = np.clip(p, eps, 1.0 - eps)\n",
    "\n",
    "            # ---------------------------\n",
    "            # p -> logit -> out_calib\n",
    "            # ---------------------------\n",
    "            logits = np.log(p / (1.0 - p))  # logit\n",
    "            logits_shift = logits - t0  # recentrage\n",
    "            out_calib = 1.0 / (1.0 + np.exp(-logits_shift))\n",
    "\n",
    "            # w = p/q ≈ odds\n",
    "            w =  # TODO\n",
    "            M =  # TODO\n",
    "            n_done += b\n",
    "\n",
    "    return M"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "263faaa5",
   "metadata": {},
   "source": [
    "### 6.1. Rejection sampling et visualisation\n",
    "\n",
    "Nous allons maintenant utiliser $ M $ et $( t_0 )$ pour :\n",
    "\n",
    "1. Générer un batch de **100 images**.\n",
    "2. Pour chaque image générée :\n",
    "   - calculer $ D_\\mathrm{calib}(x) $,\n",
    "   - calculer $ w(x) = D_\\mathrm{calib}(x) / (1 - D_\\mathrm{calib}(x)) $,\n",
    "   - calculer la **probabilité d’acceptation** $ a(x) = \\min(1, w(x)/M) $,\n",
    "   - tirer un $ u \\sim \\mathcal{U}(0,1) $ et accepter si $ u < a(x) $.\n",
    "3. Afficher les 100 images dans une grille :\n",
    "   - **cadre vert** = échantillon accepté,\n",
    "   - **cadre rouge** = échantillon rejeté,\n",
    "   - titre = probabilité d’acceptation $ a(x) $.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Dans `drs_sample_and_visualize` : Calculer les probabilités d'acceptation et décider de l'acceptation/rejet des échantillons en tirant une variable uniforme `u`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f56d40dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def drs_sample_and_visualize(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    z_dim,\n",
    "    t0,\n",
    "    M,\n",
    "    device,\n",
    "    n_samples=100,\n",
    "    conditional=False,\n",
    "    num_classes=10,\n",
    "    class_id=None,\n",
    "    use_pg=False,\n",
    "    alpha=1.0,\n",
    "    border_size=2,\n",
    "):\n",
    "    \"\"\"\n",
    "    Génère n_samples images, applique DRS avec M et t0, et affiche une grille 10x10\n",
    "    avec un cadre vert (accepté) ou rouge (rejeté) autour de chaque image.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    generator, discriminator : nn.Module\n",
    "    z_dim : int\n",
    "    t0 : float\n",
    "        Décalage de logit pour calibration.\n",
    "    M : float\n",
    "        max p/q estimé par estimate_M_p_over_q.\n",
    "    device : torch.device\n",
    "    n_samples : int\n",
    "        Nombre d'images à générer (idéalement 100).\n",
    "    conditional : bool\n",
    "    num_classes : int\n",
    "    class_id : int or None\n",
    "        Si conditionnel : classe fixe ou labels aléatoires.\n",
    "    use_pg : bool\n",
    "    alpha : float\n",
    "    border_size : int\n",
    "        Largeur du cadre en pixels.\n",
    "    \"\"\"\n",
    "    G = generator.to(device)\n",
    "    D = discriminator.to(device)\n",
    "    G.eval()\n",
    "    D.eval()\n",
    "\n",
    "    eps = 1e-6\n",
    "    n = n_samples\n",
    "\n",
    "    with torch.no_grad():\n",
    "        z = torch.randn(n, z_dim, device=device)\n",
    "        if conditional:\n",
    "            if class_id is None:\n",
    "                y = torch.randint(0, num_classes, (n,), device=device)\n",
    "            else:\n",
    "                y = torch.full((n,), int(class_id), device=device)\n",
    "        else:\n",
    "            y = None\n",
    "\n",
    "        # Génération\n",
    "        if conditional:\n",
    "            if use_pg:\n",
    "                x_fake = G(z, y, alpha=alpha)\n",
    "                out = D(x_fake, y, alpha=alpha)\n",
    "            else:\n",
    "                x_fake = G(z, y)\n",
    "                out = D(x_fake, y)\n",
    "        else:\n",
    "            if use_pg:\n",
    "                x_fake = G(z, alpha=alpha)\n",
    "                out = D(x_fake, alpha=alpha)\n",
    "            else:\n",
    "                x_fake = G(z)\n",
    "                out = D(x_fake)\n",
    "\n",
    "        p = out.view(-1).detach().cpu().numpy()\n",
    "        p = np.clip(p, eps, 1.0 - eps)\n",
    "\n",
    "        logits = np.log(p / (1.0 - p))\n",
    "        logits_shift = logits - t0\n",
    "        D_calib = 1.0 / (1.0 + np.exp(-logits_shift))\n",
    "\n",
    "        w = # TODO\n",
    "        p_accept =   # TODO\n",
    "\n",
    "        # tirage des décisions d'acceptation\n",
    "        u =  # TODO\n",
    "        accepted =  # TODO\n",
    "\n",
    "        # préparation des images avec cadre\n",
    "        x_np = x_fake.detach().cpu().numpy()  # (N,1,H,W) ou (N,C,H,W)\n",
    "        if x_np.ndim == 4 and x_np.shape[1] == 1:\n",
    "            # grayscale -> RGB\n",
    "            x_np = (x_np + 1.0) / 2.0  # [-1,1] -> [0,1]\n",
    "            x_np = np.clip(x_np, 0.0, 1.0)\n",
    "            x_np = np.repeat(x_np, 3, axis=1)  # (N,3,H,W)\n",
    "        elif x_np.ndim == 4 and x_np.shape[1] == 3:\n",
    "            x_np = (x_np + 1.0) / 2.0\n",
    "            x_np = np.clip(x_np, 0.0, 1.0)\n",
    "        else:\n",
    "            raise ValueError(\"Format d'image inattendu pour x_fake\")\n",
    "\n",
    "        N, C, H, W = x_np.shape\n",
    "\n",
    "        framed_imgs = []\n",
    "        colors = []  # \"green\" / \"red\"\n",
    "        for i in range(N):\n",
    "            img = np.transpose(x_np[i], (1, 2, 0))  # (H,W,3)\n",
    "\n",
    "            # canvas avec bordure\n",
    "            canvas = np.zeros(\n",
    "                (H + 2 * border_size, W + 2 * border_size, 3), dtype=np.float32\n",
    "            )\n",
    "\n",
    "            # couleur de la bordure\n",
    "            if accepted[i]:\n",
    "                border_color = np.array([0.0, 1.0, 0.0])  # vert\n",
    "            else:\n",
    "                border_color = np.array([1.0, 0.0, 0.0])  # rouge\n",
    "\n",
    "            # remplir la bordure\n",
    "            canvas[:, :] = border_color\n",
    "            # placer l'image au centre\n",
    "            canvas[border_size:-border_size, border_size:-border_size, :] = img\n",
    "\n",
    "            framed_imgs.append(canvas)\n",
    "            colors.append(\"green\" if accepted[i] else \"red\")\n",
    "\n",
    "    # affichage en grille\n",
    "    n_rows = int(np.ceil(np.sqrt(n_samples)))\n",
    "    n_cols = int(np.ceil(n_samples / n_rows))\n",
    "\n",
    "    fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols, n_rows))\n",
    "    axes = np.array(axes).reshape(-1)\n",
    "\n",
    "    for i in range(n_samples):\n",
    "        ax = axes[i]\n",
    "        ax.imshow(framed_imgs[i])\n",
    "        ax.axis(\"off\")\n",
    "        ax.set_title(f\"{p_accept[i]:.2f}\", fontsize=6)\n",
    "    # masquer axes en trop\n",
    "    for j in range(n_samples, len(axes)):\n",
    "        axes[j].axis(\"off\")\n",
    "\n",
    "    fig.suptitle(\n",
    "        \"Discriminator Rejection Sampling (cadre vert=accepté, rouge=rejeté)\",\n",
    "        fontsize=12,\n",
    "    )\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60779e66",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Lancez les cellules ci-dessous pour observer l'effet du DRS sur les échantillons générés."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc23a334",
   "metadata": {},
   "outputs": [],
   "source": [
    "M_dcgan = estimate_M_p_over_q(\n",
    "    generator=G,\n",
    "    discriminator=D,\n",
    "    z_dim=z_dim,\n",
    "    t0=0,  # récupéré de eval_discriminator_global\n",
    "    device=device,\n",
    "    n_samples=1000,\n",
    "    batch_size=128,\n",
    "    conditional=False,\n",
    "    use_pg=False,\n",
    ")\n",
    "print(\"M (DCGAN) =\", M_dcgan)\n",
    "drs_sample_and_visualize(\n",
    "    generator=G,\n",
    "    discriminator=D,\n",
    "    z_dim=z_dim,\n",
    "    t0=0,\n",
    "    M=M_dcgan,\n",
    "    device=device,\n",
    "    n_samples=100,  # 100 images\n",
    "    conditional=False,\n",
    "    use_pg=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09942fe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "M_dcgan = estimate_M_p_over_q(\n",
    "    generator=G,\n",
    "    discriminator=D,\n",
    "    z_dim=z_dim,\n",
    "    t0=res_dcgan[\"t0\"],  # récupéré de eval_discriminator_global\n",
    "    device=device,\n",
    "    n_samples=1000,\n",
    "    batch_size=128,\n",
    "    conditional=False,\n",
    "    use_pg=False,\n",
    ")\n",
    "print(\"M (DCGAN) =\", M_dcgan)\n",
    "drs_sample_and_visualize(\n",
    "    generator=G,\n",
    "    discriminator=D,\n",
    "    z_dim=z_dim,\n",
    "    t0=res_dcgan[\"t0\"],\n",
    "    M=M_dcgan,\n",
    "    device=device,\n",
    "    n_samples=100,  # 100 images\n",
    "    conditional=False,\n",
    "    use_pg=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97279b54",
   "metadata": {},
   "outputs": [],
   "source": [
    "M_cgan = estimate_M_p_over_q(\n",
    "    generator=Gc,\n",
    "    discriminator=Dc,\n",
    "    z_dim=z_dim,\n",
    "    t0=res_cgan[\"t0\"],\n",
    "    device=device,\n",
    "    n_samples=1000,\n",
    "    conditional=True,\n",
    "    num_classes=10,\n",
    ")\n",
    "print(\"M (cGAN) =\", M_cgan)\n",
    "drs_sample_and_visualize(\n",
    "    generator=Gc,\n",
    "    discriminator=Dc,\n",
    "    z_dim=z_dim,\n",
    "    t0=res_cgan[\"t0\"],\n",
    "    M=M_cgan,\n",
    "    device=device,\n",
    "    n_samples=100,\n",
    "    conditional=True,\n",
    "    num_classes=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eedc8319",
   "metadata": {},
   "outputs": [],
   "source": [
    "M_pggan = estimate_M_p_over_q(\n",
    "    generator=Gpg,\n",
    "    discriminator=Dpg,\n",
    "    z_dim=z_dim,\n",
    "    t0=res_pg[\"t0\"],\n",
    "    device=device,\n",
    "    n_samples=1000,\n",
    "    conditional=False,\n",
    "    num_classes=10,\n",
    ")\n",
    "print(\"M (PGGAN) =\", M_pggan)\n",
    "drs_sample_and_visualize(\n",
    "    generator=Gpg,\n",
    "    discriminator=Dpg,\n",
    "    z_dim=z_dim,\n",
    "    t0=res_pg[\"t0\"],\n",
    "    M=M_pggan,\n",
    "    device=device,\n",
    "    n_samples=100,\n",
    "    conditional=False,\n",
    "    num_classes=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1293798",
   "metadata": {},
   "source": [
    "### 6.2. Visualiser les 100 échantillons et surligner les top-10\n",
    "\n",
    "Le fait que le DRS soit stochastique peut rendre difficile l'observation de son effet sur un petit échantillon (100 images). Pour mieux visualiser l'impact du DRS, nous allons afficher les 100 images générées en mettant en évidence les 10 échantillons avec la plus haute probabilité d'acceptation.\n",
    "\n",
    "La fonction suivante `drs_visualize_100_with_top10_green` :\n",
    "\n",
    "1. Génère **100 images**.\n",
    "2. Calcule `p_accept` pour chacune (comme dans `drs_sample_and_visualize`).\n",
    "3. Trie les échantillons par `p_accept` décroissante.\n",
    "4. Affiche les **100 images** en grille :\n",
    "   - **cadre vert** sur les `top_k` échantillons (les préférés du discriminateur),\n",
    "   - **cadre rouge** sur les autres,\n",
    "   - titre = `p_accept`.\n",
    "\n",
    "<font color='blue'>**TODO:**</font>  Lancer les cellules ci-dessous pour observer l'effet du DRS observant quelles images sont privilégiées par le discriminateur.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8047257e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def drs_visualize_100_with_top10_green(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    z_dim,\n",
    "    t0,\n",
    "    M,\n",
    "    device,\n",
    "    n_samples=100,\n",
    "    top_k=10,\n",
    "    conditional=False,\n",
    "    num_classes=10,\n",
    "    class_id=None,\n",
    "    use_pg=False,\n",
    "    alpha=1.0,\n",
    "    border_size=2,\n",
    "):\n",
    "    \"\"\"\n",
    "    Génère n_samples images, calcule p_accept pour chacune (DRS),\n",
    "    trie les 10 meilleures mais AFFICHE TOUTES les 100 :\n",
    "        - Cadre VERT = top_k images\n",
    "        - Cadre ROUGE = les 100 - top_k\n",
    "    Chaque image affiche sa p_accept en titre.\n",
    "    \"\"\"\n",
    "\n",
    "    G = generator.to(device)\n",
    "    D = discriminator.to(device)\n",
    "    G.eval()\n",
    "    D.eval()\n",
    "\n",
    "    eps = 1e-6\n",
    "\n",
    "    with torch.no_grad():\n",
    "\n",
    "        # ===========================\n",
    "        # 1. Génération\n",
    "        # ===========================\n",
    "        z = torch.randn(n_samples, z_dim, device=device)\n",
    "\n",
    "        if conditional:\n",
    "            if class_id is None:\n",
    "                y = torch.randint(0, num_classes, (n_samples,), device=device)\n",
    "            else:\n",
    "                y = torch.full((n_samples,), int(class_id), device=device)\n",
    "        else:\n",
    "            y = None\n",
    "\n",
    "        if conditional:\n",
    "            if use_pg:\n",
    "                x_fake = G(z, y, alpha=alpha)\n",
    "                out = D(x_fake, y, alpha=alpha)\n",
    "            else:\n",
    "                x_fake = G(z, y)\n",
    "                out = D(x_fake, y)\n",
    "        else:\n",
    "            if use_pg:\n",
    "                x_fake = G(z, alpha=alpha)\n",
    "                out = D(x_fake, alpha=alpha)\n",
    "            else:\n",
    "                x_fake = G(z)\n",
    "                out = D(x_fake)\n",
    "\n",
    "        # p(real)\n",
    "        p = out.view(-1).detach().cpu().numpy()\n",
    "        p = np.clip(p, eps, 1.0 - eps)\n",
    "\n",
    "        # ===========================\n",
    "        # 2. Calibration logit + t0\n",
    "        # ===========================\n",
    "        logits = np.log(p / (1.0 - p))\n",
    "        logits_shift = logits - t0\n",
    "        p_calib = 1.0 / (1.0 + np.exp(-logits_shift))\n",
    "\n",
    "        # Poids DRS\n",
    "        w = p_calib / (1.0 - p_calib)\n",
    "        p_accept = np.minimum(1.0, w / (M + eps))\n",
    "\n",
    "        # ===========================\n",
    "        # 3. Tri : top_k meilleurs\n",
    "        # ===========================\n",
    "        idx_sorted = np.argsort(-p_accept)\n",
    "        idx_topk = set(idx_sorted[:top_k])  # set pour lookup rapide\n",
    "\n",
    "        # ===========================\n",
    "        # 4. Mise en forme des images\n",
    "        # ===========================\n",
    "        x_np = x_fake.detach().cpu().numpy()  # (N,1,H,W) ou (N,3,H,W)\n",
    "        if x_np.ndim == 4 and x_np.shape[1] == 1:\n",
    "            x_np = (x_np + 1.0) / 2.0\n",
    "            x_np = np.clip(x_np, 0.0, 1.0)\n",
    "            x_np = np.repeat(x_np, 3, axis=1)\n",
    "        elif x_np.ndim == 4 and x_np.shape[1] == 3:\n",
    "            x_np = (x_np + 1.0) / 2.0\n",
    "            x_np = np.clip(x_np, 0.0, 1.0)\n",
    "        else:\n",
    "            raise ValueError(\"Format d'image inattendu pour x_fake\")\n",
    "\n",
    "        N, C, H, W = x_np.shape\n",
    "\n",
    "        framed_imgs = []\n",
    "        titles = []\n",
    "\n",
    "        for i in range(n_samples):\n",
    "            img = np.transpose(x_np[i], (1, 2, 0))  # (H,W,3)\n",
    "\n",
    "            canvas = np.zeros(\n",
    "                (H + 2 * border_size, W + 2 * border_size, 3), dtype=np.float32\n",
    "            )\n",
    "\n",
    "            # Couleur du cadre\n",
    "            if i in idx_topk:\n",
    "                border_color = np.array([0.0, 1.0, 0.0])  # VERT (top-k)\n",
    "            else:\n",
    "                border_color = np.array([1.0, 0.0, 0.0])  # ROUGE sinon\n",
    "\n",
    "            canvas[:, :] = border_color\n",
    "            canvas[border_size:-border_size, border_size:-border_size, :] = img\n",
    "\n",
    "            framed_imgs.append(canvas)\n",
    "            titles.append(f\"{p_accept[i]:.2f}\")\n",
    "\n",
    "    # ===========================\n",
    "    # 5. Affichage grille 10x10\n",
    "    # ===========================\n",
    "    n_rows = 10\n",
    "    n_cols = 10\n",
    "\n",
    "    fig, axes = plt.subplots(n_rows, n_cols, figsize=(12, 12))\n",
    "    axes = np.array(axes).reshape(-1)\n",
    "\n",
    "    for i in range(n_samples):\n",
    "        ax = axes[i]\n",
    "        ax.imshow(framed_imgs[i])\n",
    "        ax.axis(\"off\")\n",
    "        ax.set_title(titles[i], fontsize=6)\n",
    "\n",
    "    fig.suptitle(f\"DRS — 100 images (cadre vert = top {top_k})\", fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5df9ca21",
   "metadata": {},
   "outputs": [],
   "source": [
    "drs_visualize_100_with_top10_green(\n",
    "    generator=G,\n",
    "    discriminator=D,\n",
    "    z_dim=z_dim,\n",
    "    t0=res_dcgan[\"t0\"],\n",
    "    M=M_dcgan,\n",
    "    device=device,\n",
    "    n_samples=100,\n",
    "    top_k=10,\n",
    "    conditional=False,\n",
    "    use_pg=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7371e5a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "drs_visualize_100_with_top10_green(\n",
    "    generator=Gc,\n",
    "    discriminator=Dc,\n",
    "    z_dim=z_dim,\n",
    "    t0=res_cgan[\"t0\"],\n",
    "    M=M_cgan,\n",
    "    device=device,\n",
    "    n_samples=100,\n",
    "    top_k=10,\n",
    "    conditional=True,\n",
    "    use_pg=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8df0373",
   "metadata": {},
   "outputs": [],
   "source": [
    "drs_visualize_100_with_top10_green(\n",
    "    generator=Gpg,\n",
    "    discriminator=Dpg,\n",
    "    z_dim=z_dim,\n",
    "    t0=res_pg[\"t0\"],\n",
    "    M=M_pggan,\n",
    "    device=device,\n",
    "    n_samples=100,\n",
    "    top_k=10,\n",
    "    conditional=False,\n",
    "    use_pg=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73a77826",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 7. Discussion & pistes d'amélioration\n",
    "\n",
    "Nous avons vu dans ce TP :\n",
    "\n",
    "1. Comment entraîner :\n",
    "   - un **GAN convolutionnel** simple,\n",
    "   - un **Conditional GAN (cGAN)** avec embeddings de classe,\n",
    "   - un **mini PG-GAN** (16×16 → 32×32) avec paramètre de fade-in $\\alpha$.\n",
    "\n",
    "2. Comment analyser le **discriminateur** :\n",
    "   - accuracy sur **vraies** vs **fausses**,\n",
    "   - accuracy **par classe**,\n",
    "   - **calibration globale** via un décalage de logit $ t_0 $.\n",
    "\n",
    "3. Comment utiliser le discriminateur comme un **estimateur du rapport de densité** $ p(x) / q(x) $ :\n",
    "   - approximation $ p/q \\approx \\frac{D_\\mathrm{calib}(x)}{1 - D_\\mathrm{calib}(x)} $,\n",
    "   - estimation d’un max $ M \\approx \\max_x p(x)/q(x) $,\n",
    "   - **Discriminator Rejection Sampling (DRS)** pour re-pondérer ou filtrer les échantillons.\n",
    "\n",
    "---\n",
    "\n",
    "### 7.1. Limites de l’algorithme actuel\n",
    "\n",
    "Quelques limites (très importantes en pratique) :\n",
    "\n",
    "- Le **discriminateur peut être mal calibré** :\n",
    "  - même après le shift $ t_0 $, ses probabilités restent parfois surestimées ou sous-estimées ;\n",
    "  - la médiane des logits force 50/50 globalement, mais pas forcément sur des régions “intéressantes” de l’espace.\n",
    "\n",
    "- L’approximation $ p/q \\approx \\frac{D_\\mathrm{calib}(x)}{1 - D_\\mathrm{calib}(x)} $ est très heuristique :\n",
    "  - si $ D $ est trop fort ou trop faible, cela fausse complètement les poids ;\n",
    "  - en cas de **mode collapse**, on ne voit jamais certains modes, donc on ne peut pas les “réparer” par DRS.\n",
    "\n",
    "- Le **rejection sampling peut être très inefficace** :\n",
    "  - si $ M $ est très grand, la **probabilité moyenne d’acceptation** devient très faible ;\n",
    "  - en dimension élevée, c’est souvent un gros problème.\n",
    "\n",
    "---\n",
    "\n",
    "### 7.2. Questions ouvertes pour améliorer l’algorithme\n",
    "\n",
    "Voici quelques pistes à discuter / explorer :\n",
    "\n",
    "1. **Meilleure calibration du discriminateur**\n",
    "   - Peut-on utiliser des techniques plus avancées (par ex. *Platt scaling*, *temperature scaling*, *isotonic regression*)  \n",
    "     pour rendre $ D(x) $ mieux calibré ?\n",
    "   - Comment valider cette calibration (par ex. via des *reliability plots* ou des scores de calibration) ?\n",
    "\n",
    "2. **Améliorer l’estimation de $ p/q $**\n",
    "   - Peut-on combiner **plusieurs discriminateurs** (ensemble) pour lisser le rapport de densité ?\n",
    "   - Peut-on apprendre un modèle dédié au **density ratio** plutôt qu’utiliser directement $ D $ ?\n",
    "\n",
    "3. **Réduire l’inefficacité du rejection sampling**\n",
    "   - Au lieu de faire du “tout ou rien” (accepter/rejeter),  \n",
    "     peut-on utiliser les poids $ w(x) $ pour faire de la **re-pondération** (importance sampling) ?\n",
    "   - Peut-on adapter $ M $ dynamiquement ou utiliser une version “tempérée” de $ w(x) $ :\n",
    "     $$\n",
    "     w_T(x) = \\left( \\frac{D_\\mathrm{calib}(x)}{1 - D_\\mathrm{calib}(x)} \\right)^{1/T}\n",
    "     $$\n",
    "     avec une *température* $ T > 1 $ pour limiter les extrêmes ?\n",
    "4. **Lien avec d’autres familles de modèles**\n",
    "   - Comment ce DRS se compare-t-il à :\n",
    "     - le **truncation trick** dans les GANs,\n",
    "     - la **guidance** dans les modèles de diffusion,\n",
    "     - les méthodes basées sur des **energies** (EBM) ?\n",
    "   - Peut-on voir le discriminateur comme une **approximation d’une fonction d’énergie**,\n",
    "     et le DRS comme une forme simplifiée de MCMC ?\n",
    "\n",
    "---\n",
    "\n",
    "### 7.3. Pour aller plus loin (questions pour vous)\n",
    "\n",
    "- Que se passe-t-il si on utilise un discriminateur **par classe** (un D spécialisé sur une seule classe)  \n",
    "  pour faire du DRS conditionnel, plutôt qu’un D global ?\n",
    "- Comment pourrait-on intégrer la **calibration par classe** (vue plus tôt) dans le DRS ?\n",
    "- Si vous aviez un budget de calcul illimité, comment combiner :\n",
    "  - plusieurs générateurs,\n",
    "  - plusieurs discriminateurs,\n",
    "  - plusieurs niveaux de résolution\n",
    "  pour obtenir une **meilleure couverture des modes** et des **échantillons de meilleure qualité** ?\n",
    "\n",
    "> *Note :* le but de cette section est de montrer que, même si l’algorithme vu ici est relativement simple,\n",
    "> il ouvre sur beaucoup de questions de recherche et d’ingénierie dans les modèles génératifs modernes."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "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.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
