notebooks/train_diffusion_monkey_unconditional.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unconditional Diffusion Training on Monkey Neural Data\n",
    "\n",
    "Trains an unconditional diffusion model on latent representations of monkey neural data. The notebook loads a pretrained autoencoder, creates latent datasets from spike data, and trains a denoiser using DDPM. Training includes EMA model updates and periodic sample visualization. Final model is saved for downstream conditional training or evaluation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# append parent directory to path (../notebooks -> ..)\n",
    "import os\n",
    "import sys\n",
    "sys.path.append(os.path.dirname(os.getcwd()))\n",
    "os.chdir(os.path.dirname(os.getcwd()))\n",
    "\n",
    "import accelerate\n",
    "import lovely_tensors as lt\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "import torch\n",
    "import yaml\n",
    "from diffusers.optimization import get_scheduler\n",
    "from omegaconf import OmegaConf\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from ldns.networks import AutoEncoder, CountWrapper\n",
    "from ldns.utils.plotting_utils import *\n",
    "from ldns.networks import Denoiser\n",
    "from diffusers.training_utils import EMAModel\n",
    "from diffusers.schedulers import DDPMScheduler\n",
    "\n",
    "lt.monkey_patch()\n",
    "matplotlib.rc_file('matplotlibrc') # mackelab plotting style\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "## load config and model path\n",
    "\n",
    "cfg_ae = OmegaConf.load(\"conf/autoencoder-monkey_z=16.yaml\")\n",
    "\n",
    "cfg_yaml = \"\"\"\n",
    "denoiser_model:\n",
    "  C_in: 16\n",
    "  C: 256\n",
    "  kernel: s4\n",
    "  num_blocks: 6\n",
    "  bidirectional: True\n",
    "  num_train_timesteps: 1000\n",
    "training:\n",
    "  lr: 0.001\n",
    "  weight_decay: 0.0\n",
    "  num_epochs: 2000\n",
    "  num_warmup_epochs: 50\n",
    "  batch_size: 512\n",
    "  random_seed: 42\n",
    "  precision: bf16\n",
    "exp_name: diffusion_monkey_unconditional\n",
    "\"\"\"\n",
    "\n",
    "cfg = OmegaConf.create(yaml.safe_load(cfg_yaml))\n",
    "cfg.dataset = cfg_ae.dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ldns.data.monkey import get_monkey_dataloaders\n",
    "\n",
    "# initialize autoencoder model    \n",
    "ae_model = AutoEncoder(\n",
    "    C_in=cfg_ae.model.C_in,\n",
    "    C=cfg_ae.model.C,\n",
    "    C_latent=cfg_ae.model.C_latent,\n",
    "    L=cfg_ae.dataset.signal_length,\n",
    "    num_blocks=cfg_ae.model.num_blocks,\n",
    "    num_blocks_decoder=cfg_ae.model.get(\"num_blocks_decoder\", cfg_ae.model.num_blocks),\n",
    "    num_lin_per_mlp=cfg_ae.model.get(\"num_lin_per_mlp\", 2),  # default 2\n",
    "    bidirectional=cfg_ae.model.get(\"bidirectional\", True),\n",
    ")\n",
    "\n",
    "ae_model = CountWrapper(ae_model)\n",
    "\n",
    "# load pretrained autoencoder\n",
    "ae_model.load_state_dict(torch.load(f\"exp/{cfg_ae.exp_name}/model.pt\"))\n",
    "\n",
    "# set random seeds\n",
    "torch.manual_seed(cfg.training.random_seed)\n",
    "np.random.seed(cfg.training.random_seed)\n",
    "\n",
    "# get dataloaders\n",
    "train_dataloader, val_dataloader, test_dataloader = get_monkey_dataloaders(\n",
    "        cfg_ae.dataset.task, cfg_ae.dataset.datapath, bin_width=5, batch_size=cfg_ae.training.batch_size\n",
    "    )\n",
    "\n",
    "# setup accelerator\n",
    "accelerator = accelerate.Accelerator(\n",
    "    mixed_precision=cfg_ae.training.get(\"precision\", \"no\"),\n",
    ")\n",
    "\n",
    "# prepare model and data for training\n",
    "ae_model = accelerator.prepare(ae_model)\n",
    "\n",
    "(\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    test_dataloader,\n",
    ") = accelerator.prepare(\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    test_dataloader,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_path = f'exp/{cfg.exp_name}'\n",
    "os.makedirs(save_path, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "accelerator.load_state(f\"exp/epoch_140/epoch_140\") # best checkpoint in our case, after this, it overfits on val Poisson loss\n",
    "\n",
    "(\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    test_dataloader,\n",
    ") = accelerator.prepare(\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    test_dataloader,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dataset containing behavior, behavior angle, spike dataset, latents from ae\n",
    "\n",
    "from ldns.data.monkey import LatentMonkeyDataset\n",
    "\n",
    "latent_dataset_train = LatentMonkeyDataset(train_dataloader, ae_model, clip=False)\n",
    "\n",
    "latent_dataset_val = LatentMonkeyDataset(\n",
    "    val_dataloader,\n",
    "    ae_model,\n",
    "    latent_means=latent_dataset_train.latent_means,\n",
    "    latent_stds=latent_dataset_train.latent_stds,\n",
    "    clip=False,\n",
    ")\n",
    "\n",
    "latent_dataset_test = LatentMonkeyDataset(\n",
    "    test_dataloader,\n",
    "    ae_model,\n",
    "    latent_means=latent_dataset_train.latent_means,\n",
    "    latent_stds=latent_dataset_train.latent_stds,\n",
    "    clip=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up dataloaders for diffusion training\n",
    "\n",
    "train_latent_dataloader = torch.utils.data.DataLoader(\n",
    "    latent_dataset_train,\n",
    "    batch_size=cfg.training.batch_size,\n",
    "    shuffle=True,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "val_latent_dataloader = torch.utils.data.DataLoader(\n",
    "    latent_dataset_val,\n",
    "    batch_size=cfg.training.batch_size,\n",
    "    shuffle=False,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "test_latent_dataloader = torch.utils.data.DataLoader(\n",
    "    latent_dataset_test,\n",
    "    batch_size=cfg.training.batch_size,\n",
    "    shuffle=False,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "num_batches = len(train_latent_dataloader)\n",
    "\n",
    "# check if signal length is power of 2\n",
    "if cfg.dataset.signal_length & (cfg.dataset.signal_length - 1) != 0:\n",
    "    cfg.training.precision = \"no\"  # torch.fft doesnt support half if L!=2^x\n",
    "\n",
    "# prepare the denoiser model and dataset\n",
    "(\n",
    "    train_latent_dataloader,\n",
    "    val_latent_dataloader,\n",
    "    test_latent_dataloader,\n",
    ") = accelerator.prepare(\n",
    "    train_latent_dataloader,\n",
    "    val_latent_dataloader,\n",
    "    test_latent_dataloader,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## initialize (unconditional) denoiser\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "denoiser = Denoiser(\n",
    "    C_in=cfg.denoiser_model.C_in,\n",
    "    C=cfg.denoiser_model.C,\n",
    "    L=cfg.dataset.signal_length,\n",
    "    num_blocks=cfg.denoiser_model.num_blocks,\n",
    "    bidirectional=cfg.denoiser_model.get(\"bidirectional\", True),\n",
    ")\n",
    "\n",
    "# initial values may be way off, scaling down the output layer makes training faster\n",
    "denoiser.conv_out.weight.data = denoiser.conv_out.weight.data * 0.1\n",
    "denoiser.conv_out.bias.data = denoiser.conv_out.bias.data * 0.1\n",
    "\n",
    "scheduler = DDPMScheduler(\n",
    "    num_train_timesteps=cfg.denoiser_model.num_train_timesteps,\n",
    "    clip_sample=False,\n",
    "    beta_schedule=\"linear\"\n",
    ")\n",
    "\n",
    "\n",
    "optimizer = torch.optim.AdamW(\n",
    "    denoiser.parameters(), lr=cfg.training.lr\n",
    ")  # default wd=0.01 for now\n",
    "\n",
    "\n",
    "\n",
    "num_batches = len(train_latent_dataloader)\n",
    "lr_scheduler = get_scheduler(\n",
    "    name=\"cosine\",\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=num_batches * cfg.training.num_warmup_epochs,  # warmup for 10% of epochs\n",
    "    num_training_steps=num_batches * cfg.training.num_epochs * 1.3,  # total number of steps until 0 lr, we use 1.3 to no go all the way to 0 lr\n",
    ")\n",
    "\n",
    "# prepare the denoiser model and dataset\n",
    "(\n",
    "    denoiser,\n",
    "    optimizer,\n",
    "    lr_scheduler,\n",
    ") = accelerator.prepare(\n",
    "    denoiser,\n",
    "    optimizer,\n",
    "    lr_scheduler,\n",
    ")\n",
    "\n",
    "ema_model = EMAModel(denoiser)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some helper functions for checking perf during training of diffusion model\n",
    "\n",
    "def sample(\n",
    "    ema_denoiser,\n",
    "    scheduler,\n",
    "    cfg,\n",
    "    batch_size=1,\n",
    "    generator=None,\n",
    "    device=\"cuda\",\n",
    "    signal_length=None\n",
    "):\n",
    "    \"\"\"Sample latents from the diffusion model.\n",
    "\n",
    "    Args:\n",
    "        ema_denoiser: EMA model wrapper around denoiser\n",
    "        scheduler: DDPM noise scheduler\n",
    "        cfg: Config dictionary containing model parameters\n",
    "        batch_size: Number of samples to generate\n",
    "        generator: Random number generator for reproducibility\n",
    "        device: Device to run sampling on\n",
    "        signal_length: Length of signal to generate. If None, uses cfg.dataset.signal_length\n",
    "\n",
    "    Returns:\n",
    "        Sampled latent tensors of shape (batch_size, C_in, signal_length)\n",
    "    \"\"\"\n",
    "    if signal_length is None:\n",
    "        signal_length = cfg.dataset.signal_length\n",
    "\n",
    "    # sample initial noise\n",
    "    z_t = torch.randn(\n",
    "        (batch_size, cfg.denoiser_model.C_in, signal_length),\n",
    "        device=device\n",
    "    )\n",
    "\n",
    "    # get averaged ema model\n",
    "    ema_denoiser_avg = ema_denoiser.averaged_model\n",
    "    ema_denoiser_avg.eval()\n",
    "\n",
    "    # set up sampling timesteps\n",
    "    scheduler.set_timesteps(cfg.denoiser_model.num_train_timesteps)\n",
    "\n",
    "    # iteratively denoise\n",
    "    for t in tqdm(scheduler.timesteps, desc=\"Sampling DDPM\"):\n",
    "        with torch.no_grad():\n",
    "            model_output = ema_denoiser_avg(\n",
    "                z_t, \n",
    "                torch.tensor([t] * batch_size, device=device).long()\n",
    "            )\n",
    "        z_t = scheduler.step(\n",
    "            model_output, t, z_t, generator=generator, return_dict=False\n",
    "        )[0]\n",
    "\n",
    "    return z_t\n",
    "\n",
    "\n",
    "def sample_spikes(\n",
    "    ema_denoiser,\n",
    "    scheduler,\n",
    "    ae,\n",
    "    cfg,\n",
    "    batch_size=1, \n",
    "    device=\"cuda\"\n",
    "):\n",
    "    \"\"\"Sample spike trains from the diffusion model.\n",
    "\n",
    "    Args:\n",
    "        ema_denoiser: EMA model wrapper around denoiser\n",
    "        scheduler: DDPM noise scheduler\n",
    "        ae: Autoencoder model for decoding latents to rates\n",
    "        cfg: Config dictionary containing model parameters\n",
    "        batch_size: Number of samples to generate\n",
    "        device: Device to run sampling on\n",
    "\n",
    "    Returns:\n",
    "        Tuple containing:\n",
    "            - Sampled spike trains of shape (batch_size, C_in, signal_length)\n",
    "            - Firing rates used to generate spikes\n",
    "    \"\"\"\n",
    "    # sample initial noise\n",
    "    z_t = torch.randn(\n",
    "        (batch_size, cfg.denoiser_model.C_in, cfg.dataset.signal_length),\n",
    "        device=device\n",
    "    )\n",
    "\n",
    "    # get averaged ema model\n",
    "    ema_denoiser_avg = ema_denoiser.averaged_model\n",
    "    ema_denoiser_avg.eval()\n",
    "\n",
    "    # set up sampling timesteps\n",
    "    scheduler.set_timesteps(cfg.denoiser_model.num_train_timesteps)\n",
    "\n",
    "    # iteratively denoise\n",
    "    for t in tqdm(scheduler.timesteps, desc=\"Sampling DDPM\"):\n",
    "        with torch.no_grad():\n",
    "            model_output = ema_denoiser_avg(\n",
    "                z_t,\n",
    "                torch.tensor([t] * batch_size, device=device).long()\n",
    "            )\n",
    "        z_t = scheduler.step(model_output, t, z_t, return_dict=False)[0]\n",
    "\n",
    "    # unnormalize latents\n",
    "    z_t = z_t * latent_dataset_train.latent_stds.to(z_t.device) + \\\n",
    "          latent_dataset_train.latent_means.to(z_t.device)\n",
    "\n",
    "    # decode latents to rates and sample spikes\n",
    "    with torch.no_grad():\n",
    "        rates = ae.decode(z_t).cpu()\n",
    "    spikes = torch.poisson(rates)\n",
    "\n",
    "    return spikes, rates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Diffusion training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use smooth l1 loss for faster convergence than mse\n",
    "loss_fn = torch.nn.SmoothL1Loss(beta=0.04, reduction=\"mean\")\n",
    "\n",
    "# training loop\n",
    "\n",
    "# sampled rates will be very high at the beginning,\n",
    "# then converge to inferred rates during training\n",
    "# this is expected behavior\n",
    "\n",
    "\n",
    "# flags for different eval plots\n",
    "plot_sample_comparison = True\n",
    "plot_spike_count_dist = False \n",
    "plot_per_neuron_dist = False\n",
    "\n",
    "pbar = tqdm(range(0, cfg.training.num_epochs), desc=\"epochs\")\n",
    "for epoch in pbar:\n",
    "    for i, batch in enumerate(train_latent_dataloader):\n",
    "        denoiser.train()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # get batch and add noise\n",
    "        z = batch[\"latent\"] \n",
    "        t = torch.randint(0, cfg.denoiser_model.num_train_timesteps, (z.shape[0],), device=\"cpu\").long()\n",
    "        noise = torch.randn_like(z)\n",
    "        noisy_z = scheduler.add_noise(z, noise, t)\n",
    "        \n",
    "        # predict noise and compute loss\n",
    "        noise_pred = denoiser(noisy_z, t)\n",
    "        loss = loss_fn(noise, noise_pred)\n",
    "\n",
    "        # optimization step\n",
    "        accelerator.backward(loss)\n",
    "        accelerator.clip_grad_norm_(denoiser.parameters(), 1.0)\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "\n",
    "        # update progress bar\n",
    "        if i % 10 == 0:\n",
    "            pbar.set_postfix({\"loss\": loss.item(), \"lr\": lr_scheduler.get_last_lr()[0]})\n",
    "\n",
    "        # update ema model\n",
    "        ema_model.step(denoiser)\n",
    "\n",
    "    # evaluation and plotting\n",
    "    if (epoch) % 100 == 0 and plot_sample_comparison:\n",
    "        denoiser.eval()\n",
    "\n",
    "\n",
    "        # generate samples\n",
    "        sampled_latents = sample(ema_denoiser=ema_model, scheduler=scheduler, \n",
    "                               cfg=cfg, batch_size=2, device=\"cuda\")\n",
    "        \n",
    "        # unnormalize samples\n",
    "        sampled_latents = sampled_latents * latent_dataset_train.latent_stds.to(sampled_latents.device) + \\\n",
    "                         latent_dataset_train.latent_means.to(sampled_latents.device)\n",
    "\n",
    "        # get real samples for comparison\n",
    "        real_latents = latent_dataset_train.latents[:2].cuda()\n",
    "        real_latents = real_latents * latent_dataset_train.latent_stds.to(real_latents.device) + \\\n",
    "                      latent_dataset_train.latent_means.to(real_latents.device)\n",
    "\n",
    "\n",
    "        # decode latents to rates\n",
    "        with torch.no_grad():\n",
    "            sampled_rates = ae_model.decode(sampled_latents).cpu()\n",
    "            decoded_rates_from_real_latents = ae_model.decode(real_latents).cpu()\n",
    "        # plot comparison between sampled and real rates\n",
    "        fig, ax = plt.subplots(1, 2, figsize=(6,4))\n",
    "        \n",
    "        # plot sampled rates\n",
    "        im = ax[0].imshow(sampled_rates[0], aspect=\"auto\")\n",
    "        ax[0].set_title(\"sampled rates\")\n",
    "        fig.colorbar(im, ax=ax[0], orientation=\"vertical\", fraction=0.046, pad=0.04)\n",
    "\n",
    "        # plot real rates\n",
    "        im = ax[1].imshow(decoded_rates_from_real_latents[0], aspect=\"auto\")\n",
    "        ax[1].set_title(\"real rates\") \n",
    "        fig.colorbar(im, ax=ax[1], orientation=\"vertical\", fraction=0.046, pad=0.04)\n",
    "        \n",
    "        fig.tight_layout()\n",
    "        plt.show()\n",
    "\n",
    "pbar.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save model\n",
    "torch.save(accelerator.unwrap(ema_model.averaged_model), f\"exp/{cfg.exp_name}/model.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train conditional diffusion head to `notebooks/train_diffusion_monkey_angle_conditioned.ipynb` and train conditional diffusion head to `notebooks/train_diffusion_monkey_velocity_conditioned.ipynb`.\n",
    "\n",
    "To evaluate the model, go to `notebooks/plotting_diffusion_monkey_unconditional.ipynb`.\n",
    "\n",
    "To train a spike history model, go to `notebooks/train_with_spike_history_monkey.ipynb`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "timeseries",
   "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.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}