Nightly build

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
 
 

tests/helpers/model.py

@@ -15,7 +15,7 @@
 # and limitations under the License.
 
 import os
-from typing import Dict, Tuple, Union
+from typing import Dict, List, Tuple, Union
 
 import numpy as np
 from omegaconf import DictConfig, ListConfig
@@ -36,8 +36,9 @@ def setup_model_train(
     nncf: bool,
     category: str,
     score_type: str = None,
-    weight_file: str = "weights/last.ckpt",
+    weight_file: str = "weights/model.ckpt",
     fast_run: bool = False,
+    device: Union[List[int], int] = [0],
 ) -> Tuple[Union[DictConfig, ListConfig], LightningDataModule, AnomalyModule, Trainer]:
     """Train the model based on the parameters passed.
 
@@ -51,6 +52,7 @@ def setup_model_train(
         weight_file (str, optional): Path to weight file.
         fast_run (bool, optional): If set to true, the model trains for only 1 epoch. We train for one epoch as
             this ensures that both anomalous and non-anomalous images are present in the validation step.
+        device (List[int], int, optional): Select which device you want to train the model on. Defaults to first GPU.
 
     Returns:
         Tuple[DictConfig, LightningDataModule, AnomalyModule, Trainer]: config, datamodule, trained model, trainer
@@ -62,12 +64,13 @@ def setup_model_train(
     config.dataset.category = category
     config.dataset.path = dataset_path
     config.project.log_images_to = []
+    config.trainer.gpus = device
 
     # If weight file is empty, remove the key from config
     if "weight_file" in config.model.keys() and weight_file == "":
         config.model.pop("weight_file")
     else:
-        config.model.weight_file = weight_file
+        config.model.weight_file = weight_file if not fast_run else "weights/last.ckpt"
 
     if nncf:
         config.optimization.nncf.apply = True