guide/transfer_learning-bayesian.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kNguCLDMCDxD"
      },
      "source": [
        "# Train a neural network with very little data\n",
        "\n",
        "In this notebook, we will use Nobrainer to train a Bayesian neural network with limited data. We will start off with a pre-trained model. You can find available pre-trained Nobrainer models at https://github.com/neuronets/nobrainer-models.\n",
        "\n",
        "The pre-trained models can be used to train models for the same task as they were trained for or to transfer learn a new task. For instance, a pre-trained brain labelling model can be re-trained for tumor labeling. In this notebook, we will train a brain labeling model, but keep in mind that you can retrain these models for many 3D semantic segmentation tasks.\n",
        "\n",
        "In the following cells, we will:\n",
        "\n",
        "1. Get sample T1-weighted MR scans as features and FreeSurfer segmentations as labels.\n",
        "2. Convert the data to TFRecords format.\n",
        "3. Create two Datasets of the features and labels.\n",
        "    - One dataset will be for training and the other will be for evaluation.\n",
        "4. Load a pre-trained 3D semantic segmentation Bayesian model.\n",
        "5. Choose a loss function and metrics to use.\n",
        "6. Train on part of the data.\n",
        "7. Evaluate on the rest of the data.\n",
        "\n",
        "## Google Colaboratory\n",
        "\n",
        "If you are using Colab, please switch your runtime to GPU. To do this, select `Runtime > Change runtime type` in the top menu. Then select GPU under `Hardware accelerator`. A GPU greatly speeds up training."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install nobrainer"
      ],
      "metadata": {
        "id": "JQ-P_PzHCJb6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O8zmcBiACDxI"
      },
      "outputs": [],
      "source": [
        "import nobrainer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kfHn42x_CDxJ"
      },
      "source": [
        "# Get sample features and labels\n",
        "\n",
        "We use 9 pairs of volumes for training and 1 pair of volumes for evaulation. Many more volumes would be required to train a model for any useful purpose."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WnBCrxRUCDxK"
      },
      "outputs": [],
      "source": [
        "csv_of_filepaths = nobrainer.utils.get_data()\n",
        "filepaths = nobrainer.io.read_csv(csv_of_filepaths)\n",
        "\n",
        "train_paths = filepaths[:9]\n",
        "evaluate_paths = filepaths[9:]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6dsTOMVhCDxL"
      },
      "source": [
        "# Convert medical images to TFRecords\n",
        "\n",
        "Remember how many full volumes are in the TFRecords files. This will be necessary to know how many steps are in on training epoch. The default training method needs to know this number, because Datasets don't always know how many items they contain."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lHefZL7XCDxL"
      },
      "outputs": [],
      "source": [
        "# Verify that all volumes have the same shape and that labels are integer-ish.\n",
        "\n",
        "invalid = nobrainer.io.verify_features_labels(train_paths, num_parallel_calls=2)\n",
        "assert not invalid\n",
        "\n",
        "invalid = nobrainer.io.verify_features_labels(evaluate_paths)\n",
        "assert not invalid"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FMH1QG1ACDxL"
      },
      "outputs": [],
      "source": [
        "!mkdir -p data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x1AdoUfLCDxM"
      },
      "outputs": [],
      "source": [
        "# Convert training and evaluation data to TFRecords.\n",
        "\n",
        "nobrainer.tfrecord.write(\n",
        "    features_labels=train_paths,\n",
        "    filename_template='data/data-train_shard-{shard:03d}.tfrec',\n",
        "    examples_per_shard=3)\n",
        "\n",
        "nobrainer.tfrecord.write(\n",
        "    features_labels=evaluate_paths,\n",
        "    filename_template='data/data-evaluate_shard-{shard:03d}.tfrec',\n",
        "    examples_per_shard=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Q5FvIGUCDxN"
      },
      "outputs": [],
      "source": [
        "!ls data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-4hBG14lCDxO"
      },
      "source": [
        "# Create Datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pkOvJBT8CDxO"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "n_classes = 50\n",
        "batch_size = 2\n",
        "volume_shape = (256, 256, 256)\n",
        "block_shape = (32, 32, 32)\n",
        "n_epochs = 2\n",
        "\n",
        "def _to_blocks(x, y):\n",
        "    \"\"\"Separate `x` into blocks and repeat `y` by number of blocks.\"\"\"\n",
        "    x = nobrainer.volume.to_blocks(x, block_shape)\n",
        "    y = nobrainer.volume.to_blocks(y, block_shape)\n",
        "    return (x, y)\n",
        "\n",
        "def process_dataset(dset):\n",
        "    # Standard score the features.\n",
        "    dset = dset.map(lambda x, y: (nobrainer.volume.standardize(x), y))\n",
        "    # Separate features into blocks.\n",
        "    dset = dset.map(_to_blocks)\n",
        "    # This step is necessary because separating into blocks adds a dimension.\n",
        "    dset = dset.unbatch()\n",
        "    # Add a grayscale channel to the features.\n",
        "    dset = dset.map(lambda x, y: (tf.expand_dims(x, -1), y))\n",
        "    # Batch features and labels.\n",
        "    dset = dset.batch(batch_size, drop_remainder=True)\n",
        "    dset = dset.repeat(n_epochs)\n",
        "    return dset\n",
        "\n",
        "# Create a `tf.data.Dataset` instance.\n",
        "dataset_train = nobrainer.dataset.tfrecord_dataset(\n",
        "    file_pattern=\"data/data-train_shard-*.tfrec\",\n",
        "    volume_shape=volume_shape,\n",
        "    shuffle=True,\n",
        "    scalar_label=False,\n",
        "    num_parallel_calls=2\n",
        ")\n",
        "dataset_train = process_dataset(dataset_train)\n",
        "\n",
        "# Create a `tf.data.Dataset` instance.\n",
        "dataset_evaluate = nobrainer.dataset.tfrecord_dataset(\n",
        "    file_pattern=\"data/data-evaluate_shard-*.tfrec\",\n",
        "    volume_shape=volume_shape,\n",
        "    shuffle=False,\n",
        "    scalar_label=False,\n",
        "    num_parallel_calls=2\n",
        ")\n",
        "dataset_evaluate = process_dataset(dataset_evaluate)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y4pabFvcCDxP"
      },
      "outputs": [],
      "source": [
        "dataset_train"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zBKDgf93CDxP"
      },
      "outputs": [],
      "source": [
        "dataset_evaluate"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XScYNuNoCDxP"
      },
      "source": [
        "# Load pre-trained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "STBCfp_NCDxQ"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "from nobrainer.models.bayesian import variational_meshnet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dyz8SXmCCDxQ"
      },
      "outputs": [],
      "source": [
        "model = variational_meshnet(\n",
        "    n_classes=50, \n",
        "    input_shape=(32, 32, 32, 1),\n",
        "    filters=96, \n",
        "    dropout=\"concrete\", \n",
        "    receptive_field=37, \n",
        "    is_monte_carlo=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "69XPBEAzCDxR"
      },
      "outputs": [],
      "source": [
        "weights_path = tf.keras.utils.get_file(\n",
        "    fname=\"nobrainer_spikeslab_32iso_weights.h5\",\n",
        "    origin=\"https://dl.dropbox.com/s/rojjoio9jyyfejy/nobrainer_spikeslab_32iso_weights.h5\")\n",
        "\n",
        "model.load_weights(weights_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": false,
        "id": "_Zn_WWLgCDxR"
      },
      "outputs": [],
      "source": [
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sr_SsvDkCDxR"
      },
      "source": [
        "# Considerations for transfer learning\n",
        "\n",
        "Training a neural network changes the model's weights. A pre-trained network has learned weights for a task, and we do not want to forget these weights during training. In other words, we do not want to ruin the pre-trained weights when using our new data. To avoid dramatic changes in the learnable parameters, we can use a relatively small learning rate.\n",
        "\n",
        "We also want to optimize the evidence lower bound ([ELBO](https://en.wikipedia.org/wiki/Evidence_lower_bound)). Specifically, we will minimize $-ELBO$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QjIgEvuPCDxS"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "loss_fn = nobrainer.losses.ELBO(model=model, num_examples=np.prod(block_shape))\n",
        "\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-06)\n",
        "\n",
        "model.compile(\n",
        "    optimizer=optimizer,\n",
        "    loss=loss_fn,\n",
        "    # See https://github.com/tensorflow/probability/issues/519\n",
        "    experimental_run_tf_function=False\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rmGh44FTCDxS"
      },
      "source": [
        "# Train and evaluate model\n",
        "\n",
        "$$\n",
        "steps = \\frac{nBlocks}{volume} * \\frac{nVolumes}{batchSize}\n",
        "$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0ZeMaDNaCDxS"
      },
      "outputs": [],
      "source": [
        "steps_per_epoch = nobrainer.dataset.get_steps_per_epoch(\n",
        "    n_volumes=len(train_paths),\n",
        "    volume_shape=volume_shape,\n",
        "    block_shape=block_shape,\n",
        "    batch_size=batch_size)\n",
        "\n",
        "steps_per_epoch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GxFb0nKaCDxS"
      },
      "outputs": [],
      "source": [
        "validation_steps = nobrainer.dataset.get_steps_per_epoch(\n",
        "    n_volumes=len(evaluate_paths),\n",
        "    volume_shape=volume_shape,\n",
        "    block_shape=block_shape,\n",
        "    batch_size=batch_size)\n",
        "\n",
        "validation_steps"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": false,
        "id": "5IRNkG-bCDxT"
      },
      "outputs": [],
      "source": [
        "model.fit(\n",
        "    dataset_train,\n",
        "    epochs=n_epochs,\n",
        "    steps_per_epoch=steps_per_epoch, \n",
        "    validation_data=dataset_evaluate, \n",
        "    validation_steps=validation_steps)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "Z23HRhuLDRjS"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.6.9"
    },
    "colab": {
      "name": "transfer_learning-bayesian.ipynb",
      "provenance": []
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}