train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from dataset import ReconstructionDataset\n",
    "from encoder import Encoder\n",
    "from decoder import IMDecoder\n",
    "from utils import create_coord_map, initialize_enc_dec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training parameters and device\n",
    "DEV = torch.device(\"mps\")\n",
    "BATCH_SIZE = 64\n",
    "Z_DIM = 32\n",
    "LR = 0.001\n",
    "\n",
    "# We can train progressively with increasing resolutions, but MNIST is simple enough that we can directly train at 28x28\n",
    "PROGRESSIVE_TRAINING_RESOLUTIONS = [28]\n",
    "TRAINING_EPOCHS_PER_RESOLUTION = [i * 2 for i in PROGRESSIVE_TRAINING_RESOLUTIONS]\n",
    "\n",
    "dataset = ReconstructionDataset()\n",
    "enc, dec = initialize_enc_dec(Z_DIM, DEV)\n",
    "opt = optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr = LR)\n",
    "crit = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Main training loop\n",
    "for resolution, epochs in zip(PROGRESSIVE_TRAINING_RESOLUTIONS, TRAINING_EPOCHS_PER_RESOLUTION):\n",
    "    # Setting target resolution for dataset object and creating a new loader\n",
    "    dataset.set_target_resolution(resolution)\n",
    "    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "\n",
    "    # Same coordinate map is used for all predictions of a resolution, so init outside loop\n",
    "    coord_map = create_coord_map(resolution).to(DEV)\n",
    "    for e in range(epochs):\n",
    "        loop = tqdm(loader, total=len(loader), position=0)\n",
    "        loop.set_description_str(f\"Resolution: {resolution}x{resolution} | Epoch: {e}\")\n",
    "        for input_img, target_img in loop:\n",
    "            opt.zero_grad()\n",
    "            input_img, target_img = input_img.to(DEV), target_img.to(DEV)\n",
    "            target_img = target_img.round()\n",
    "            feature_vector = enc(input_img)\n",
    "            predicted_img = dec(feature_vector, coord_map.unsqueeze(0).repeat(feature_vector.shape[0], 1, 1))\n",
    "            loss = crit(predicted_img, target_img.view(target_img.shape[0], -1))\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "            loop.set_postfix(loss = loss.item())\n",
    "\n",
    "        # Output a single example each epoch for sanity check\n",
    "        with torch.no_grad():\n",
    "            sample_input = dataset[0][0].unsqueeze(0)\n",
    "            features = enc(sample_input.to(DEV))\n",
    "            inference_coord_map = create_coord_map(resolution)\n",
    "            output = dec(features, inference_coord_map.unsqueeze(0).to(DEV)).view(resolution, resolution)\n",
    "\n",
    "            fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(6, 2))\n",
    "            ax = ax.flatten()\n",
    "            ax[0].imshow(sample_input[0][0])\n",
    "            ax[0].axis(False)\n",
    "            ax[1].imshow(output.detach().cpu())\n",
    "            ax[1].axis(False)\n",
    "            plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ml",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}