Gaussian diff Model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6d98e7fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import copy\n",
    "import sys\n",
    "import time\n",
    "from random import seed\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib import animation\n",
    "from torch import optim\n",
    "\n",
    "import dataset\n",
    "import evaluation\n",
    "from GaussianDiffusion import GaussianDiffusionModel, get_beta_schedule\n",
    "from helpers import *\n",
    "from UNet import UNetModel, update_ema_params\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "ROOT_DIR = \"./\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7817d0ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = {\n",
    "  \"img_size\": [\n",
    "    256,\n",
    "    256\n",
    "  ],\n",
    "  \"Batch_Size\": 16,\n",
    "  \"EPOCHS\": 3000,\n",
    "  \"T\": 1000,\n",
    "  \"base_channels\": 128,\n",
    "  \"beta_schedule\": \"linear\",\n",
    "  \"channel_mults\": \"\",\n",
    "  \"loss-type\": \"l2\",\n",
    "  \"loss_weight\": \"none\",\n",
    "  \"train_start\": True,\n",
    "  \"lr\": 1e-4,\n",
    "  \"random_slice\": False,\n",
    "  \"sample_distance\": 800,\n",
    "  \"weight_decay\": 0.0,\n",
    "  \"save_imgs\": False,\n",
    "  \"save_vids\": True,\n",
    "  \"dropout\": 0,\n",
    "  \"attention_resolutions\": \"16,8\",\n",
    "  \"num_heads\": 2,\n",
    "  \"num_head_channels\": -1,\n",
    "  \"noise_fn\": \"simplex\",\n",
    "  \"dataset\": \"mri\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "47d5231f",
   "metadata": {},
   "outputs": [],
   "source": [
    "in_channels=3\n",
    "unet = UNetModel(\n",
    "        args['img_size'][0], args['base_channels'], channel_mults=args['channel_mults'], in_channels=in_channels\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "85846468",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(10, 3, 256, 256)\n",
    "t_batch = torch.tensor([1], device=x.device).repeat(x.shape[0])\n",
    "y_hat = unet(x, t_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "49998a42",
   "metadata": {},
   "outputs": [],
   "source": [
    "betas = get_beta_schedule(args['T'], args['beta_schedule'])\n",
    "\n",
    "diff = GaussianDiffusionModel(\n",
    "        args['img_size'], betas, loss_weight=args['loss_weight'],\n",
    "        loss_type=args['loss-type'], noise=args[\"noise_fn\"], img_channels=in_channels\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5bdd5467",
   "metadata": {},
   "outputs": [],
   "source": [
    "# t_batch.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7b319da2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# diff.noise_fn(x, t_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32aa4c90",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}