brain_age_deep_starting_kit.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Brain age regression with deep learning\n",
    "\n",
    "Predict age from brain grey matter (regression) using Deep Learning (DL).\n",
    "Aging is associated with grey matter (GM) atrophy. Each year, an adult lose\n",
    "0.1% of GM. We will try to learn a predictor of the chronological age (true age)\n",
    "using GM measurements on a population of healthy control participants.\n",
    "\n",
    "Such a predictor provides the expected **brain age** of a subject. Deviation from\n",
    "this expected brain age indicates acceleration or slowdown of the aging process\n",
    "which may be associated with a pathological neurobiological process or protective factor of aging.\n",
    "\n",
    "## Dataset\n",
    "\n",
    "There are 357 samples in the training set and 90 samples in the test set.\n",
    "\n",
    "### Input data\n",
    "\n",
    "Voxel-based_morphometry [VBM](https://en.wikipedia.org/wiki/Voxel-based_morphometry)\n",
    "using [cat12](http://www.neuro.uni-jena.de/cat/) software which provides:\n",
    "\n",
    "- Regions Of Interest (`rois`) of Grey Matter (GM) scaled for the Total\n",
    "  Intracranial Volume (TIV): `[train|test]_rois.csv` 284 features.\n",
    "\n",
    "- VBM GM 3D maps or images (`vbm3d`) of [voxels](https://en.wikipedia.org/wiki/Voxel) in the\n",
    "  [MNI](https://en.wikipedia.org/wiki/Talairach_coordinates) space:\n",
    "  `[train|test]_vbm.npz` contains 3D images of shapes (121, 145, 121).\n",
    "  This npz contains also the 3D mask and the affine transformation to MNI\n",
    "  referential. Masking the brain provides *flat* 331 695 input features (masked voxels)\n",
    "  for each participant.\n",
    "\n",
    "By default `problem.get_[train|test]_data()` return the concatenation of 284 ROIs of\n",
    "Grey Matter (GM) features with 331 695 features (voxels) within a brain mask.\n",
    "Those two blocks are higly redundant.\n",
    "To select only `rois` features do:\n",
    "\n",
    "```\n",
    "X[:, :284]\n",
    "```\n",
    "\n",
    "To select only `vbm` features do:\n",
    "\n",
    "```\n",
    "X[:, 284:]\n",
    "```\n",
    "\n",
    "### Target\n",
    "\n",
    "The target can be found in `[test|train]_participants.csv` files, selecting the\n",
    "`age` column for regression problem.\n",
    "\n",
    "## Evaluation metrics\n",
    "\n",
    "[sklearn metrics](https://scikit-learn.org/stable/modules/model_evaluation.html)\n",
    "\n",
    "The main Evaluation metrics is the Root-mean-square deviation\n",
    "[RMSE](https://en.wikipedia.org/wiki/Root-mean-square_deviation). We will also\n",
    "look at the R-squared\n",
    "[R2](https://en.wikipedia.org/wiki/Coefficient_of_determination).\n",
    "\n",
    "\n",
    "## Links\n",
    "\n",
    "- [RAMP-workflow’s documentation](https://paris-saclay-cds.github.io/ramp-workflow)\n",
    "- [RAMP-workflow’s github](https://github.com/paris-saclay-cds/ramp-workflow)\n",
    "- [RAMP Kits](https://github.com/ramp-kits)\n",
    "\n",
    "## Installation\n",
    "\n",
    "This starting kit requires Python 3 and the following dependencies:\n",
    "\n",
    "* `numpy`\n",
    "* `scipy`\n",
    "* `pandas`\n",
    "* `scikit-learn`\n",
    "* `matplolib`\n",
    "* `seaborn`\n",
    "* `jupyter`\n",
    "* `torch`\n",
    "* `ramp-workflow`\n",
    "\n",
    "You can install the dependencies with the following command-line:\n",
    "\n",
    "```\n",
    "pip install -U -r requirements.txt\n",
    "```\n",
    "\n",
    "If you are using conda, we provide an environment.yml file for similar usage.\n",
    "\n",
    "```\n",
    "conda env create -n ramp-brainage -f environment.yml\n",
    "```\n",
    "\n",
    "Then, you can activate/desactivate the conda environment using:\n",
    "\n",
    "```\n",
    "conda activate ramp-brainage\n",
    "conda deactivate\n",
    "```\n",
    "\n",
    "## Getting started\n",
    "\n",
    "1. Download the data\n",
    "\n",
    "```\n",
    "python download_data.py\n",
    "```\n",
    "\n",
    "The train/test data will be available in the `data` directory.\n",
    "\n",
    "2. Execute the jupyter notebook\n",
    "\n",
    "```\n",
    "jupyter notebook brain_age_deep_starting_kit.ipynb\n",
    "```\n",
    "\n",
    "Play with this notebook to create your new model.\n",
    "\n",
    "3. Test Submission\n",
    "\n",
    "The submissions need to be located in the `submissions` folder.\n",
    "For instance to create a `linear_regression_rois` submission, start by\n",
    "copying the starting kit\n",
    "\n",
    "```\n",
    "cp -r submissions/starting_kit submissions/linear_regression_rois`.\n",
    "```\n",
    " \n",
    "Tune the estimator in the`submissions/linear_regression_rois/estimator.py` file.\n",
    "This file must contain a function `get_estimator()` that returns a scikit learn Pipeline.\n",
    "\n",
    "Then, test your submission locally:\n",
    "\n",
    "```\n",
    "ramp-test --submission linear_regression_rois\n",
    "```\n",
    "\n",
    "Note that the weights of the model are expected in a file called `weights.pth` located in your submission folder.\n",
    "\n",
    "4. Submission on RAMP.studio\n",
    "\n",
    "Connect to your RAMP.studio account, select the `brain_age_deep` event, and submit your estimator in the\n",
    "sandbox section.\n",
    "Note that no training will be performed on the server side.\n",
    "You need to join the weights of the model in a file called `weights.pth`.\n",
    "\n",
    "## The MLP estimator in details\n",
    "\n",
    "Let's play with the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import time\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import sklearn.metrics as metrics\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.model_selection import cross_validate\n",
    "from sklearn.base import BaseEstimator, TransformerMixin\n",
    "from joblib.externals.loky.backend.context import get_context\n",
    "import problem\n",
    "\n",
    "X_train, y_train = problem.get_train_data()\n",
    "X_test, y_test = problem.get_test_data()\n",
    "assert X_train.shape[1] == 284 + 331695"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define here some utility functions to compute scores:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cv_train_test_scores(rmse_cv_test, rmse_cv_train, r2_cv_test, r2_cv_train,\n",
    "                         y_train, y_pred_train, y_test, y_pred_test):\n",
    "    \"\"\" Compute CV score, train and test score from a cv grid search model.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    rmse_cv_test : array\n",
    "        Test rmse across CV folds.\n",
    "    rmse_cv_train : array\n",
    "        Train rmse across CV folds.\n",
    "    r2_cv_test : array\n",
    "        Test R2 across CV folds.\n",
    "    r2_cv_train : array\n",
    "        Train R2 across CV folds.\n",
    "    y_train : array\n",
    "        True train values.\n",
    "    y_pred_train : array\n",
    "        Predicted train values.\n",
    "    y_test : array\n",
    "        True test values.\n",
    "    y_pred_test : array\n",
    "        Predicted test values.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    info : TYPE\n",
    "        DataFrame(r2_cv, r2_train, mae_train, mse_train).\n",
    "    \"\"\"\n",
    "    # CV scores\n",
    "    rmse_cv_test_mean, rmse_cv_test_sd = np.mean(rmse_cv_test), np.std(rmse_cv_test)\n",
    "    rmse_cv_train_mean, rmse_cv_train_sd = np.mean(rmse_cv_train), np.std(rmse_cv_train)\n",
    "\n",
    "    r2_cv_test_mean, r2_cv_test_sd = np.mean(r2_cv_test), np.std(r2_cv_test)\n",
    "    r2_cv_train_mean, r2_cv_train_sd = np.mean(r2_cv_train), np.std(r2_cv_train)\n",
    "\n",
    "    # Test scores\n",
    "    rmse_test = np.sqrt(metrics.mean_squared_error(y_test, y_pred_test))\n",
    "    r2_test = metrics.r2_score(y_test, y_pred_test)\n",
    "\n",
    "    # Train scores\n",
    "    rmse_train = np.sqrt(metrics.mean_squared_error(y_train, y_pred_train))\n",
    "    r2_train = metrics.r2_score(y_train, y_pred_train)\n",
    "\n",
    "    scores = pd.DataFrame(\n",
    "        [[rmse_cv_test_mean, rmse_cv_test_sd, rmse_cv_train_mean, rmse_cv_train_sd,\n",
    "          r2_cv_test_mean, rmse_cv_test_sd, r2_cv_train_mean, r2_cv_train_sd,\n",
    "          rmse_test, r2_test, rmse_train, r2_train]],\n",
    "        columns=('rmse_cv_test_mean', 'rmse_cv_test_sd', 'rmse_cv_train_mean', 'rmse_cv_train_sd',\n",
    "                 'r2_cv_test_mean', 'rmse_cv_test_sd', 'r2_cv_train_mean', 'r2_cv_train_sd',\n",
    "                 'rmse_test', 'r2_test', 'rmse_train', 'r2_train'))\n",
    "\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's select the input features of our analysis: `rois` or `vbm`.\n",
    "This can be achieved using the `ROIsFeatureExtractor` or `VBMFeatureExtractor` extractors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ROIsFeatureExtractor(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"Select only the 284 ROIs features:\"\"\"\n",
    "    def fit(self, X, y):\n",
    "        return self\n",
    "\n",
    "    def transform(self, X):\n",
    "        return X[:, :284]\n",
    "\n",
    "    \n",
    "class VBMFeatureExtractor(BaseEstimator, TransformerMixin):\n",
    "    \"\"\"Select only the 284 ROIs features:\"\"\"\n",
    "    def fit(self, X, y):\n",
    "        return self\n",
    "\n",
    "    def transform(self, X):\n",
    "        return X[:, 284:]\n",
    "\n",
    "\n",
    "fe = ROIsFeatureExtractor()\n",
    "print(fe.transform(X_train).shape)\n",
    "\n",
    "fe = VBMFeatureExtractor()\n",
    "print(fe.transform(X_train).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's design a simple MLP age predictor.\n",
    "The framework is evaluated with a cross-validation approach.\n",
    "The metrics used are the root-mean-square error (RMSE).\n",
    "In this example the predictor is trained on the `roi` features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    \"\"\" Define a simple one hidden layer MLP.\n",
    "    \"\"\"\n",
    "    def __init__(self, in_features):\n",
    "        super(MLP, self).__init__()\n",
    "        self.layers = nn.Sequential(\n",
    "          nn.Linear(in_features, 120),\n",
    "          nn.ReLU(),\n",
    "          nn.Linear(120, 84),\n",
    "          nn.ReLU(),\n",
    "          nn.Linear(84, 1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.layers(x)\n",
    "\n",
    "\n",
    "class Dataset(torch.utils.data.Dataset):\n",
    "    \"\"\" A torch dataset for regression.\n",
    "    \"\"\"\n",
    "    def __init__(self, X, y=None):\n",
    "        \"\"\" Init class.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        X: array-like (n_samples, n_features)\n",
    "            training data.\n",
    "        y: array-like (n_samples, ), default None\n",
    "            target values.\n",
    "        \"\"\"\n",
    "        self.X = torch.from_numpy(X)\n",
    "        if y is not None:\n",
    "            self.y = torch.from_numpy(y)\n",
    "        else:\n",
    "            self.y = None\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.X)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        if self.y is not None:\n",
    "            return self.X[i], self.y[i]\n",
    "        else:\n",
    "            return self.X[i]\n",
    "\n",
    "\n",
    "class RegressionModel(object):\n",
    "    \"\"\" Base class for Regression models.\n",
    "    \"\"\"\n",
    "    def __init__(self, model, batch_size=10, n_epochs=10, print_freq=40):\n",
    "        \"\"\" Init class.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        model: nn.Module\n",
    "            the input model.\n",
    "        batch_size:int, default 10\n",
    "            the mini_batch size.\n",
    "        n_epochs: int, default 5\n",
    "            the number of epochs.\n",
    "        print_freq: int, default 100\n",
    "            the print frequency.\n",
    "        \"\"\"\n",
    "        self.model = model\n",
    "        self.batch_size = batch_size\n",
    "        self.n_epochs = n_epochs\n",
    "        self.print_freq = print_freq\n",
    "\n",
    "    def fit(self, X, y):\n",
    "        \"\"\" Fit model.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        X: array-like (n_samples, n_features)\n",
    "            training data.\n",
    "        y: array-like (n_samples, )\n",
    "            target values.\n",
    "        fold: int\n",
    "            the fold index.\n",
    "        \"\"\"\n",
    "        self.model.train()\n",
    "        self.reset_weights()\n",
    "        print(\"-- training model...\")\n",
    "        dataset = Dataset(X, y)\n",
    "        loader = torch.utils.data.DataLoader(\n",
    "            dataset, batch_size=self.batch_size, shuffle=True,\n",
    "            num_workers=1, multiprocessing_context=get_context(\"loky\"))\n",
    "        loss_function = nn.L1Loss()\n",
    "        optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)\n",
    "        start_time = time.time()\n",
    "        current_loss = 0.\n",
    "        for epoch in range(self.n_epochs):\n",
    "            for step, data in enumerate(loader, start=epoch * len(loader)):\n",
    "                inputs, targets = data\n",
    "                inputs, targets = inputs.float(), targets.float()\n",
    "                targets = targets.reshape((targets.shape[0], 1))\n",
    "                optimizer.zero_grad()\n",
    "                outputs = self.model(inputs)\n",
    "                loss = loss_function(outputs, targets)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                current_loss += loss.item()\n",
    "                if step % self.print_freq == 0:\n",
    "                    stats = dict(epoch=epoch, step=step,\n",
    "                                 lr=optimizer.param_groups[0][\"lr\"],\n",
    "                                 loss=loss.item(),\n",
    "                                 time=int(time.time() - start_time))\n",
    "                    print(json.dumps(stats))\n",
    "        current_loss /= (len(loader) * self.n_epochs)\n",
    "\n",
    "    def predict(self, X):\n",
    "        \"\"\" Predict using the input model.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        X: array-like (n_samples, n_features)\n",
    "            samples.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        C: array (n_samples, )\n",
    "            returns predicted values.\n",
    "        \"\"\"\n",
    "        self.model.eval()\n",
    "        dataset = Dataset(X)\n",
    "        testloader = torch.utils.data.DataLoader(\n",
    "            dataset, batch_size=self.batch_size, shuffle=False, num_workers=1,\n",
    "            multiprocessing_context=get_context(\"loky\"))\n",
    "        with torch.no_grad():\n",
    "            C = []\n",
    "            for i, inputs in enumerate(testloader):\n",
    "                inputs = inputs.float() \n",
    "                C.append(self.model(inputs))\n",
    "            C = torch.cat(C, dim=0)\n",
    "        return C.numpy().squeeze()\n",
    "\n",
    "    def reset_weights(self):\n",
    "        \"\"\" Reset all the weights of the model.\n",
    "        \"\"\"\n",
    "        def weight_reset(m):\n",
    "            if hasattr(m, \"reset_parameters\"):\n",
    "                m.reset_parameters()\n",
    "        self.model.apply(weight_reset)\n",
    "\n",
    "\n",
    "cv = problem.get_cv(X_train, y_train)\n",
    "mlp = MLP(284)\n",
    "estimator = make_pipeline(ROIsFeatureExtractor(), StandardScaler(), RegressionModel(mlp))\n",
    "\n",
    "cv_results = cross_validate(\n",
    "    estimator, X_train, y_train, scoring=['neg_root_mean_squared_error', 'r2'],\n",
    "    cv=cv, verbose=1, return_train_score=True, n_jobs=5)\n",
    "\n",
    "# Refit on all train\n",
    "estimator.fit(X_train, y_train)\n",
    "\n",
    "# Apply on test\n",
    "y_pred_train = estimator.predict(X_train)\n",
    "y_pred_test = estimator.predict(X_test)\n",
    "\n",
    "print(\"Important scores are rmse_cv_test_mean and rmse_test:\")\n",
    "scores = cv_train_test_scores(\n",
    "    rmse_cv_test=-cv_results['test_neg_root_mean_squared_error'],\n",
    "    rmse_cv_train=-cv_results['train_neg_root_mean_squared_error'],\n",
    "    r2_cv_test=cv_results['test_r2'],\n",
    "    r2_cv_train=cv_results['train_r2'],\n",
    "    y_train=y_train, y_pred_train=y_pred_train, y_test=y_test, y_pred_test=y_pred_test).T.round(3)\n",
    "print(scores)"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}