example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example Usage\n",
    "Simple example of how to get the FLD of a generative model for CIFAR10 (e.g. in this example RandomGAN).\n",
    "\n",
    "We assume the generative model returns images in range $[0, 1]$ with shape $[B, C, W, H]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "\n",
    "\n",
    "class RandomGAN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(RandomGAN, self).__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.randn((128, 3, 32, 32)).clip(0, 1)\n",
    "\n",
    "\n",
    "GAN = RandomGAN()\n",
    "\n",
    "\n",
    "# Create a no-argument function that returns batches of images\n",
    "def generate_imgs():\n",
    "    x = torch.randn((128, 100))\n",
    "    return GAN(x).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mapping images to features\n",
    "To work in a more perceptually space, we must first map samples to meaningful features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 79/79 [01:06<00:00,  1.18it/s]  \n",
      "                  \r"
     ]
    }
   ],
   "source": [
    "from fld.features.InceptionFeatureExtractor import (\n",
    "    InceptionFeatureExtractor,\n",
    ")  # or DINOv2FeatureExtractor/CLIPFeatureExtractor\n",
    "\n",
    "feature_extractor = InceptionFeatureExtractor()\n",
    "\n",
    "# FLD needs 3 sets of samples: train, test and gen\n",
    "train_dataset = torchvision.datasets.CIFAR10(root=\"./data\", train=True, download=True)\n",
    "train_dataset.name = \"CIFAR10_train\"  # Dataset needs a name to cache features\n",
    "\n",
    "test_dataset = torchvision.datasets.CIFAR10(root=\"./data\", train=False, download=True)\n",
    "test_dataset.name = \"CIFAR10_test\"\n",
    "\n",
    "# get_dataset_features to compute features of a torch.utils.Dataset\n",
    "train_feat = feature_extractor.get_dataset_features(train_dataset)\n",
    "test_feat = feature_extractor.get_dataset_features(test_dataset)\n",
    "\n",
    "# Can get features directly from a model (e.g. RandomGAN) with get_model_features\n",
    "gen_feat = feature_extractor.get_model_features(generate_imgs, num_samples=10_000)\n",
    "\n",
    "# If you've already generated images, features can be obtained from a directory as well!\n",
    "gen_dir_feat = feature_extractor.get_dir_features(\"/path/to/images\", extension=\"png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing metrics\n",
    "Once the three sets of features are ready, metrics can be computed!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random GAN FLD: 55.90\n",
      "Random GAN FID: 440.23\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random GAN FLD Generalization Gap: 0.06\n"
     ]
    }
   ],
   "source": [
    "# And many others!\n",
    "from fld.metrics.FLD import FLD\n",
    "from fld.metrics.FID import FID\n",
    "\n",
    "# Compute FLD, FID\n",
    "cifar_fld = FLD().compute_metric(train_feat, test_feat, gen_feat)\n",
    "cifar_fid = FID().compute_metric(train_feat, None, gen_feat)\n",
    "\n",
    "print(f\"Random GAN FLD: {cifar_fld:.2f}\")\n",
    "print(f\"Random GAN FID: {cifar_fid:.2f}\")\n",
    "\n",
    "# Make sure RandomGAN isn't overfitting (more negative is more overfit)\n",
    "gen_gap = FLD(\"gap\").compute_metric(train_feat, test_feat, gen_feat)\n",
    "print(f\"Random GAN FLD Generalization Gap: {gen_gap:.2f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.9.15"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}