🐞 Fix tensor detach and gpu count issues in benchmarking script (#100)

Co-authored-by: Ashwin Vaidya <[email protected]>

Author: ashwinvaidya17

anomalib/utils/sweep/helpers/inference.py

@@ -19,6 +19,7 @@
 from typing import Dict, Iterable, List, Tuple, Union
 
 import numpy as np
+import torch
 from omegaconf import DictConfig, ListConfig
 from torch.utils.data import DataLoader
 
@@ -106,6 +107,7 @@ def get_torch_throughput(
     Returns:
         float: Inference throughput
     """
+    torch.set_grad_enabled(False)
     model.eval()
     inferencer = TorchInferencer(config, model)
     torch_dataloader = MockImageLoader(config.dataset.image_size, len(test_dataset))
@@ -118,6 +120,7 @@ def get_torch_throughput(
     inference_time = time.time() - start_time
     throughput = len(test_dataset) / inference_time
 
+    torch.set_grad_enabled(True)
     return throughput
 
 

tools/benchmarking/benchmark.py

@@ -16,6 +16,7 @@
 
 
 import logging
+import math
 import multiprocessing
 import time
 from concurrent.futures import ProcessPoolExecutor, as_completed
@@ -136,12 +137,12 @@ def distribute_over_gpus():
         run_configs = list(get_run_config(sweep_config.grid_search))
         jobs = []
         for device_id, run_split in enumerate(
-            range(0, len(run_configs), len(run_configs) // torch.cuda.device_count())
+            range(0, len(run_configs), math.ceil(len(run_configs) / torch.cuda.device_count()))
         ):
             jobs.append(
                 executor.submit(
                     compute_on_gpu,
-                    run_configs[run_split : run_split + len(run_configs) // torch.cuda.device_count()],
+                    run_configs[run_split : run_split + math.ceil(len(run_configs) / torch.cuda.device_count())],
                     device_id + 1,
                     sweep_config.seed,
                     sweep_config.writer,