Add AMP inference logic in Brats tutorial (#335)

* [DLMED] add AMP inference

Signed-off-by: Nic Ma <nma@nvidia.com>

* [DLMED] fix format issue

Signed-off-by: Nic Ma <nma@nvidia.com>

Author: Nic-Ma

3d_segmentation/brats_segmentation_3d.ipynb

@@ -435,6 +435,7 @@
    "source": [
     "max_epochs = 300\n",
     "val_interval = 1\n",
+    "VAL_AMP = True\n",
     "\n",
     "# standard PyTorch program style: create SegResNet, DiceLoss and Adam optimizer\n",
     "device = torch.device(\"cuda:0\")\n",
@@ -457,9 +458,28 @@
     "    [EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold_values=True)]\n",
     ")\n",
     "\n",
+    "\n",
+    "# define inference method\n",
+    "def inference(input):\n",
+    "\n",
+    "    def _compute(input):\n",
+    "        return sliding_window_inference(\n",
+    "            inputs=input,\n",
+    "            roi_size=(240, 240, 160),\n",
+    "            sw_batch_size=1,\n",
+    "            predictor=model,\n",
+    "            overlap=0.5,\n",
+    "        )\n",
+    "\n",
+    "    if VAL_AMP:\n",
+    "        with torch.cuda.amp.autocast():\n",
+    "            return _compute(input)\n",
+    "    else:\n",
+    "        return _compute(input)\n",
+    "\n",
+    "\n",
     "# use amp to accelerate training\n",
     "scaler = torch.cuda.amp.GradScaler()\n",
-    "VAL_AMP = True\n",
     "# enable cuDNN benchmark\n",
     "torch.backends.cudnn.benchmark = True"
    ]
@@ -531,17 +551,7 @@
     "                    val_data[\"image\"].to(device),\n",
     "                    val_data[\"label\"].to(device),\n",
     "                )\n",
-    "\n",
-    "                def _compute_sliding_window(input):\n",
-    "                    return sliding_window_inference(\n",
-    "                        inputs=input, roi_size=(240, 240, 160), sw_batch_size=1, predictor=model, overlap=0.5\n",
-    "                    )\n",
-    "\n",
-    "                if VAL_AMP:\n",
-    "                    with torch.cuda.amp.autocast():\n",
-    "                        val_outputs = _compute_sliding_window(val_inputs)\n",
-    "                else:\n",
-    "                    val_outputs = _compute_sliding_window(val_inputs)\n",
+    "                val_outputs = inference(val_inputs)\n",
     "                val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]\n",
     "                dice_metric(y_pred=val_outputs, y=val_labels)\n",
     "                dice_metric_batch(y_pred=val_outputs, y=val_labels)\n",
@@ -732,9 +742,7 @@
     "    val_input = val_ds[6][\"image\"].unsqueeze(0).to(device)\n",
     "    roi_size = (128, 128, 64)\n",
     "    sw_batch_size = 4\n",
-    "    val_output = sliding_window_inference(\n",
-    "        inputs=val_input, roi_size=(240, 240, 160), sw_batch_size=1, predictor=model, overlap=0.5\n",
-    "    )\n",
+    "    val_output = inference(val_input)\n",
     "    val_output = post_trans(val_output[0])\n",
     "    plt.figure(\"image\", (24, 6))\n",
     "    for i in range(4):\n",
@@ -835,9 +843,7 @@
     "with torch.no_grad():\n",
     "    for val_data in val_org_loader:\n",
     "        val_inputs = val_data[\"image\"].to(device)\n",
-    "        val_data[\"pred\"] = sliding_window_inference(\n",
-    "            inputs=val_inputs, roi_size=(240, 240, 160), sw_batch_size=1, predictor=model, overlap=0.5\n",
-    "        )\n",
+    "        val_data[\"pred\"] = inference(val_inputs)\n",
     "        val_data = [post_transforms(i) for i in decollate_batch(val_data)]\n",
     "        val_outputs, val_labels = from_engine([\"pred\", \"label\"])(val_data)\n",
     "        dice_metric(y_pred=val_outputs, y=val_labels)\n",
@@ -879,7 +885,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },