Update no_grad usage to inference_mode (#1164)

Author: GuillaumeErhard

examples/binary_segmentation_buildings.py

@@ -269,7 +269,7 @@ def train_and_evaluate_one_epoch(
     # Set the model to evaluation mode
     model.eval()
     val_loss = 0
-    with torch.no_grad():
+    with torch.inference_mode():
         for batch in tqdm(valid_dataloader, desc="Evaluating"):
             images, masks = batch
             images, masks = images.to(device), masks.to(device)
@@ -325,7 +325,7 @@ def test_model(model, output_dir, test_dataloader, loss_fn, device):
     model.eval()
     test_loss = 0
     tp, fp, fn, tn = 0, 0, 0, 0
-    with torch.no_grad():
+    with torch.inference_mode():
         for batch in tqdm(test_dataloader, desc="Evaluating"):
             images, masks = batch
             images, masks = images.to(device), masks.to(device)

examples/binary_segmentation_intro.ipynb

@@ -1026,7 +1026,7 @@
       ],
       "source": [
         "batch = next(iter(test_dataloader))\n",
-        "with torch.no_grad():\n",
+        "with torch.inference_mode():\n",
         "    model.eval()\n",
         "    logits = model(batch[\"image\"])\n",
         "pr_masks = logits.sigmoid()\n",

examples/camvid_segmentation_multiclass.ipynb

@@ -1683,7 +1683,7 @@
     "images, masks = next(iter(test_loader))\n",
     "\n",
     "# Switch the model to evaluation mode\n",
-    "with torch.no_grad():\n",
+    "with torch.inference_mode():\n",
     "    model.eval()\n",
     "    logits = model(images)  # Get raw logits from the model\n",
     "\n",

examples/cars segmentation (camvid).ipynb

@@ -1209,7 +1209,7 @@
    ],
    "source": [
     "images, masks = next(iter(test_loader))\n",
-    "with torch.no_grad():\n",
+    "with torch.inference_mode():\n",
     "    model.eval()\n",
     "    logits = model(images)\n",
     "pr_masks = logits.sigmoid()\n",

examples/convert_to_onnx.ipynb

@@ -189,7 +189,7 @@
    ],
    "source": [
     "# compute PyTorch output prediction\n",
-    "with torch.no_grad():\n",
+    "with torch.inference_mode():\n",
     "    torch_out = model(sample)\n",
     "\n",
     "# compare ONNX Runtime and PyTorch results\n",

examples/dpt_inference_pretrained.ipynb

@@ -70,7 +70,7 @@
     "input_tensor = input_tensor.to(device)\n",
     "\n",
     "# Perform inference\n",
-    "with torch.no_grad():\n",
+    "with torch.inference_mode():\n",
     "    output_mask = model(input_tensor)\n",
     "\n",
     "# Postprocess mask\n",

examples/segformer_inference_pretrained.ipynb

@@ -63,7 +63,7 @@
     "input_tensor = input_tensor.to(device)\n",
     "\n",
     "# Perform inference\n",
-    "with torch.no_grad():\n",
+    "with torch.inference_mode():\n",
     "    output_mask = model(input_tensor)\n",
     "\n",
     "# Postprocess mask\n",

examples/upernet_inference_pretrained.ipynb

@@ -85,7 +85,7 @@
     "input_tensor = input_tensor.to(device)\n",
     "\n",
     "# Perform inference\n",
-    "with torch.no_grad():\n",
+    "with torch.inference_mode():\n",
     "    output_mask = model(input_tensor)\n",
     "\n",
     "# Postprocess mask\n",

misc/generate_test_models.py

@@ -52,7 +52,7 @@ def save_and_push(model, inputs, outputs, model_name, encoder_name):
     torch.manual_seed(423553)
     sample = torch.rand(1, 3, 256, 256)
 
-    with torch.no_grad():
+    with torch.inference_mode():
         output = model(sample)
 
     save_and_push(model, sample, output, model_name, encoder_name)

scripts/models-conversions/segformer-original-decoder-to-smp.py

@@ -107,7 +107,7 @@ def main(args):
     tensor = torch.tensor(normalized_image).permute(2, 0, 1).unsqueeze(0).float()
 
     # Forward pass
-    with torch.no_grad():
+    with torch.inference_mode():
         mask = model(tensor)
 
     # Postprocessing