Make iteration-base training params epoch-base in auto3dseg (#1213)

Signed-off-by: Mingxin Zheng
<[email protected]>

Fixes #1212 

### Description
A few sentences describing the changes proposed in this pull request.

### Checks
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Avoid including large-size files in the PR.
- [x] Clean up long text outputs from code cells in the notebook.
- [x] For security purposes, please check the contents and remove any
sensitive info such as user names and private key.
- [x] Ensure (1) hyperlinks and markdown anchors are working (2) use
relative paths for tutorial repo files (3) put figure and graphs in the
`./figure` folder
- [x] Notebook runs automatically
    - `./runner.sh -t 'auto3dseg/notebooks/auto3dseg_hello_world.ipynb'`
- `./runner.sh -t
'auto3dseg/notebooks/auto3dseg_autorunner_ref_api.ipynb'`
    - `./runner.sh -t 'auto3dseg/notebooks/auto_runner.ipynb'`
    - `./runner.sh -t 'auto3dseg/notebooks/hpo_optuna.ipynb'`

---------

Signed-off-by: Mingxin Zheng <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Wenqi Li <[email protected]>

Author: 3 people

3d_segmentation/unetr_btcv_segmentation_3d.ipynb

@@ -33,7 +33,7 @@
     "\n",
     "Under Institutional Review Board (IRB) supervision, 50 abdomen CT scans of were randomly selected from a combination of an ongoing colorectal cancer chemotherapy trial, and a retrospective ventral hernia study. The 50 scans were captured during portal venous contrast phase with variable volume sizes (512 x 512 x 85 - 512 x 512 x 198) and field of views (approx. 280 x 280 x 280 mm3 - 500 x 500 x 650 mm3). The in-plane resolution varies from 0.54 x 0.54 mm2 to 0.98 x 0.98 mm2, while the slice thickness ranges from 2.5 mm to 5.0 mm. \n",
     "\n",
-    "Target: 13 abdominal organs including 1. Spleen 2. Right Kidney 3. Left Kideny 4.Gallbladder 5.Esophagus 6. Liver 7. Stomach 8.Aorta 9. IVC 10. Portal and Splenic Veins 11. Pancreas 12 Right adrenal gland 13 Left adrenal gland.\n",
+    "Target: 13 abdominal organs including 1. Spleen 2. Right Kidney 3. Left Kidney 4.Gallbladder 5.Esophagus 6. Liver 7. Stomach 8.Aorta 9. IVC 10. Portal and Splenic Veins 11. Pancreas 12 Right adrenal gland 13 Left adrenal gland.\n",
     "\n",
     "Modality: CT\n",
     "Size: 30 3D volumes (24 Training + 6 Testing)  \n",

3d_segmentation/unetr_btcv_segmentation_3d_lightning.ipynb

@@ -36,7 +36,7 @@
     "\n",
     "Under Institutional Review Board (IRB) supervision, 50 abdomen CT scans of were randomly selected from a combination of an ongoing colorectal cancer chemotherapy trial, and a retrospective ventral hernia study. The 50 scans were captured during portal venous contrast phase with variable volume sizes (512 x 512 x 85 - 512 x 512 x 198) and field of views (approx. 280 x 280 x 280 mm3 - 500 x 500 x 650 mm3). The in-plane resolution varies from 0.54 x 0.54 mm2 to 0.98 x 0.98 mm2, while the slice thickness ranges from 2.5 mm to 5.0 mm. \n",
     "\n",
-    "Target: 13 abdominal organs including 1. Spleen 2. Right Kidney 3. Left Kideny 4.Gallbladder 5.Esophagus 6. Liver 7. Stomach 8.Aorta 9. IVC 10. Portal and Splenic Veins 11. Pancreas 12 Right adrenal gland 13 Left adrenal gland.\n",
+    "Target: 13 abdominal organs including 1. Spleen 2. Right Kidney 3. Left Kidney 4.Gallbladder 5.Esophagus 6. Liver 7. Stomach 8.Aorta 9. IVC 10. Portal and Splenic Veins 11. Pancreas 12 Right adrenal gland 13 Left adrenal gland.\n",
     "\n",
     "Modality: CT\n",
     "Size: 30 3D volumes (24 Training + 6 Testing)  \n",

auto3dseg/notebooks/auto3dseg_autorunner_ref_api.ipynb

@@ -52,7 +52,6 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "import torch\n",
     "import tempfile\n",
     "\n",
     "from monai.apps import download_and_extract\n",
@@ -64,23 +63,21 @@
     "    export_bundle_algo_history,\n",
     "    import_bundle_algo_history,\n",
     ")\n",
-    "from monai.auto3dseg import algo_to_pickle, datafold_read\n",
+    "from monai.auto3dseg import algo_to_pickle\n",
     "from monai.bundle.config_parser import ConfigParser\n",
     "from monai.config import print_config\n",
     "\n",
     "print_config()"
    ]
   },
   {
+   "attachments": {},
    "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.\n",
-    "\n",
-    "> NOTE: Each validation set only has 6 images in one fold of training.\n",
-    "> Therefore, we need to set a limit on the total number of GPUs we're using in this notebook."
+    "We provide a toy datalist file that splits a subset of the downloaded datasets into five folds."
    ]
   },
   {
@@ -101,11 +98,7 @@
     "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\")\n",
-    "\n",
-    "if torch.cuda.device_count() > 6:\n",
-    "    os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
-    "    os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1,2,3,4,5\""
+    "datalist_file = os.path.join(\"..\", \"tasks\", \"msd\", msd_task, \"msd_\" + msd_task.lower() + \"_folds.json\")"
    ]
   },
   {
@@ -231,14 +224,15 @@
    ]
   },
   {
+   "attachments": {},
    "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 make sure the algorithm templates are also copied there. Some functions require the path to instantiate the algorithm class properly."
+    "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."
    ]
   },
   {
@@ -252,6 +246,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -266,7 +261,15 @@
     "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."
+    "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."
    ]
   },
   {
@@ -277,24 +280,11 @@
    "source": [
     "max_epochs = 2  # change epoch number to 2 to cut down the notebook running time\n",
     "\n",
-    "# safeguard to ensure max_epochs is greater or equal to 2\n",
-    "max_epochs = max(max_epochs, 2)\n",
-    "\n",
-    "num_gpus = 1 if \"multigpu\" in input and not input[\"multigpu\"] else torch.cuda.device_count()\n",
-    "\n",
-    "num_epoch = max_epochs\n",
-    "num_images_per_batch = 2\n",
-    "files_train_fold0, _ = datafold_read(datalist_file, \"\", 0)\n",
-    "n_data = len(files_train_fold0)\n",
-    "n_iter = int(num_epoch * n_data / num_images_per_batch / max(num_gpus, 1))\n",
-    "n_iter_val = int(n_iter / 2)\n",
-    "\n",
     "train_param = {\n",
-    "    \"num_iterations\": n_iter,\n",
-    "    \"num_iterations_per_validation\": n_iter_val,\n",
-    "    \"num_images_per_batch\": num_images_per_batch,\n",
-    "    \"num_epochs\": num_epoch,\n",
-    "    \"num_warmup_iterations\": n_iter_val,\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)"

auto3dseg/notebooks/auto3dseg_hello_world.ipynb

@@ -54,7 +54,6 @@
     "import nibabel as nib\n",
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
-    "import torch\n",
     "\n",
     "from monai.apps.auto3dseg import AutoRunner\n",
     "from monai.config import print_config\n",
@@ -64,17 +63,15 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Simulate a special dataset\n",
     "\n",
     "It is well known that AI takes time to train. To provide the \"Hello World!\" experience of Auto3D in this notebook, we will simulate a small dataset and run training only for multiple epochs. Due to the nature of AI, the performance shouldn't be highly expected, but the entire pipeline will be completed within minutes!\n",
     "\n",
-    "`sim_datalist` provides the information of the simulated datasets. It lists 12 training and 2 testing images and labels. The training data are split into 3 folds. Each fold will use 8 images to train and 4 images to validate. The size of the dimension is defined by the `sim_dim` .\n",
-    "\n",
-    "> NOTE: Each validation set only has 4 images in one fold of training.\n",
-    "> Therefore, we need to set a limit on the total number of GPUs we're using in this notebook."
+    "`sim_datalist` provides the information of the simulated datasets. It lists 12 training and 2 testing images and labels. The training data are split into 3 folds. Each fold will use 8 images to train and 4 images to validate. The size of the dimension is defined by the `sim_dim` ."
    ]
   },
   {
@@ -104,11 +101,7 @@
     "    ],\n",
     "}\n",
     "\n",
-    "sim_dim = (64, 64, 64)\n",
-    "\n",
-    "if torch.cuda.device_count() > 4:\n",
-    "    os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
-    "    os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1,2,3\""
+    "sim_dim = (64, 64, 64)"
    ]
   },
   {
@@ -216,10 +209,15 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Override the training parameters so that we can complete the pipeline in minutes"
+    "## Override the training parameters so that we can complete the pipeline in minutes\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",
+    "If users would like to use more than one GPU, they can change the `CUDA_VISIBLE_DEVICES`, or just remove the key to use all available devices.\n",
+    "Users also need to ensure the number of GPUs is not greater than the number that the training dataset can be partitioned."
    ]
   },
   {
@@ -230,16 +228,12 @@
    "source": [
     "max_epochs = 2\n",
     "\n",
-    "# safeguard to ensure max_epochs is greater or equal to 2\n",
-    "max_epochs = max(max_epochs, 2)\n",
-    "\n",
     "train_param = {\n",
     "    \"CUDA_VISIBLE_DEVICES\": [0],  # use only 1 gpu\n",
-    "    \"num_iterations\": 4 * max_epochs,\n",
-    "    \"num_iterations_per_validation\": 2 * max_epochs,\n",
+    "    \"num_epochs_per_validation\": 1,\n",
     "    \"num_images_per_batch\": 2,\n",
     "    \"num_epochs\": max_epochs,\n",
-    "    \"num_warmup_iterations\": 2 * max_epochs,\n",
+    "    \"num_warmup_epochs\": 1,\n",
     "}\n",
     "runner.set_training_params(train_param)\n",
     "runner.set_num_fold(num_fold=1)"

auto3dseg/notebooks/auto_runner.ipynb

@@ -60,28 +60,24 @@
    "source": [
     "import os\n",
     "import tempfile\n",
-    "import torch\n",
     "\n",
     "from monai.bundle.config_parser import ConfigParser\n",
     "from monai.apps import download_and_extract\n",
     "\n",
     "from monai.apps.auto3dseg import AutoRunner\n",
-    "from monai.auto3dseg import datafold_read\n",
     "from monai.config import print_config\n",
     "\n",
     "print_config()"
    ]
   },
   {
+   "attachments": {},
    "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.\n",
-    "\n",
-    "> NOTE: Each validation set only has 6 images in one fold of training.\n",
-    "> Therefore, we need to set a limit on the total number of GPUs we're using in this notebook."
+    "We provide a toy datalist file that splits a subset of the downloaded datasets into five folds."
    ]
   },
   {
@@ -102,11 +98,7 @@
     "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\")\n",
-    "\n",
-    "if torch.cuda.device_count() > 6:\n",
-    "    os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
-    "    os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1,2,3,4,5\""
+    "datalist_file = os.path.join(\"..\", \"tasks\", \"msd\", msd_task, \"msd_\" + msd_task.lower() + \"_folds.json\")"
    ]
   },
   {
@@ -267,18 +259,26 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Customize training parameters by override the default values\n",
     "\n",
     "`set_training_params` in `AutoRunner` provides an interface to change all algorithms' training parameters in one line. \n",
     "\n",
-    "> NOTE **Auto3DSeg** uses MONAI bundle templates to perform training, validation, and inference. The number of epochs/iterations of training is specified by the config files in each template.\n",
-    "> Users can override these these values in the bundle templates.\n",
-    "> But users should consider that some bundle templates may use `num_iterations` and other may use `num_epochs` to iterate.\n",
+    "NOTE: \n",
+    "**Auto3DSeg** uses MONAI bundle templates to perform training, validation, and inference.\n",
+    "The number of epochs/iterations of training is specified by the config files in each template.\n",
+    "Users can override these these values in the bundle templates.\n",
+    "But users should consider that some bundle templates may use `num_iterations` and other may use `num_epochs` to iterate.\n",
     "\n",
-    "For demo purpose, below is a code block to convert num_epoch to iteration style and override all algorithms with the same training parameters for 1-GPU/2-GPU machine. \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.\n"
    ]
   },
   {
@@ -289,25 +289,13 @@
    "source": [
     "max_epochs = 2\n",
     "\n",
-    "# safeguard to ensure max_epochs is greater or equal to 2\n",
-    "max_epochs = max(max_epochs, 2)\n",
-    "\n",
-    "num_gpus = 1 if \"multigpu\" in input_cfg and not input_cfg[\"multigpu\"] else torch.cuda.device_count()\n",
-    "\n",
-    "num_epoch = max_epochs\n",
-    "num_images_per_batch = 2\n",
-    "files_train_fold0, _ = datafold_read(datalist_file, \"\", 0)\n",
-    "n_data = len(files_train_fold0)\n",
-    "n_iter = int(num_epoch * n_data / num_images_per_batch / num_gpus)\n",
-    "n_iter_val = int(n_iter / 2)\n",
-    "\n",
     "train_param = {\n",
-    "    \"num_iterations\": n_iter,\n",
-    "    \"num_iterations_per_validation\": n_iter_val,\n",
-    "    \"num_images_per_batch\": num_images_per_batch,\n",
-    "    \"num_epochs\": num_epoch,\n",
-    "    \"num_warmup_iterations\": n_iter_val,\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",
     "runner = AutoRunner(input=input)\n",
     "runner.set_training_params(params=train_param)\n",
     "# runner.run()"
@@ -360,25 +348,26 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Train model with HPO\n",
     "\n",
     "**Auto3DSeg** supports hyper parameter optimization (HPO) via `NNI` and `Optuna` backends.\n",
-    "If you wound like to the use `Optuna`, please check the [notebook](hpo_optuna.ipynb) for detailed usage.\n",
+    "If you would like to the use `Optuna`, please check the [notebook](hpo_optuna.ipynb) for detailed usage.\n",
     "\n",
     "Here we demonstrate the HPO option with `NNI` by Microsoft.\n",
     "Please install it via `pip install nni` if you hope to execute HPO with it in tutorial and haven't done so in the beginning of the notebook.\n",
     "AutoRunner supports `NNI` backend with a grid search method via automatically generating a the `NNI` config and run `nnictl` commands in subprocess.\n",
     "\n",
     "## Use `AutoRunner` with `NNI` backend to perform grid search\n",
     "\n",
-    "After `runner.run()` is executed, `nni` will attempt to start a web service using port 8088 by default. If you are running the tutorial in a remote host, please make sure the port is available on the system.\n",
+    "After `runner.run()` is executed, `nni` will attempt to start a web service using port 8088 by default. If you are running the tutorial in a remote host, please ensure the port is available on the system.\n",
     "\n",
     "> NOTE: it is recommended to turn off ensemble if the users are using HPO features.\n",
     "> By default, all the models are saved under the working directory, including the ones tuned by the HPO package.\n",
-    "> Users may want to read the HPO results before the taking the next step.\n",
+    "> Users may want to read the HPO results before taking the next step.\n",
     "> If the users want to ensemble all the models, the `ensemble` option can be set to True."
    ]
   },
@@ -395,6 +384,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -403,6 +393,7 @@
     "The default `NNI` config that `AutoRunner` looks like below. User can override some of the parameters via the `set_hpo_params` interface:\n",
     "\n",
     "```python\n",
+    "import torch\n",
     "default_nni_config = {\n",
     "    \"trialCodeDirectory\": \".\",\n",
     "    \"trialGpuNumber\": torch.cuda.device_count(),\n",
@@ -449,19 +440,19 @@
    "outputs": [],
    "source": [
     "runner = AutoRunner(input=input, hpo=True, ensemble=False)\n",
+    "num_epoch = 2\n",
     "hpo_params = {\n",
     "    \"maxTrialNumber\": 20,\n",
     "    \"maxExperimentDuration\": \"30m\",\n",
-    "    \"num_iterations\": n_iter,\n",
-    "    \"num_iterations_per_validation\": n_iter_val,\n",
-    "    \"num_images_per_batch\": num_images_per_batch,\n",
-    "    \"num_epochs\": num_epoch,\n",
-    "    \"num_warmup_iterations\": n_iter_val,\n",
-    "    \"training#num_iterations\": n_iter,\n",
-    "    \"training#num_iterations_per_validation\": n_iter_val,\n",
-    "    \"searching#num_iterations\": n_iter,\n",
-    "    \"searching#num_iterations_per_validation\": n_iter_val,\n",
-    "    \"searching#num_warmup_iterations\": n_iter,\n",
+    "    \"num_epochs_per_validation\": 1,\n",
+    "    \"num_images_per_batch\": 1,\n",
+    "    \"num_epochs\": 2,\n",
+    "    \"num_warmup_epochs\": 1,\n",
+    "    \"training#num_epochs\": 2,\n",
+    "    \"training#num_epochs_per_validation\": 1,\n",
+    "    \"searching#num_epochs\": 2,\n",
+    "    \"searching#num_epochs_per_validation\": 1,\n",
+    "    \"searching#num_warmup_epochs\": 1,\n",
     "}\n",
     "search_space = {\"learning_rate\": {\"_type\": \"choice\", \"_value\": [0.0001, 0.01]}}\n",
     "runner.set_num_fold(num_fold=1)\n",