03_COSDD/exercise-1-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise 1. Training COSDD<br>\n",
    "In this section, we will train a COSDD model to remove row correlated and signal-dependent imaging noise. \n",
    "You will load noisy data and examine the noise for spatial correlation, then initialise a model and monitor its training.\n",
    "\n",
    "COSDD is a Ladder VAE with an autoregressive decoder, a type of deep generative model. Deep generative models are trained with the objective of capturing all the structures and characteristics present in a dataset, i.e., modelling the dataset. In our case the dataset will be a collection of noisy microscopy images. \n",
    "\n",
    "When COSDD is trained to model noisy images, it exploits differences between the structure of imaging noise and the structure of the clean signal to separate them, capturing each with different components of the model. Specifically, the noise will be captured by the autoregressive decoder and the signal will be captured by the VAE's latent variables. We can then feed an image into the model and sample a latent variable, which will describe the image's clean signal content. This latent variable is then fed through a second network, which was trained alongside the main VAE, to reveal an estimate of the denoised image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "import tifffile\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.callbacks import EarlyStopping\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import utils\n",
    "from models.lvae import LadderVAE\n",
    "from models.pixelcnn import PixelCNN\n",
    "from models.s_decoder import SDecoder\n",
    "from models.unet import UNet\n",
    "from hub import Hub\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_cuda = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1. Load data\n",
    "\n",
    "In this example we will be using the Penicillium dataset, provided by: <br>\n",
    "Eom, Minho, et al. \"Statistically unbiased prediction enables accurate denoising of voltage imaging data.\" Nature Methods 20.10 (2023): 1581-1592."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.1.\n",
    "\n",
    "The low signal-to-noise ratio data that we will be using in this exercise is stored as a tiff file at `/group/dl4miacourse/image_regression/penicillium/penicillium_low_snr.tif`. \n",
    "1. Use the function `tifffile.imread` to load the data as a numpy array.\n",
    "2. Then use np.newaxis to add a channel axis. *Hint* We're treating the data as a stack of 2D images, so the channel axis should be the second dimension (dimension 1 if we start counting from zero).\n",
    "3. Finally, use `torch.from_numpy` to convert it into a pytorch tensor and set the data type as float32.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lowsnr_path = \"/group/dl4miacourse/image_regression/penicillium/penicillium_low_snr.tif\"\n",
    "\n",
    "low_snr = ...\n",
    "low_snr = ...\n",
    "low_snr = ...\n",
    "\n",
    "print(f\"Noisy data dimensions: {low_snr.shape}\")\n",
    "correct_dims = [1001, 1, 1024, 1024]\n",
    "msg = f\"Data dimensions should be {correct_dims} but are {[*low_snr.shape]}\"\n",
    "assert [*low_snr.shape] == correct_dims, msg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "COSDD can handle 1-, 2- and 3-dimensional data, as long as it's loaded as a PyTorch tensor with a batch and channel dimension. For 1D data, it should have dimensions [Number of images, Channels, X], for 2D data: [Number of images, Channels, Y, X] and for 3D: [Number of images, Channels, Z, Y, X]. This applies even if the data has only one channel. <br>\n",
    "In this example, we're treating a single 3D image as a stack of 2D images. This is because COSDD is much more slow and memory hungry than N2V, and we want the model to train more in the time available for the exercise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2. Examine spatial correlation of the noise\n",
    "\n",
    "This denoiser can remove noise that is correlated along rows or columns of pixels (or not spatially correlated at all). Row or column-correlated noise means that knowing the noise value in one pixel can help us guess the noise value in other pixels in the same row (or column), preventing us from using Noise2Void. Noise correlated along rows of pixels is common in scanning imaging techniques such as scanning confocal microscopy and scanning electron microscopy.\n",
    "\n",
    "To remove this type of noise, we need to tell the denoiser whether the noise is correlated along rows or along columns. We can establish this by estimating the spatial autocorrelation of the noise. In the context of images, spatial autocorrelation is a measure of what the value of a pixel in one location tells us about the value of pixels in other locations. It is a scale from -1 to +1 and is measured between all pairs of pixel locations. A negative autocorrelation means that the higher the intensity of the one pixel, the lower we should expect the intensity of the other to be. A positive autocorrelation means that the higher the intensity of the one pixel, the higher we should expect the intensity of the other to be. Zero autocorrelation means that the intensity of the one pixel being high or low tells us nothing about whether the intensity of the other will be high or low. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.2.\n",
    "\n",
    "To estimate the spatial correlation of the noise, we need a sample of pure noise. This can be a patch of `low_snr` with no signal. Adjust the values for `img_idx`, `top`, `bottom`, `left` and `right` to explore slices of the data and identify a suitable dark patch. When decided, set the `dark_patch` in the following cell and pass it as an argument to `utils.autocorrelation`, then plot the result. \n",
    "\n",
    "*Hint: The bigger the dark patch, the more accurate our estimate of the spatial autocorrelation will be.*\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vmin = np.percentile(low_snr[::100], 0.1)\n",
    "vmax = np.percentile(low_snr[::100], 99.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Explore slices of the data here\n",
    "img_idx = 500\n",
    "top = 0\n",
    "bottom = 1024\n",
    "left = 0\n",
    "right = 1024\n",
    "\n",
    "crop = (img_idx, 0, slice(top, bottom), slice(left, right))\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(low_snr[crop], vmin=vmin, vmax=vmax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Define the crop of the dark image patch here\n",
    "dark_img_idx = ...\n",
    "dark_top = ...\n",
    "dark_bottom = ...\n",
    "dark_left = ...\n",
    "dark_right = ...\n",
    "\n",
    "dark_crop = (dark_img_idx, 0, slice(dark_top, dark_bottom), slice(dark_left, dark_right))\n",
    "dark_patch = low_snr[dark_crop]\n",
    "\n",
    "noise_ac = utils.autocorrelation(dark_patch, max_lag=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the autocorrelation\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(noise_ac, cmap=\"seismic\", vmin=-1, vmax=1)\n",
    "plt.colorbar()\n",
    "plt.title(\"Autocorrelation of the noise\")\n",
    "plt.xlabel(\"Horizontal lag\")\n",
    "plt.ylabel(\"Vertical lag\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a two dimensional plot describing the linear relationship between pixel values at different locations. The colour of a square is the value of the autocorrelation between pixels separated by the horizontal and vertical lag. For example, the colour of the square at a horizontal lag of 5 and a vertical lag of 10 tells us that if we looked at the value any one pixel in the image, then counted 5 pixels to the right and 10 pixels down, what value should we expect the pixel we land on to be. Correlation is symetric, so we can expect the same if we counted left or up.\n",
    "\n",
    "In this plot, all of the squares should be white, except for the top row. The autocorrelation of the square at (0, 0) will always be 1.0, as a pixel's value will always be perfectly correlated with itself. We define this type of noise as correlated along the x axis.\n",
    "\n",
    "To remove this type of noise, the autoregressive decoder of our VAE must have a receptive field spanning the x axis.\n",
    "Note that if the data contained spatially uncorrelated noise, we can still remove it, as the decoder's receptive field will become redundant."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3. Create training and validation dataloaders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data will be fed to the model by two dataloaders, `train_loader` and `val_loader`, for the training and validation set respectively. <br>\n",
    "In this example, 90% of images will be used for training and the remaining 10% for validation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`real_batch_size` (int) Number of images passed through the network at a time. <br>\n",
    "`n_grad_batches` (int) Number of batches to pass through the network before updating parameters.<br>\n",
    "`crop_size` (tuple(int)): The size of randomly cropped patches. Should be less than the dimensions of your images.<br>\n",
    "`train_split` (0 < float < 1): Fraction of images to be used in the training set, with the remainder used for the validation set.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "real_batch_size = 4\n",
    "n_grad_batches = 4\n",
    "print(f\"Effective batch size: {real_batch_size * n_grad_batches}\")\n",
    "crop_size = (256, 256)\n",
    "train_split = 0.9\n",
    "\n",
    "n_iters = np.prod(low_snr.shape[2:]) // np.prod(crop_size)\n",
    "transform = utils.RandomCrop(crop_size)\n",
    "\n",
    "low_snr = low_snr[torch.randperm(len(low_snr))]\n",
    "train_set = low_snr[: int(len(low_snr) * train_split)]\n",
    "val_set = low_snr[int(len(low_snr) * train_split) :]\n",
    "\n",
    "train_set = utils.TrainDataset(train_set, n_iters=n_iters, transform=transform)\n",
    "val_set = utils.TrainDataset(val_set, n_iters=n_iters, transform=transform)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_set, batch_size=real_batch_size, shuffle=True, pin_memory=True, num_workers=7,\n",
    ")\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    val_set, batch_size=real_batch_size, shuffle=False, pin_memory=True, num_workers=7,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4. Create the model\n",
    "\n",
    "The model we will train to denoise consists of four modules, with forth being the optional Direct Denoiser which we can train if we want to speed up inference. Each module is listed below with an explanation of their hyperparameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`lvae` The ladder variational autoencoder that will output latent variables.<br>\n",
    "* `s_code_channels` (int): Number of channels in outputted latent variable.\n",
    "* `n_layers` (int): Number of levels in the ladder vae.\n",
    "* `z_dims` (list(int)): List with the numer of latent space dimensions at each level of the hierarchy. List starts from the input/output level and works down.\n",
    "* `downsampling` (list(int)): Binary list of whether to downsample at each level of the hierarchy. 1 for do and 0 for don't.\n",
    "\n",
    "`ar_decoder` The autoregressive decoder that will decode latent variables into a distribution over the input.<br>\n",
    "* `kernel_size` (int): Length of 1D convolutional kernels.\n",
    "* `noise_direction` (str): Axis along which noise is correlated: `\"x\"`, `\"y\"` or `\"z\"`. This needs to match the orientation of the noise structures we revealed in the autocorrelation plot in Part 2.\n",
    "* `n_filters` (int): Number of feature channels.\n",
    "* `n_gaussians` (int): Number of components in Gaussian mixture used to model data.\n",
    "\n",
    "`s_decoder` A decoder that will map the latent variables into image space, giving us a denoised image. <br>\n",
    "* `n_filters` (int): The number of feature channels.<br>\n",
    "\n",
    "`direct_denoiser` The U-Net that can optionally be trained to predict the MMSE or MMAE of the denoised images. This will slow training slightly but massively speed up inference and is worthwile if you have an inference dataset in the gigabytes. See [this paper](https://arxiv.org/abs/2310.18116). Enable or disable the direct denoiser by setting `use_direct_denoiser` to `True` or `False`.\n",
    "* `n_filters` (int): Feature channels at each level of UNet. Defaults to `s_code_channel`.\n",
    "* `n_layers` (int): Number of levels in the UNet. Defaults to the number of levels in the `LadderVAE`.\n",
    "* `downsampling` (list(int)): Binary list of whether to downsample at each level of the hierarchy. 1 for do and 0 for don't. Also defaults to match the `LadderVAE`.\n",
    "* `loss_fn` (str): Whether to use `\"L1\"` or `\"L2\"` loss function to predict either the mean or pixel-wise median of denoised images respectively.\n",
    "\n",
    "`hub` The hub that will unify and train the above modules.\n",
    "* `n_grad_batches` (int): Number of batches to accumulate gradients for before updating weights of all models. If the real batch or random crop size has been reduced to lower memory consumption, increase this value for the effective batch size to stay the same.\n",
    "* `checkpointed` (bool): Whether to use activation checkpointing during training. This reduces memory consumption but increases training time. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.3.\n",
    "\n",
    "Most hyperparameters have been set to recommended values for a small sized model. The two that have been left blank are `noise_direction` under the `ar_decoder`, and `use_direct_denoiser`. Use the above description of what each hyperparameter means to determine the best value for each of these.\n",
    "\n",
    "*Hint: enabling the Direct Denoiser will give us additional results to look at in the next notebook.*\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_code_channels = 64\n",
    "\n",
    "n_layers = 6\n",
    "z_dims = [s_code_channels // 2] * n_layers\n",
    "downsampling = [1] * n_layers\n",
    "lvae = LadderVAE(\n",
    "    colour_channels=low_snr.shape[1],\n",
    "    img_size=crop_size,\n",
    "    s_code_channels=s_code_channels,\n",
    "    n_filters=s_code_channels,\n",
    "    z_dims=z_dims,\n",
    "    downsampling=downsampling,\n",
    ")\n",
    "\n",
    "ar_decoder = PixelCNN(\n",
    "    colour_channels=low_snr.shape[1],\n",
    "    s_code_channels=s_code_channels,\n",
    "    kernel_size=5,\n",
    "    noise_direction=...,  ### Insert a value here\n",
    "    n_filters=64,\n",
    "    n_layers=4,\n",
    "    n_gaussians=8,\n",
    ")\n",
    "\n",
    "s_decoder = SDecoder(\n",
    "    colour_channels=low_snr.shape[1],\n",
    "    s_code_channels=s_code_channels,\n",
    "    n_filters=s_code_channels,\n",
    ")\n",
    "\n",
    "use_direct_denoiser = ...  ### Insert a value here\n",
    "if use_direct_denoiser:\n",
    "    direct_denoiser = UNet(\n",
    "        colour_channels=low_snr.shape[1],\n",
    "        n_filters=s_code_channels,\n",
    "        n_layers=n_layers,\n",
    "        downsampling=downsampling,\n",
    "        loss_fn=\"L2\",\n",
    "    )\n",
    "else:\n",
    "    direct_denoiser = None\n",
    "\n",
    "hub = Hub(\n",
    "    vae=lvae,\n",
    "    ar_decoder=ar_decoder,\n",
    "    s_decoder=s_decoder,\n",
    "    direct_denoiser=direct_denoiser,\n",
    "    data_mean=low_snr.mean(),\n",
    "    data_std=low_snr.std(),\n",
    "    n_grad_batches=n_grad_batches,\n",
    "    checkpointed=True,\n",
    ")\n",
    "\n",
    "# Show the model the dimensionality of the data to initialise its parameters\n",
    "a = next(iter(train_loader))\n",
    "with torch.no_grad():\n",
    "    _ = hub(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.5. Train the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.4.\n",
    "\n",
    "We're going to use Tensorboard to monitor training metrics. Run the two cells below to open tensorboard in the notebook. If it doesn't work, follow the steps below to open it in a browser.\n",
    "\n",
    "1. Open a terminal.\n",
    "2. Enter `mamba activate 02_regression` to activate an environment with Tensorboard installed.\n",
    "3. Enter `tensorboard --logdir 02_regression/03_COSDD/checkpoints`\n",
    "4. Finally, open a browser and enter localhost:6006 in the address bar.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tensorboard metrics\n",
    "\n",
    "In the SCALARS tab, there will be 4 metrics to track (5 if direct denoiser is enabled). These are:<br>\n",
    "1. `kl_loss` The Kullback-Leibler divergence between the VAE's approximate posterior and its prior. This can be thought of as a measure of how much information about the input image is going into the VAE's latent variables. We want information about the input's underlying clean signal to go into the latent variables, so this metric shouldn't go all the way to zero. Instead, it can typically go either up or down during training before plateauing.<br>\n",
    "2. `reconstruction_loss` The negative log-likelihood of the AR decoder's predicted distribution given the input data. This is how accurately the AR decoder is able to predict the input. This value can go below zero and should decrease throughout training before plateauing.<br>\n",
    "3. `elbo` The Evidence Lower Bound, which is the total loss of the main VAE. This is the sum of the kl and reconstruction loss and should decrease throughout training before plateauing.<br>\n",
    "4. `sd_loss` The mean squared error between the noisy image and the image predicted by the signal decoder. This metric should steadily decrease towards zero without ever reaching it. Sometimes the loss will not go down for the first few epochs because its input (produced by the VAE) is rapidly changing. This is ok and the loss should start to decrease when the VAE stabilises. <br>\n",
    "5. `dd_loss` The mean squared error between the output of the direct denoiser and the clean images predicted by the signal decoder. This will only be present if `use_direct_denoiser` is set to `True`. The metric should steadily decrease towards zero without ever reaching it, but may be unstable at the start of training as its targets (produced by the signal decoder) are rapidly changing.\n",
    "\n",
    "There will also be an IMAGES tab. This shows noisy input images from the validation set and some outputs. These will be two randomly sampled denoised images (sample 1 and sample 2), the average of ten denoised images (mmse) and if the direct denoiser is enabled, its output (direct estimate).\n",
    "\n",
    "If noise has not been fully removed from the output images, try increasing `n_gaussians` argument of the AR decoder. This will give it more flexibility to model complex noise characteristics. However, setting the value too high can lead to unstable training. Typically, values from 3 to 5 work best.\n",
    "\n",
    "Note that the trainer is set to train for only 20 minutes in this example. Remove the line with `max_time` to train fully."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%tensorboard --logdir checkpoints/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.5.\n",
    "\n",
    "Now the model is ready to start training. Give the model a sensible name by setting `model_name` to a string then run the following cells.\n",
    "\n",
    "The `max_time` parameter in the cell below means we'll only train the model for 20 minutes, just to get idea of what to expect. In the future, to remove the time restriction, the `max_time` parameter can be set to `None`.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`model_name` (str): Should be set to something appropriate so that the trained parameters can be used later for inference.<br>\n",
    "`max_epochs` (int): The number of training epochs.<br>\n",
    "`patience` (int): If the validation loss has plateaued for this many epochs, training will stop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = ...  ### Insert a value here\n",
    "checkpoint_path = os.path.join(\"checkpoints\", model_name)\n",
    "logger = TensorBoardLogger(checkpoint_path)\n",
    "\n",
    "max_epochs = 1000\n",
    "max_time = \"00:00:20:00\"\n",
    "patience = 100\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    logger=logger,\n",
    "    accelerator=\"gpu\" if use_cuda else \"cpu\",\n",
    "    devices=1,\n",
    "    max_epochs=max_epochs,\n",
    "    max_time=max_time,  # Remove this time limit to train the model fully\n",
    "    log_every_n_steps=len(train_set) // (4 * real_batch_size),\n",
    "    callbacks=[EarlyStopping(patience=patience, monitor=\"val/elbo\")],\n",
    "    precision=\"bf16-mixed\" if use_cuda else 32,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.fit(hub, train_loader, val_loader)\n",
    "trainer.save_checkpoint(os.path.join(checkpoint_path, \"final_model.ckpt\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\">\n",
    "\n",
    "## Checkpoint 1\n",
    "We've now trained a COSDD model to denoise our data. Continue to the next notebook, exercise-2-prediction.ipynb, to use it to get some results.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}