grassy_mnist_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aef790b-3427-4bd1-87a4-d58aedb0549d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pytorch_lightning as pl\n",
    "import torch.nn as nn\n",
    "\n",
    "from contrastive_fs.data import LabeledDataset\n",
    "from torch.utils.data import DataLoader\n",
    "from contrastive_fs.models import CFS_SG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "119c043c-7484-42c4-8e52-3447bbf698e4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Grab the data from Google Drive here: https://drive.google.com/file/d/14YaCSlOM46N7aTS9_lCnsn8S2CMzz-gZ/view?usp=sharing\n",
    "# and deposit it into data/Grassy_MNIST\n",
    "\n",
    "batch_size = 128\n",
    "background = np.load(os.path.join(\"data\", \"Grassy_MNIST\", \"background.npy\")).astype(np.float32)\n",
    "target = np.load(os.path.join(\"data\", \"Grassy_MNIST\", \"target.npy\")).astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "955c23a6-0c22-49d4-91df-802d7ee5decf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Label here determines target versus background\n",
    "labels_train = np.concatenate([np.zeros(background.shape[0]), np.ones(target.shape[0])])\n",
    "data_train = np.concatenate([background, target])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23ff0a48-0764-452e-a4f0-94a6764b353e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_size = data_train.shape[1]\n",
    "output_size = data_train.shape[1]\n",
    "dataset = LabeledDataset(data_train, labels_train)\n",
    "\n",
    "model = CFS_SG(\n",
    "    input_size=input_size,\n",
    "    output_size=output_size,\n",
    "    hidden=[512, 512], # Number of units in each hidden layer\n",
    "    k_prime=20, # Background dimension size\n",
    "    lam=0.175, # Tuned to select about 10 features\n",
    "    lr=1e-3,\n",
    "    loss_fn=nn.MSELoss()\n",
    ")\n",
    "\n",
    "loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "trainer = pl.Trainer(max_epochs=10, accelerator='gpu', devices=1)\n",
    "trainer.fit(model, loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e69b7bf3-e795-44c3-a801-ba5b85f70d20",
   "metadata": {},
   "outputs": [],
   "source": [
    "indices = model.get_inds(10) # Get indices with top 10 most strongly selected features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46d39f40-d9e0-49c2-ac10-cc99d69f3dea",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(15, 10))\n",
    "\n",
    "blank_image = np.zeros(784)\n",
    "blank_image[indices] = 1\n",
    "\n",
    "axes[0].imshow(blank_image.reshape(28, 28))\n",
    "axes[0].set_xticks([])\n",
    "axes[0].set_yticks([])\n",
    "axes[0].set_title(\"CFS Features\", fontsize=16)\n",
    "\n",
    "axes[1].imshow(target[0].reshape(28, 28), cmap='gray')\n",
    "axes[1].set_title(\"Example target image\", fontsize=16)\n",
    "axes[1].set_xticks([])\n",
    "axes[1].set_yticks([])\n",
    "\n",
    "axes[2].imshow(background[0].reshape(28, 28), cmap='gray')\n",
    "axes[2].set_title(\"Example background image\", fontsize=16)\n",
    "axes[2].set_xticks([])\n",
    "axes[2].set_yticks([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb5b75a2-a264-4167-9ecf-2bf91b5923ed",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}