guide/transfer_learning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python [conda env:nb]",
      "language": "python",
      "name": "conda-env-nb-py"
    },
    "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.7.6"
    },
    "colab": {
      "name": "transfer_learning.ipynb",
      "provenance": []
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I4y7PvtMtLWG"
      },
      "source": [
        "# Train a neural network with very little data\n",
        "\n",
        "In this notebook, we will use Nobrainer to train a model 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 extraction model can be re-trained for tumor labeling. In this notebook, we will train a brain extraction 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",
        "    - We will binarize the FreeSurfer to get a precise brainmask.\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 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": "markdown",
      "metadata": {
        "id": "XaMNusTGDwAB"
      },
      "source": [
        "In this tutorial we will need to pull in external models using datalad. To support it we first install git-annex using Neurodebian.\n",
        "\n",
        "**Note: Sometimes getting the gpg key can fail. If you notice such an error, simply re-execute the cell.**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qjf_GJfxzKot"
      },
      "source": [
        "!wget -O- http://neuro.debian.net/lists/bionic.us-nh.full | tee /etc/apt/sources.list.d/neurodebian.sources.list \\\n",
        " && export GNUPGHOME=\"$(mktemp -d)\" \\\n",
        " && echo \"disable-ipv6\" >> ${GNUPGHOME}/dirmngr.conf \\\n",
        " && (apt-key adv --homedir $GNUPGHOME --recv-keys --keyserver hkp://pgpkeys.eu 0xA5D32F012649A5A9 \\\n",
        "|| { curl -sSL http://neuro.debian.net/_static/neuro.debian.net.asc | apt-key add -; } ) \\\n",
        " && apt-get update \\\n",
        " && apt-get install git-annex-standalone git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OZqLJWAp1fa1"
      },
      "source": [
        "Let's make sure the correct version of git-annex is installed. This should be at least **8.20210223** or later."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5dUDuJ3g1a4R"
      },
      "source": [
        "!git-annex version"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0LKdY1cctNQf"
      },
      "source": [
        "!pip install --no-cache-dir nilearn datalad datalad-osf nobrainer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7CzMpls0tLWJ"
      },
      "source": [
        "import nobrainer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iOPRjWy2tLWJ"
      },
      "source": [
        "# Get sample features and labels\n",
        "\n",
        "We use 9 pairs of volumes for training and 1 pair of volumes for evaluation. Many more volumes would be required to train a model for any useful purpose."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lz-h4M-4tLWJ"
      },
      "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:]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VHA91ByhtLWK"
      },
      "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",
      "metadata": {
        "id": "cz9AapF1tLWK"
      },
      "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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JP2tPSVbtLWL"
      },
      "source": [
        "!mkdir -p data"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aYhISMZytLWL"
      },
      "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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tDY0txbKtLWL"
      },
      "source": [
        "!ls data"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nOFvDm9jtLWM"
      },
      "source": [
        "# Create Datasets"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lYS-TQG8tLWM"
      },
      "source": [
        "n_classes = 1\n",
        "batch_size = 2\n",
        "volume_shape = (256, 256, 256)\n",
        "block_shape = (128, 128, 128)\n",
        "n_epochs = None\n",
        "augment = False\n",
        "shuffle_buffer_size = 10\n",
        "num_parallel_calls = 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iueawTPBtLWN"
      },
      "source": [
        "dataset_train = nobrainer.dataset.get_dataset(\n",
        "    file_pattern='data/data-train_shard-*.tfrec',\n",
        "    n_classes=n_classes,\n",
        "    batch_size=batch_size,\n",
        "    volume_shape=volume_shape,\n",
        "    block_shape=block_shape,\n",
        "    n_epochs=n_epochs,\n",
        "    augment=augment,\n",
        "    shuffle_buffer_size=shuffle_buffer_size,\n",
        "    num_parallel_calls=num_parallel_calls,\n",
        ")\n",
        "\n",
        "dataset_evaluate = nobrainer.dataset.get_dataset(\n",
        "    file_pattern='data/data-evaluate_shard-*.tfrec',\n",
        "    n_classes=n_classes,\n",
        "    batch_size=batch_size,\n",
        "    volume_shape=volume_shape,\n",
        "    block_shape=block_shape,\n",
        "    n_epochs=1,\n",
        "    augment=False,\n",
        "    shuffle_buffer_size=None,\n",
        "    num_parallel_calls=1,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q2ZN2_QHtLWN"
      },
      "source": [
        "dataset_train"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1o4oPLoDtLWN"
      },
      "source": [
        "dataset_evaluate"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QTOKiUHXtLWO"
      },
      "source": [
        "# Load pre-trained model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OET-DhW8DhYK"
      },
      "source": [
        "Use datalad to retrieve trained models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JmW5o3xyDTuM"
      },
      "source": [
        "!datalad clone https://github.com/neuronets/trained-models && \\\n",
        "  cd trained-models && git-annex enableremote osf-storage && \\\n",
        "  datalad get -s osf-storage neuronets/brainy/0.1.0/brain-extraction-unet-128iso-model.h5"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gk0VjXi4tLWO"
      },
      "source": [
        "import tensorflow as tf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "86t_scOFtLWO"
      },
      "source": [
        "model_path = \"trained-models/neuronets/brainy/0.1.0/brain-extraction-unet-128iso-model.h5\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XACTk31ytLWO"
      },
      "source": [
        "model = tf.keras.models.load_model(model_path, compile=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": false,
        "id": "sr3rOcTjtLWO"
      },
      "source": [
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NKZIaJUstLWP"
      },
      "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 apply regularization and use a relatively small learning rate."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JdCmZZiAtLWP"
      },
      "source": [
        "for layer in model.layers:\n",
        "    layer.kernel_regularizer = tf.keras.regularizers.l2(0.001)\n",
        "\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-05)\n",
        "\n",
        "model.compile(\n",
        "    optimizer=optimizer,\n",
        "    loss=nobrainer.losses.jaccard,\n",
        "    metrics=[nobrainer.metrics.dice],\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rEHeL_V8tLWP"
      },
      "source": [
        "# Train and evaluate model\n",
        "\n",
        "$$\n",
        "steps = \\frac{nBlocks}{volume} * \\frac{nVolumes}{batchSize}\n",
        "$$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XZHwKi3TtLWP"
      },
      "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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yq7hQURstLWQ"
      },
      "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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bhjErOh9m1R5"
      },
      "source": [
        "The following step may take about 10 mins on a standard colab GPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": false,
        "id": "Kj9tO1REtLWQ"
      },
      "source": [
        "model.fit(\n",
        "    dataset_train,\n",
        "    epochs=5,\n",
        "    steps_per_epoch=steps_per_epoch, \n",
        "    validation_data=dataset_evaluate, \n",
        "    validation_steps=validation_steps)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TGLd_JMftLWQ"
      },
      "source": [
        "# Predict natively without TFRecords"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EvAOMTUatLWR"
      },
      "source": [
        "from nobrainer.volume import standardize\n",
        "\n",
        "image_path = evaluate_paths[0][0]\n",
        "out = nobrainer.prediction.predict_from_filepath(image_path, \n",
        "                                                 model,\n",
        "                                                 block_shape = block_shape,\n",
        "                                                 batch_size = batch_size,\n",
        "                                                 normalizer = standardize,\n",
        "                                                 )\n",
        "out.shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b3PJdwWmt3ll"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "from nilearn import plotting\n",
        "fig = plt.figure(figsize=(12, 6))\n",
        "plotting.plot_roi(out, bg_img=image_path, alpha=0.4, figure=fig, cmap=plt.cm.gist_rainbow)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v97Pv7gHzC4J"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}