auto3dseg/notebooks/auto3dseg_autorunner_ref_api.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) MONAI Consortium  \n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");  \n",
    "you may not use this file except in compliance with the License.  \n",
    "You may obtain a copy of the License at  \n",
    "&nbsp;&nbsp;&nbsp;&nbsp;http://www.apache.org/licenses/LICENSE-2.0  \n",
    "Unless required by applicable law or agreed to in writing, software  \n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,  \n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  \n",
    "See the License for the specific language governing permissions and  \n",
    "limitations under the License.\n",
    "\n",
    "# MONAI Auto3DSeg Reference Python APIs\n",
    "\n",
    "In this notebook, we will break down the Auto3DSeg by the modules in the pipeline and introduce the API calls in Python and CLI commands. Particularly, if you have used the AutoRunner class, we will map the AutoRunner commands and configurations to each of the Auto3DSeg module APIs\n",
    "\n",
    "![workflow](../figures/workflow.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -c \"import monai\" || pip install -q \"monai-weekly[nibabel, tqdm]\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "\n",
    "from monai.apps import download_and_extract\n",
    "from monai.apps.auto3dseg import (\n",
    "    DataAnalyzer,\n",
    "    BundleGen,\n",
    "    AlgoEnsembleBestN,\n",
    "    AlgoEnsembleBuilder,\n",
    "    export_bundle_algo_history,\n",
    "    import_bundle_algo_history,\n",
    ")\n",
    "from monai.auto3dseg import algo_to_pickle\n",
    "from monai.bundle.config_parser import ConfigParser\n",
    "from monai.config import print_config\n",
    "from monai.utils.enums import AlgoKeys\n",
    "\n",
    "print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download dataset\n",
    "\n",
    "We provide a toy datalist file that splits a subset of the downloaded datasets into five folds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "directory = os.environ.get(\"MONAI_DATA_DIRECTORY\")\n",
    "if directory is not None:\n",
    "    os.makedirs(directory, exist_ok=True)\n",
    "root_dir = tempfile.mkdtemp() if directory is None else directory\n",
    "print(root_dir)\n",
    "\n",
    "msd_task = \"Task04_Hippocampus\"\n",
    "resource = \"https://msd-for-monai.s3-us-west-2.amazonaws.com/\" + msd_task + \".tar\"\n",
    "\n",
    "compressed_file = os.path.join(root_dir, msd_task + \".tar\")\n",
    "dataroot = os.path.join(root_dir, msd_task)\n",
    "if not os.path.exists(dataroot):\n",
    "    download_and_extract(resource, compressed_file, root_dir)\n",
    "\n",
    "datalist_file = os.path.join(\"..\", \"tasks\", \"msd\", msd_task, \"msd_\" + msd_task.lower() + \"_folds.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare a input YAML configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_cfg = {\n",
    "    \"name\": msd_task,  # optional, it is only for your own record\n",
    "    \"task\": \"segmentation\",  # optional, it is only for your own record\n",
    "    \"modality\": \"MRI\",  # required\n",
    "    \"datalist\": datalist_file,  # required\n",
    "    \"dataroot\": dataroot,  # required\n",
    "}\n",
    "input = \"./input.yaml\"\n",
    "ConfigParser.export_config_file(input_cfg, input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Breaking down the AutoRunner\n",
    "\n",
    "Below is the typical usage of AutoRunner\n",
    "```python\n",
    "runner = AutoRunner(input=input)\n",
    "runner.run() \n",
    "```\n",
    "\n",
    "The two lines cover the typical settings in Auto3DSeg and now we are going through the internal APIs calls inside these two lines\n",
    "\n",
    "## Data Analysis\n",
    "\n",
    "When the `analyze` flag is set to `True`, `AutoRunner` will call `DataAnalyzer` to analyze the datasets and generate a statisical report in YAML. Below is the equivalent Python API calls of `DataAnalyzer`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "work_dir = \"./ref_api_work_dir\"\n",
    "\n",
    "if not os.path.isdir(work_dir):\n",
    "    os.makedirs(work_dir)\n",
    "datastats_file = os.path.join(work_dir, \"data_stats.yaml\")\n",
    "analyser = DataAnalyzer(datalist_file, dataroot, output_path=datastats_file)\n",
    "datastat = analyser.get_all_case_stats()\n",
    "\n",
    "print(\"datalist file: \", os.path.abspath(datalist_file))\n",
    "print(\"dataroot path: \", os.path.abspath(dataroot))\n",
    "print(\"datastat path: \", os.path.abspath(datastats_file))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Besides the Python API call, user can also use command line interface (CLI) provided by the Python Fire:\n",
    "\n",
    "```bash\n",
    "python -m monai.apps.auto3dseg DataAnalyzer get_all_case_stats \\\n",
    "    --datalist=\"<datalist file>\" \\\n",
    "    --dataroot=\"<dataroot path>\" \\\n",
    "    --output_path=\"<datastat path>\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Algorithm Generation (algo_gen)\n",
    "\n",
    "When the `algo_gen` flag is set to `True`, `AutoRunner` will use `BundleGen` to generate monai bundles from templated algorithms in the working directory. \n",
    "\n",
    "The templated algorithms are customized for the datasets when the `generate` method is called. In detail, the `generate` method will fill the templates using information from the data_stats report. Also, it will copy the necessary scripts (train.py/infer.py) to the algorithm folder. Finally, it will create an algo_object.pkl to save the `Algo` so that it can be instantiated in the local or remote machine. Cross validation is used by default, and `num_fold` can be set to 1 if the users do not want cross validation.\n",
    "\n",
    "Below is the equivalent Python API calls of `BundleGen`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bundle_generator = BundleGen(\n",
    "    algo_path=work_dir,\n",
    "    data_stats_filename=datastats_file,\n",
    "    data_src_cfg_name=input,\n",
    ")\n",
    "\n",
    "bundle_generator.generate(work_dir, num_fold=5)\n",
    "\n",
    "print(\"algo path: \", os.path.abspath(work_dir))\n",
    "print(\"data_stats file: \", os.path.abspath(datastats_file))\n",
    "print(\"task input file: \", os.path.abspath(input))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Besides the Python API call, user can also use command line interface (CLI) provided by the user's OS. One example is the following bash commands:\n",
    "\n",
    "```bash\n",
    "python -m monai.apps.auto3dseg BundleGen generate \\\n",
    "    --algo_path=\"<algo path>\" \\\n",
    "    --data_stats_filename=\"<data_stats file>\" \\\n",
    "    --data_src_cfg_name=\"<task input file>\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting and saving the algorithm generation history to the local drive\n",
    "\n",
    "If the users continue to train the algorithms on local system, The history of the algorithm generation can be fetched via `get_history` method of the `BundleGen` object. There also are scenarios that users need to stop the Python process after the `algo_gen`. For example, the users may need to transfer the files to a remote cluster to start the training. `Auto3DSeg` offers a utility function `export_bundle_algo_history` to dump the history to hard drive and recall it by `import_bundle_algo_history`. \n",
    "\n",
    "If the files are copied to a remote system, please ensure the algorithm templates are also copied there. Some functions require the path to instantiate the algorithm class properly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = bundle_generator.get_history()\n",
    "export_bundle_algo_history(history)  # save the Algo objects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add training parameters to cut down the training time in this notebook (Optional)\n",
    "\n",
    "This step is not required, but for demo purposes, we'll set a limit of the epochs to train the algorithms. \n",
    "\n",
    "Some algorithms in **Auto3DSeg** use `epoch` to mark the progress of training, while the others use `iteration` to iterate the loops. \n",
    "Below is the code block to convert `num_epoch` to iteration style and override all algorithms with the same training parameters for a 1-GPU/2-GPU machine. \n",
    "\n",
    "It is not required for the users to set the `train_param`. \n",
    "The users can use either `train()` or `train({})` if no changes are needed.\n",
    "Then the algorithms will go for the full training and repeat 5 folds.\n",
    "\n",
    "On the other hand, users can also use set `train_param` for each algorithm.\n",
    "\n",
    "\n",
    "For demo purposes, below is a code block to convert num_epoch to iteration style and override all algorithms with the same training parameters.\n",
    "The setup works fine for a machine that has GPUs less than or equal to 8.\n",
    "The datalist in this example is only using a subset of the original dataset.\n",
    "Users need to ensure the number of GPUs is not greater than the number that the training dataset can be partitioned.\n",
    "For example, the following code block is not suitable for a 16-GPU system.\n",
    "In such cases, please change the code block accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_epochs = 2  # change epoch number to 2 to cut down the notebook running time\n",
    "\n",
    "train_param = {\n",
    "    \"num_epochs_per_validation\": 1,\n",
    "    \"num_images_per_batch\": 2,\n",
    "    \"num_epochs\": max_epochs,\n",
    "    \"num_warmup_epochs\": 1,\n",
    "}\n",
    "\n",
    "print(train_param)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the neural networks sequentially\n",
    "\n",
    "The algo_gen history contains `Algo` object that has multiple methods such as `train` and `predict`. We can easily use such APIs to trigger neural network training. By default, `AutoRunnner` will start a training on a single node (single or multiple GPUs) in a seqential manner.\n",
    "\n",
    "`algo_to_pickle` is optional and it will update the dumped Algo objects with the accuracies information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = import_bundle_algo_history(work_dir, only_trained=False)\n",
    "for algo_dict in history:\n",
    "    algo = algo_dict[AlgoKeys.ALGO]\n",
    "    algo.train(train_param)  # can use default params by `algo.train()`\n",
    "    acc = algo.get_score()\n",
    "    algo_to_pickle(algo, template_path=algo.template_path, best_metric=acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Information about Hyper-parameter Optimization (HPO)\n",
    "\n",
    "Another method to handle the neural network training is to perform HPO (e.g. training & searching). This is made possible by NNI or Optuna packages which are installed in the MONAI development environment. `AutoRunner` uses Microsoft `NNI` as backend via the `NNIGen`, but Optuna HPO can also be chosen via the `OptunaGen` method in the Auto3DSeg pipeline\n",
    "\n",
    "For more information, please refer to the [HPO NNI notebook](hpo_nni.ipynb) and the [HPO Optuna notebook](hpo_optuna.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ensemble\n",
    "\n",
    "Finally, after the neural networks are trained, `AutoRunner` will apply the ensemble methods in Auto3DSeg to improve the overall performance. \n",
    "\n",
    "Here we used a utility function `import_bundle_algo_history` to load the `Algo` that is trained into the ensemble. With the history loaded, we build an ensemble method and use the method to perform the inference on all testing data.\n",
    "\n",
    "> NOTE: Because we need to get the prediction in Python, there is no alternative CLI commands for this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = import_bundle_algo_history(work_dir, only_trained=True)\n",
    "builder = AlgoEnsembleBuilder(history, input)\n",
    "builder.set_ensemble_method(AlgoEnsembleBestN(n_best=5))\n",
    "ensembler = builder.get_ensemble()\n",
    "preds = ensembler()"
   ]
  }
 ],
 "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.8.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "adad466f3ae0207e0597f6bf999e86a6601b04f4db83734c81b06a9e73a01812"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}