2-train_denoisers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a Hierarchical DivNoising network with a Direct Denoiser for Convallaria data which is intrinsically noisy\n",
    "In this notebook, we train a Hierarchical DivNoising Ladder VAE alongside a Direct Denoising network for an intrinsically noisy data. This requires having a noise model (model of the imaging noise) which can be either measured from calibration data or bootstrapped from raw noisy images themselves. If you haven't done so, please first run '1-train_noise_model.ipynb', which will download the data and create a noise model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "import tifffile\n",
    "import numpy as np\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.callbacks import EarlyStopping\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "\n",
    "from hdn.lib import utils\n",
    "from hdn.models.lvae import LadderVAE\n",
    "from unet import UNet\n",
    "from hdn.lib.gaussianMixtureNoiseModel import GaussianMixtureNoiseModel\n",
    "from backbone import Backbone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_cuda = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Specify ```path``` to load training data\n",
    "Your data should be stored in the directory indicated by ```path```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"./data/Convallaria_diaphragm/\"\n",
    "observation = tifffile.imread(path + \"20190520_tl_25um_50msec_05pc_488_130EM_Conv.tif\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Data Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For training we need to follow some preprocessing steps first which will prepare the data for training purposes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first divide the data into training and validation sets with 85% images allocated to training set  and rest to validation set. Then we augment the training data 8-fold by 90 degree rotations and flips."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = observation[: int(0.85 * observation.shape[0])]\n",
    "val_data = observation[int(0.85 * observation.shape[0]) :]\n",
    "print(\n",
    "    \"Shape of training images:\",\n",
    "    train_data.shape,\n",
    "    \"Shape of validation images:\",\n",
    "    val_data.shape,\n",
    ")\n",
    "train_data = utils.augment_data(\n",
    "    train_data\n",
    ")  ### Data augmentation disabled for fast training, but can be enabled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### We extract overlapping patches of size ```patch_size x patch_size``` from training and validation images.\n",
    "### Usually 64x64 patches work well for most microscopy datasets\n",
    "patch_size = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_width = observation.shape[2]\n",
    "img_height = observation.shape[1]\n",
    "num_patches = int(float(img_width * img_height) / float(patch_size**2) * 1)\n",
    "train_images = utils.extract_patches(train_data, patch_size, num_patches)\n",
    "val_images = utils.extract_patches(val_data, patch_size, num_patches)\n",
    "val_images = val_images[\n",
    "    :1000\n",
    "]  # We limit validation patches to 1000 to speed up training but it is not necessary\n",
    "img_shape = (train_images.shape[1], train_images.shape[2])\n",
    "print(\n",
    "    \"Shape of training images:\",\n",
    "    train_images.shape,\n",
    "    \"Shape of validation images:\",\n",
    "    val_images.shape,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<code>batch_size</code>: Number of patches for which loss will be calculated before updating weights.<br>\n",
    "<code>virtual_batch</code>: Number of patches that will be passed through the network at a time. Increase to save time, decrease to save memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### We create PyTorch dataloaders for training and validation data\n",
    "batch_size = 64 \n",
    "virtual_batch = 8\n",
    "n_grad_batches = batch_size // virtual_batch\n",
    "\n",
    "train_images = torch.from_numpy(train_images[:, np.newaxis]).float()\n",
    "val_images = torch.from_numpy(val_images[:, np.newaxis]).float()\n",
    "train_dataset = TensorDataset(train_images)\n",
    "val_dataset = TensorDataset(val_images)\n",
    "train_loader = DataLoader(train_dataset, batch_size=virtual_batch, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=virtual_batch, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configure Hierarchical DivNoising model and the Direct Denoiser"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<code>model_name</code> specifies the name of the model with which the weights will be saved and wil be loaded later for prediction.<br>\n",
    "<code>checkpoint_path</code> specifies the directory where the model weights will be saved. <br>\n",
    "<code>gaussian_noise_std</code> is only applicable if dataset is synthetically corrupted with Gaussian noise of known std. For real datasets, it should be set to ```None```.<br>\n",
    "<code>noiseModel</code> specifies a noise model for training. If noisy data is generated synthetically using Gaussian noise, set it to None. Else set it to the GMM based noise model (.npz file)  generated from '1-CreateNoiseModel.ipynb'.<br>\n",
    "<code>lr</code> specifies the learning rate.<br>\n",
    "<code>max_epochs</code> specifies the total number of training epochs. Around $150-200$ epochs work well generally.<br>\n",
    "<code>steps_per_epoch</code> specifies how many steps to take per epoch of training. Around $400-500$ steps work well for most datasets.<br>\n",
    "<code>num_latents</code> specifies the number of stochastic layers. The default setting of $6$ works well for most datasets but quite good results can also be obtained with as less as $4$ layers. However, more stochastic layers may improve performance for some datasets at the cost of increased training time.<br>\n",
    "<code>z_dims</code> specifies the number of bottleneck dimensions (latent space dimensions) at each stochastic layer per pixel. The default setting of $32$ works well for most datasets.<br>\n",
    "<code>blocks_per_layer</code> specifies how many residual blocks to use per stochastic layer. Usually, setting it to be $4$ or more works well. However, more residual blocks improve performance at the cost of increased training time.<br>\n",
    "<code>batchnorm</code> specifies if batch normalization is used or not. Turning it to True is recommended.<br>\n",
    "<code>free_bits</code> specifies the threshold below which KL loss is not optimized for. This prevents the [KL-collapse problem](https://arxiv.org/pdf/1511.06349.pdf%3Futm_campaign%3DRevue%2520newsletter%26utm_medium%3DNewsletter%26utm_source%3Drevue). The default setting of $1.0$ works well for most datasets.<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"convallaria\"\n",
    "checkpoint_path = os.path.join(\"checkpoints\", model_name)\n",
    "\n",
    "# Load trained noise model\n",
    "gaussian_noise_std = None\n",
    "noise_model_params = np.load(\n",
    "    \"./data/Convallaria_diaphragm/GMMNoiseModel_convallaria_3_2_calibration.npz\"\n",
    ")\n",
    "noiseModel = GaussianMixtureNoiseModel(params=noise_model_params)\n",
    "\n",
    "# Training specific\n",
    "lr = 3e-4\n",
    "max_epochs = 500\n",
    "steps_per_epoch = 400\n",
    "limit_train_batches = steps_per_epoch * n_grad_batches\n",
    "\n",
    "# VAE specific\n",
    "num_latents = 6\n",
    "z_dims = [32]*int(num_latents)\n",
    "blocks_per_layer = 5\n",
    "batchnorm = True\n",
    "free_bits = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<code>lvae</code>: The traditional Hierarchical DivNoising model which uses the <code>noiseModel</code>, $p(\\text{observation}|\\text{signal})$, to train the approximate posterior, $q(\\text{signal}|\\text{observation})$.<br>\n",
    "<code>direct_denoiser</code>: A deterministic network that is trained by DivNoising to estimate $\\mathbb{E}_{q(\\text{signal}|\\text{observation})}[\\text{signal}]$.<br>\n",
    "<code>backbone</code>: The <code>lvae</code> and <code>direct_denoiser</code> models are trained simulataneously but with their own optimizers. This module handles their co-training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_mean = observation.mean()\n",
    "data_std = observation.std()\n",
    "\n",
    "lvae = LadderVAE(\n",
    "    z_dims=z_dims,\n",
    "    blocks_per_layer=blocks_per_layer,\n",
    "    data_mean=data_mean,\n",
    "    data_std=data_std,\n",
    "    noiseModel=noiseModel,\n",
    "    batchnorm=batchnorm,\n",
    "    free_bits=free_bits,\n",
    "    img_shape=img_shape,\n",
    ")\n",
    "\n",
    "direct_denoiser = UNet(depth=4, start_filters=32)\n",
    "\n",
    "backbone = Backbone(\n",
    "    lvae,\n",
    "    direct_denoiser,\n",
    "    data_mean=data_mean,\n",
    "    data_std=data_std,\n",
    "    gaussian_noise_std=gaussian_noise_std,\n",
    "    n_grad_batches=n_grad_batches,\n",
    "    lr=lr,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train networks<br>\n",
    "Training can be monitored in Tensorboard. Enter `tensorboard --logdir <path-to-this-directory>/checkpoints/convallaria` to run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = TensorBoardLogger(checkpoint_path, name=model_name)\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    default_root_dir=checkpoint_path,\n",
    "    accelerator=\"gpu\" if use_cuda else \"cpu\",\n",
    "    devices=1,\n",
    "    limit_train_batches=limit_train_batches,\n",
    "    max_epochs=max_epochs,\n",
    "    callbacks=[EarlyStopping(patience=30, monitor=\"val/elbo\")],\n",
    "    logger=logger,\n",
    ")\n",
    "\n",
    "trainer.fit(backbone, train_loader, val_loader)\n",
    "trainer.save_checkpoint(os.path.join(checkpoint_path, \"final_model.ckpt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}