Merge branch 'main' into add-infer-accelerate-tutorial

Author: yiheng-wang-nv

2d_classification/mednist_tutorial.ipynb

@@ -575,7 +575,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "y_true = []\n",
     "y_pred = []\n",

2d_segmentation/torch/unet_evaluation_array.py

@@ -47,7 +47,7 @@ def main(tempdir):
     val_loader = DataLoader(val_ds, batch_size=1, num_workers=1, pin_memory=torch.cuda.is_available())
     dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
     post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
-    saver = SaveImage(output_dir="./output", output_ext=".png", output_postfix="seg")
+    saver = SaveImage(output_dir="./output", output_ext=".png", output_postfix="seg", scale=255)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = UNet(
         spatial_dims=2,
@@ -58,7 +58,7 @@ def main(tempdir):
         num_res_units=2,
     ).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_segmentation2d_array.pth"))
+    model.load_state_dict(torch.load("best_metric_model_segmentation2d_array.pth", weights_only=True))
     model.eval()
     with torch.no_grad():
         for val_data in val_loader:

2d_segmentation/torch/unet_evaluation_dict.py

@@ -61,7 +61,7 @@ def main(tempdir):
     val_loader = DataLoader(val_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate)
     dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
     post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
-    saver = SaveImage(output_dir="./output", output_ext=".png", output_postfix="seg")
+    saver = SaveImage(output_dir="./output", output_ext=".png", output_postfix="seg", scale=255)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = UNet(
         spatial_dims=2,
@@ -72,7 +72,7 @@ def main(tempdir):
         num_res_units=2,
     ).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_segmentation2d_dict.pth"))
+    model.load_state_dict(torch.load("best_metric_model_segmentation2d_dict.pth", weights_only=True))
 
     model.eval()
     with torch.no_grad():

3d_classification/torch/densenet_evaluation_array.py

@@ -57,7 +57,7 @@ def main():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = monai.networks.nets.DenseNet121(spatial_dims=3, in_channels=1, out_channels=2).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_classification3d_array.pth"))
+    model.load_state_dict(torch.load("best_metric_model_classification3d_array.pth", weights_only=True))
     model.eval()
     with torch.no_grad():
         num_correct = 0.0

3d_classification/torch/densenet_evaluation_dict.py

@@ -63,7 +63,7 @@ def main():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = monai.networks.nets.DenseNet121(spatial_dims=3, in_channels=1, out_channels=2).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_classification3d_dict.pth"))
+    model.load_state_dict(torch.load("best_metric_model_classification3d_dict.pth", weights_only=True))
     model.eval()
     with torch.no_grad():
         num_correct = 0.0

3d_registration/learn2reg_nlst_paired_lung_ct.ipynb

@@ -872,7 +872,7 @@
    "source": [
     "# Automatic mixed precision (AMP) for faster training\n",
     "amp_enabled = True\n",
-    "scaler = torch.cuda.amp.GradScaler()\n",
+    "scaler = torch.GradScaler(\"cuda\")\n",
     "\n",
     "# Tensorboard\n",
     "if do_save:\n",
@@ -1127,7 +1127,7 @@
     "    )\n",
     "    # load model weights\n",
     "    filename_best_model = glob.glob(os.path.join(dir_load, \"segresnet_kpt_loss_best_tre*\"))[0]\n",
-    "    model.load_state_dict(torch.load(filename_best_model))\n",
+    "    model.load_state_dict(torch.load(filename_best_model, weights_only=True))\n",
     "    # to GPU\n",
     "    model.to(device)\n",
     "\n",
@@ -1139,7 +1139,7 @@
     "# Forward pass\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
-    "    with torch.cuda.amp.autocast(enabled=amp_enabled):\n",
+    "    with torch.autocast(\"cuda\", enabled=amp_enabled):\n",
     "        ddf_image, ddf_keypoints, pred_image, pred_label = forward(\n",
     "            check_data[\"fixed_image\"].to(device),\n",
     "            check_data[\"moving_image\"].to(device),\n",

3d_registration/learn2reg_oasis_unpaired_brain_mr.ipynb

@@ -610,7 +610,7 @@
    "source": [
     "# Automatic mixed precision (AMP) for faster training\n",
     "amp_enabled = True\n",
-    "scaler = torch.cuda.amp.GradScaler()\n",
+    "scaler = torch.GradScaler(\"cuda\")\n",
     "\n",
     "# Tensorboard\n",
     "if do_save:\n",
@@ -646,7 +646,7 @@
     "\n",
     "        # Forward pass and loss\n",
     "        optimizer.zero_grad()\n",
-    "        with torch.cuda.amp.autocast(enabled=amp_enabled):\n",
+    "        with torch.autocast(\"cuda\", enabled=amp_enabled):\n",
     "            ddf_image, pred_image, pred_label_one_hot = forward(\n",
     "                fixed_image, moving_image, moving_label, model, warp_layer, num_classes=4\n",
     "            )\n",
@@ -694,7 +694,7 @@
     "                # moving_label_35 = batch_data[\"moving_label_35\"].to(device)\n",
     "                n_steps += 1\n",
     "                # Infer\n",
-    "                with torch.cuda.amp.autocast(enabled=amp_enabled):\n",
+    "                with torch.autocast(\"cuda\", enabled=amp_enabled):\n",
     "                    ddf_image, pred_image, pred_label_one_hot = forward(\n",
     "                        fixed_image, moving_image, moving_label_4, model, warp_layer, num_classes=4\n",
     "                    )\n",
@@ -840,7 +840,7 @@
     "    model = VoxelMorph()\n",
     "    # load model weights\n",
     "    filename_best_model = glob.glob(os.path.join(dir_load, \"voxelmorph_loss_best_dice_*\"))[0]\n",
-    "    model.load_state_dict(torch.load(filename_best_model))\n",
+    "    model.load_state_dict(torch.load(filename_best_model, weights_only=True))\n",
     "    # to GPU\n",
     "    model.to(device)\n",
     "\n",
@@ -860,7 +860,7 @@
     "# Forward pass\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
-    "    with torch.cuda.amp.autocast(enabled=amp_enabled):\n",
+    "    with torch.autocast(\"cuda\", enabled=amp_enabled):\n",
     "        ddf_image, pred_image, pred_label_one_hot = forward(\n",
     "            fixed_image, moving_image, moving_label_35, model, warp_layer, num_classes=35\n",
     "        )"

3d_registration/paired_lung_ct.ipynb

@@ -860,7 +860,7 @@
     "resource = \"https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/pair_lung_ct.pth\"\n",
     "dst = f\"{root_dir}/pretrained_weight.pth\"\n",
     "download_url(resource, dst)\n",
-    "model.load_state_dict(torch.load(dst))"
+    "model.load_state_dict(torch.load(dst, weights_only=True))"
    ]
   },
   {

3d_segmentation/brats_segmentation_3d.ipynb

@@ -473,14 +473,14 @@
     "        )\n",
     "\n",
     "    if VAL_AMP:\n",
-    "        with torch.cuda.amp.autocast():\n",
+    "        with torch.autocast(\"cuda\"):\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",
+    "scaler = torch.GradScaler(\"cuda\")\n",
     "# enable cuDNN benchmark\n",
     "torch.backends.cudnn.benchmark = True"
    ]
@@ -526,7 +526,7 @@
     "            batch_data[\"label\"].to(device),\n",
     "        )\n",
     "        optimizer.zero_grad()\n",
-    "        with torch.cuda.amp.autocast():\n",
+    "        with torch.autocast(\"cuda\"):\n",
     "            outputs = model(inputs)\n",
     "            loss = loss_function(outputs, labels)\n",
     "        scaler.scale(loss).backward()\n",
@@ -733,7 +733,7 @@
     }
    ],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "with torch.no_grad():\n",
     "    # select one image to evaluate and visualize the model output\n",
@@ -835,7 +835,7 @@
     }
    ],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "\n",
     "with torch.no_grad():\n",
@@ -924,7 +924,7 @@
     "        )\n",
     "\n",
     "    if VAL_AMP:\n",
-    "        with torch.cuda.amp.autocast():\n",
+    "        with torch.autocast(\"cuda\"):\n",
     "            return _compute(input)\n",
     "    else:\n",
     "        return _compute(input)"
@@ -977,7 +977,7 @@
    "source": [
     "onnx_model_path = os.path.join(root_dir, \"best_metric_model.onnx\")\n",
     "ort_session = onnxruntime.InferenceSession(onnx_model_path)\n",
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))\n",
     "model.eval()\n",
     "\n",
     "with torch.no_grad():\n",

3d_segmentation/challenge_baseline/run_net.py

@@ -219,7 +219,7 @@ def infer(data_folder=".", model_folder="runs", prediction_folder="output"):
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     net = get_net().to(device)
-    net.load_state_dict(torch.load(ckpt, map_location=device))
+    net.load_state_dict(torch.load(ckpt, map_location=device, weights_only=True))
     net.eval()
 
     image_folder = os.path.abspath(data_folder)