experiments/TrainGeneration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9a9cd0dc-9151-49f2-ba8a-c5628f4517f6",
   "metadata": {},
   "source": [
    "# Imports and environment setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b5c082b-b044-4211-9ba8-696f79ee1024",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change Directory to be outside the experiments folder\n",
    "%cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc4a2726-2925-4afd-9a5e-39036b3841cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#import \n",
    "from datasets.shapenet_pointflow_sparse_cond import get_dataloaders\n",
    "from pclab.utils import DataLoaders\n",
    "from utils.callbacks import *\n",
    "from pclab.learner import TrainLearner, ProgressCB, BatchSchedCB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16e33c69-70a3-4c5e-83f1-871aa48d379a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "import torchsparse\n",
    "from torchsparse.utils.collate import sparse_collate_fn\n",
    "from pclab.utils import DataLoaders\n",
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56569dca-025d-438a-9ec9-8761d8e97e00",
   "metadata": {},
   "source": [
    "## Datasets and Dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dacac80f-f81b-4b17-b49d-b85b0db74f4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"/home/vvrbeast/Desktop/Giannis/Data/ShapeNetCore.v2.PC15k\"\n",
    "path = '/home/tourloid/Desktop/PhD/Data/ShapeNetCore.v2.PC15k/'\n",
    "\n",
    "categories = ['chair']\n",
    "\n",
    "tr_dl, te_dl = get_dataloaders(path, categories)\n",
    "dls = DataLoaders(tr_dl, te_dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c252c1d-a35e-4e53-8102-4e6c3016a933",
   "metadata": {},
   "source": [
    "# Load the model\n",
    "\n",
    "Uncomment and run one of the following cells to select the model version."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88c43cb3-d1cd-44b6-af3d-62f3ab56835e",
   "metadata": {},
   "source": [
    "## SVD-S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf8d084e-1dc7-484d-aa0a-22fe31dd8f97",
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.ddpm_unet_attn import SPVUnet\n",
    "model = SPVUnet(voxel_size=0.1, nfs=(32, 64, 128, 256), num_layers=1, attn_chans=8, attn_start=3)\n",
    "checkpoint_name = 'ddpm_unet_attn_64_128_256_256'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc12dea7-ba08-4d83-9de3-55acb7175a91",
   "metadata": {},
   "source": [
    "## SPVD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c9ff345-50bc-4933-a679-5e1e9dcea2f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from models import SPVD\n",
    "# model=SPVD\n",
    "# checkpoint_name = 'SPVD'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c87d320-1b9a-47da-8f11-3ca7f4768e69",
   "metadata": {},
   "source": [
    "## SPVD-L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cea1a4c8-3cde-451c-8eec-5245323f2582",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from models import SPVD_L\n",
    "# model=SPVD_L\n",
    "# checkpoint_name = 'SPVD_L'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f9a70d4-67a7-4da7-9cb1-e2144cab5271",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "120b51b6-0cb7-43d0-8d00-349bcd87bf25",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.001\n",
    "epochs = 2000 \n",
    "\n",
    "# scheduler\n",
    "total_steps = epochs * len(dls.train)\n",
    "sched = partial(torch.optim.lr_scheduler.OneCycleLR, max_lr=lr, total_steps = total_steps)\n",
    "\n",
    "# Callbacks\n",
    "ddpm_cb = DDPMCB()\n",
    "checkpoint_cb = CheckpointCB(2000, checkpoint_name, run_params={})\n",
    "cbs = [ddpm_cb, DeviceCBSparse(), ProgressCB(plot=False), LossCB(), GradientClipCB(), checkpoint_cb, BatchSchedCB(sched)]\n",
    "\n",
    "learn = TrainLearner(model, dls, nn.MSELoss(), lr=lr, cbs=cbs, opt_func=torch.optim.Adam)\n",
    "learn.fit(epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c163fb7-9382-42a3-bbae-db332ce320fb",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d6c4218-ba1a-466f-b182-7ce16c50cfd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.schedulers import DDPMSparseSchedulerGPU\n",
    "from utils.visualization import quick_vis_batch\n",
    "vis_batch = partial(quick_vis_batch, x_offset = 8, y_offset=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd37291b-4ea0-4f9b-8c6e-70bb780df1c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddpm_sched = DDPMSparseSchedulerGPU(n_steps=1000, beta_min=0.0001, beta_max=0.02, pres=1e-5)\n",
    "preds = ddpm_sched.sample(model, 32, 2048)\n",
    "vis_batch(preds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8651ac0-7757-488a-aaf7-1fa0bdd23a9d",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc88287-d6c5-4ebf-b1d6-472d62c24ed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddpm_sched = DDPMSparseSchedulerGPU(n_steps=1000, beta_min=0.0001, beta_max=0.02, sigma='coef_bt')\n",
    "evaluate_gen(path, model, ddpm_sched, save_path='./results/', cates=categories)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "246c1fdd-6d54-403e-be66-0a6f887e3bfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddpm_sched = DDPMSparseSchedulerGPU(n_steps=1000, beta_min=0.0001, beta_max=0.02, sigma='coef_bt')\n",
    "evaluate_gen(path, model, ddpm_sched, save_path='./results/', cates=categories)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ace29c65-a4e2-475f-bd6f-71b6d1f41e35",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddpm_sched = DDPMSparseSchedulerGPU(n_steps=1000, beta_min=0.0001, beta_max=0.02, sigma='coef_bt')\n",
    "evaluate_gen(path, model, ddpm_sched, save_path='./results/', cates=categories)"
   ]
  }
 ],
 "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}