Add support for pytorch 2.6+

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

@@ -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

@@ -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

@@ -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",

3d_registration/learn2reg_oasis_unpaired_brain_mr.ipynb

@@ -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",

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

@@ -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",
@@ -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)

3d_segmentation/spleen_segmentation_3d.ipynb

@@ -640,7 +640,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",
     "    for i, val_data in enumerate(val_loader):\n",
@@ -730,7 +730,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",
@@ -827,7 +827,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",
     "\n",
     "with torch.no_grad():\n",

3d_segmentation/spleen_segmentation_3d_visualization_basic.ipynb

@@ -823,7 +823,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",

3d_segmentation/swin_unetr_brats21_segmentation_3d.ipynb

@@ -885,7 +885,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"model.pt\"))[\"state_dict\"])\n",
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"model.pt\"), weights_only=True)[\"state_dict\"])\n",
     "model.to(device)\n",
     "model.eval()\n",
     "\n",

3d_segmentation/swin_unetr_btcv_segmentation_3d.ipynb

@@ -472,7 +472,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "weight = torch.load(\"./model_swinvit.pt\")\n",
+    "weight = torch.load(\"./model_swinvit.pt\", weights_only=True)\n",
     "model.load_from(weights=weight)\n",
     "print(\"Using pretrained self-supervied Swin UNETR backbone weights !\")"
    ]
@@ -590,7 +590,7 @@
     "metric_values = []\n",
     "while global_step < max_iterations:\n",
     "    global_step, dice_val_best, global_step_best = train(global_step, train_loader, dice_val_best, global_step_best)\n",
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))"
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))"
    ]
   },
   {
@@ -679,7 +679,7 @@
    ],
    "source": [
     "case_num = 4\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",
     "with torch.no_grad():\n",
     "    img_name = os.path.split(val_ds[case_num][\"image\"].meta[\"filename_or_obj\"])[1]\n",

3d_segmentation/torch/unet_evaluation_array.py

@@ -63,7 +63,7 @@ def main(tempdir):
         num_res_units=2,
     ).to(device)
 
-    model.load_state_dict(torch.load("best_metric_model_segmentation3d_array.pth"))
+    model.load_state_dict(torch.load("best_metric_model_segmentation3d_array.pth", weights_only=True))
     model.eval()
     with torch.no_grad():
         for val_data in val_loader:

3d_segmentation/torch/unet_evaluation_dict.py

@@ -81,7 +81,7 @@ def main(tempdir):
         num_res_units=2,
     ).to(devices[0])
 
-    model.load_state_dict(torch.load("best_metric_model_segmentation3d_dict.pth"))
+    model.load_state_dict(torch.load("best_metric_model_segmentation3d_dict.pth", weights_only=True))
 
     # if we have multiple GPUs, set data parallel to execute sliding window inference
     if len(devices) > 1:

3d_segmentation/torch/unet_inference_dict.py

@@ -91,7 +91,7 @@ def main(tempdir):
         strides=(2, 2, 2, 2),
         num_res_units=2,
     ).to(device)
-    net.load_state_dict(torch.load("best_metric_model_segmentation3d_dict.pth"))
+    net.load_state_dict(torch.load("best_metric_model_segmentation3d_dict.pth", weights_only=True))
 
     net.eval()
     with torch.no_grad():

3d_segmentation/unetr_btcv_segmentation_3d.ipynb

@@ -680,7 +680,7 @@
     "metric_values = []\n",
     "while global_step < max_iterations:\n",
     "    global_step, dice_val_best, global_step_best = train(global_step, train_loader, dice_val_best, global_step_best)\n",
-    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))"
+    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\"), weights_only=True))"
    ]
   },
   {
@@ -769,7 +769,7 @@
    ],
    "source": [
     "case_num = 4\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",
     "with torch.no_grad():\n",
     "    img_name = os.path.split(val_ds[case_num][\"image\"].meta[\"filename_or_obj\"])[1]\n",

acceleration/TensorRT_inference_acceleration.ipynb

@@ -284,7 +284,7 @@
     "device = workflow.device\n",
     "spatial_shape = (1, 3, 736, 480)\n",
     "model = workflow.network_def\n",
-    "model.load_state_dict(torch.load(model_weight))\n",
+    "model.load_state_dict(torch.load(model_weight, weights_only=True))\n",
     "model.to(device)\n",
     "model.eval()\n",
     "\n",

acceleration/distributed_training/unet_evaluation_ddp.py

@@ -119,7 +119,7 @@ def evaluate(args):
     # config mapping to expected GPU device
     map_location = {"cuda:0": f"cuda:{args.local_rank}"}
     # load model parameters to GPU device
-    model.load_state_dict(torch.load("final_model.pth", map_location=map_location))
+    model.load_state_dict(torch.load("final_model.pth", map_location=map_location, weights_only=True))
 
     model.eval()
     with torch.no_grad():

acceleration/distributed_training/unet_evaluation_horovod.py

@@ -133,7 +133,7 @@ def evaluate(args):
     ).to(device)
     if hvd.rank() == 0:
         # load model parameters for evaluation
-        model.load_state_dict(torch.load("final_model.pth"))
+        model.load_state_dict(torch.load("final_model.pth", weights_only=True))
     # Horovod broadcasts parameters
     hvd.broadcast_parameters(model.state_dict(), root_rank=0)
 

automl/DiNTS/decode_plot.py

@@ -56,7 +56,7 @@ def plot_graph(
     Return:
         graphviz graph.
     """
-    code = torch.load(codepath)
+    code = torch.load(codepath, weights_only=True)
     arch_code_a = code["arch_code_a"]
     arch_code_c = code["arch_code_c"]
     ga = Digraph("G", filename=filename, engine="neato")

automl/DiNTS/search_dints.py

@@ -422,7 +422,7 @@ def main():
 
     if args.checkpoint != None and os.path.isfile(args.checkpoint):
         print("[info] fine-tuning pre-trained checkpoint {0:s}".format(args.checkpoint))
-        model.load_state_dict(torch.load(args.checkpoint, map_location=device))
+        model.load_state_dict(torch.load(args.checkpoint, map_location=device, weights_only=True))
         torch.cuda.empty_cache()
     else:
         print("[info] training from scratch")

automl/DiNTS/train_dints.py

@@ -346,7 +346,7 @@ def main():
     train_loader = ThreadDataLoader(train_ds, num_workers=8, batch_size=num_images_per_batch, shuffle=True)
     val_loader = ThreadDataLoader(val_ds, num_workers=4, batch_size=1, shuffle=False)
 
-    ckpt = torch.load(args.arch_ckpt)
+    ckpt = torch.load(args.arch_ckpt, weights_only=True)
     node_a = ckpt["node_a"]
     arch_code_a = ckpt["arch_code_a"]
     arch_code_c = ckpt["arch_code_c"]
@@ -399,7 +399,7 @@ def main():
 
     if args.checkpoint != None and os.path.isfile(args.checkpoint):
         print("[info] fine-tuning pre-trained checkpoint {0:s}".format(args.checkpoint))
-        model.load_state_dict(torch.load(args.checkpoint, map_location=device))
+        model.load_state_dict(torch.load(args.checkpoint, map_location=device, weights_only=True))
         torch.cuda.empty_cache()
     else:
         print("[info] training from scratch")

bundle/02_mednist_classification.ipynb

@@ -575,7 +575,7 @@
     "\n",
     "# loads the weights from the given file (which needs to be set on the command line) then calls \"evaluate\"\n",
     "evaluate:\n",
-    "- '$@net.load_state_dict(torch.load(@ckpt_file))'\n",
+    "- '$@net.load_state_dict(torch.load(@ckpt_file, weights_only=True))'\n",
     "- '$scripts.evaluate(@net, @eval_dl, @class_names, @device)'\n"
    ]
   },

bundle/04_integrating_code.ipynb

@@ -597,7 +597,7 @@
     "dataloader: '$scripts.dataloaders.get_dataloader(False, @transforms)'\n",
     "\n",
     "test:\n",
-    "- $@net.load_state_dict(torch.load('./cifar_net.pth'))\n",
+    "- $@net.load_state_dict(torch.load('./cifar_net.pth', weights_only=True))\n",
     "- $scripts.test.test(@net, @dataloader)\n"
    ]
   },
@@ -723,7 +723,7 @@
     "transforms: '$scripts.transforms.transform'\n",
     "\n",
     "inference:\n",
-    "- $@net.load_state_dict(torch.load('./cifar_net.pth'))\n",
+    "- $@net.load_state_dict(torch.load('./cifar_net.pth', weights_only=True))\n",
     "- $scripts.inference.inference(@net, @transforms, @input_files)"
    ]
   },

bundle/hybrid_programming/scripts/inference.py

@@ -31,7 +31,7 @@ def run(config_file: Union[str, Sequence[str]], ckpt_path: str):
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     # instantialize the components
     model = parser.get_parsed_content("network").to(device)
-    model.load_state_dict(torch.load(ckpt_path))
+    model.load_state_dict(torch.load(ckpt_path, weights_only=True))
 
     dataloader = parser.get_parsed_content("dataloader")
     if len(dataloader) == 0:

competitions/MICCAI/surgtoolloc/classification_files/train.py

@@ -62,7 +62,9 @@ def main(cfg):
     model.to(cfg.device)
 
     if cfg.weights is not None:
-        model.load_state_dict(torch.load(os.path.join(f"{cfg.output_dir}/fold{cfg.fold}", cfg.weights))["model"])
+        model.load_state_dict(
+            torch.load(os.path.join(f"{cfg.output_dir}/fold{cfg.fold}", cfg.weights), weights_only=True)["model"]
+        )
         print(f"weights from: {cfg.weights} are loaded.")
 
     # set optimizer, lr scheduler

competitions/kaggle/RANZCR/4th_place_solution/models/seg_model.py

@@ -208,7 +208,7 @@ def __init__(self, cfg):
 
         if cfg.pretrained_weights is not None:
             self.load_state_dict(
-                torch.load(cfg.pretrained_weights, map_location="cpu")["model"],
+                torch.load(cfg.pretrained_weights, map_location="cpu", weights_only=True)["model"],
                 strict=True,
             )
             print("weights loaded from", cfg.pretrained_weights)

competitions/kaggle/RANZCR/4th_place_solution/train.py

@@ -248,7 +248,7 @@ def run_infer(weights_folder_path, cfg):
 
     nets = []
     for path in all_path:
-        state_dict = torch.load(path)["model"]
+        state_dict = torch.load(path, weights_only=True)["model"]
         new_state_dict = {}
         for k, v in state_dict.items():
             new_state_dict[k.replace("module.", "")] = v

deep_atlas/deep_atlas_tutorial.ipynb

@@ -1347,7 +1347,7 @@
    "outputs": [],
    "source": [
     "# CHECKPOINT CELL; LOAD\n",
-    "# seg_net.load_state_dict(torch.load('seg_net_pretrained.pth'))"
+    "# seg_net.load_state_dict(torch.load('seg_net_pretrained.pth', weights_only=True))"
    ]
   },
   {
@@ -2035,8 +2035,8 @@
    "outputs": [],
    "source": [
     "# CHECKPOINT CELL; LOAD\n",
-    "# seg_net.load_state_dict(torch.load('seg_net.pth'))\n",
-    "# reg_net.load_state_dict(torch.load('reg_net.pth'))"
+    "# seg_net.load_state_dict(torch.load('seg_net.pth', weights_only=True))\n",
+    "# reg_net.load_state_dict(torch.load('reg_net.pth', weights_only=True))"
    ]
   },
   {

deepedit/ignite/infoANDinference.ipynb

@@ -476,7 +476,7 @@
    "source": [
     "# Evaluation\n",
     "model_path = \"pretrained_deepedit_dynunet-final.pt\"\n",
-    "model.load_state_dict(torch.load(model_path))\n",
+    "model.load_state_dict(torch.load(model_path, weights_only=True))\n",
     "model.cuda()\n",
     "model.eval()\n",
     "\n",

deepedit/ignite/train.py

@@ -219,7 +219,7 @@ def create_trainer(args):
     if args.resume:
         logging.info("{}:: Loading Network...".format(local_rank))
         map_location = {"cuda:0": "cuda:{}".format(local_rank)}
-        network.load_state_dict(torch.load(args.model_filepath, map_location=map_location))
+        network.load_state_dict(torch.load(args.model_filepath, map_location=map_location, weights_only=True))
 
     # define event-handlers for engine
     val_handlers = [
@@ -333,7 +333,7 @@ def run(args):
         network = get_network(args.network, args.labels, args.spatial_size).to(device)
 
         map_location = {"cuda:0": "cuda:{}".format(args.local_rank)}
-        network.load_state_dict(torch.load(args.input, map_location=map_location))
+        network.load_state_dict(torch.load(args.input, map_location=map_location, weights_only=True))
 
         logging.info("{}:: Saving TorchScript Model".format(args.local_rank))
         model_ts = torch.jit.script(network)