standalone val (#56)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/yolo/engine/exporter.py

@@ -0,0 +1,19 @@
+import pandas as pd
+
+
+def export_formats():
+    # YOLOv5 export formats
+    x = [
+        ['PyTorch', '-', '.pt', True, True],
+        ['TorchScript', 'torchscript', '.torchscript', True, True],
+        ['ONNX', 'onnx', '.onnx', True, True],
+        ['OpenVINO', 'openvino', '_openvino_model', True, False],
+        ['TensorRT', 'engine', '.engine', False, True],
+        ['CoreML', 'coreml', '.mlmodel', True, False],
+        ['TensorFlow SavedModel', 'saved_model', '_saved_model', True, True],
+        ['TensorFlow GraphDef', 'pb', '.pb', True, True],
+        ['TensorFlow Lite', 'tflite', '.tflite', True, False],
+        ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False],
+        ['TensorFlow.js', 'tfjs', '_web_model', False, False],
+        ['PaddlePaddle', 'paddle', '_paddle_model', True, True],]
+    return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU'])

ultralytics/yolo/engine/trainer.py

@@ -25,7 +25,7 @@ import ultralytics.yolo.utils as utils
 import ultralytics.yolo.utils.callbacks as callbacks
 from ultralytics.yolo.data.utils import check_dataset, check_dataset_yaml
 from ultralytics.yolo.utils import LOGGER, ROOT, TQDM_BAR_FORMAT, colorstr
-from ultralytics.yolo.utils.checks import print_args
+from ultralytics.yolo.utils.checks import check_file, print_args
 from ultralytics.yolo.utils.files import increment_path, save_yaml
 from ultralytics.yolo.utils.modeling import get_model
 from ultralytics.yolo.utils.torch_utils import ModelEMA, de_parallel, init_seeds, one_cycle, strip_optimizer
@@ -299,13 +299,16 @@ class BaseTrainer:
         """
         Get train, val path from data dict if it exists. Returns None if data format is not recognized
         """
-        return data["train"], data["val"]
+        return data["train"], data.get("val") or data.get("test")
 
     def get_model(self, model: Union[str, Path]):
         """
         load/create/download model for any task
         """
-        pretrained = not str(model).endswith(".yaml")
+        pretrained = True
+        if str(model).endswith(".yaml"):
+            model = check_file(model)
+            pretrained = False
         return self.load_model(model_cfg=None if pretrained else model,
                                weights=get_model(model) if pretrained else None,
                                data=self.data)  # model
@@ -376,7 +379,7 @@ class BaseTrainer:
         """
         To set or update model parameters before training.
         """
-        pass
+        self.model.names = self.data["names"]
 
     def build_targets(self, preds, targets):
         pass

ultralytics/yolo/engine/validator.py

@@ -5,11 +5,14 @@ import torch
 from omegaconf import OmegaConf
 from tqdm import tqdm
 
+from ultralytics.yolo.data.utils import check_dataset, check_dataset_yaml
 from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG
-from ultralytics.yolo.utils import TQDM_BAR_FORMAT
+from ultralytics.yolo.utils import LOGGER, TQDM_BAR_FORMAT
 from ultralytics.yolo.utils.files import increment_path
+from ultralytics.yolo.utils.modeling import get_model
+from ultralytics.yolo.utils.modeling.autobackend import AutoBackend
 from ultralytics.yolo.utils.ops import Profile
-from ultralytics.yolo.utils.torch_utils import de_parallel, select_device
+from ultralytics.yolo.utils.torch_utils import check_img_size, de_parallel, select_device
 
 
 class BaseValidator:
@@ -17,17 +20,18 @@ class BaseValidator:
     Base validator class.
     """
 
-    def __init__(self, dataloader, save_dir=None, pbar=None, logger=None, args=None):
+    def __init__(self, dataloader=None, save_dir=None, pbar=None, logger=None, args=None):
         self.dataloader = dataloader
         self.pbar = pbar
-        self.logger = logger or logging.getLogger()
+        self.logger = logger or LOGGER
         self.args = args or OmegaConf.load(DEFAULT_CONFIG)
-        self.device = select_device(self.args.device, dataloader.batch_size)
-        self.save_dir = save_dir if save_dir is not None else \
-                increment_path(Path(self.args.project) / self.args.name, exist_ok=self.args.exist_ok)
-        self.cuda = self.device.type != 'cpu'
+        self.model = None
+        self.data = None
+        self.device = None
         self.batch_i = None
         self.training = True
+        self.save_dir = save_dir if save_dir is not None else \
+                increment_path(Path(self.args.project) / self.args.name, exist_ok=self.args.exist_ok)
 
     def __call__(self, trainer=None, model=None):
         """
@@ -36,14 +40,35 @@ class BaseValidator:
         """
         self.training = trainer is not None
         if self.training:
+            self.device = trainer.device
+            self.data = trainer.data
             model = trainer.ema.ema or trainer.model
             self.args.half &= self.device.type != 'cpu'
             model = model.half() if self.args.half else model.float()
+            self.model = model
             loss = torch.zeros_like(trainer.loss_items, device=trainer.device)
         else:  # TODO: handle this when detectMultiBackend is supported
             assert model is not None, "Either trainer or model is needed for validation"
-            # model = DetectMultiBacked(model)
-            # TODO: implement init_model_attributes()
+            self.device = select_device(self.args.device, self.args.batch_size)
+            self.args.half &= self.device.type != 'cpu'
+            model = AutoBackend(model, device=self.device, dnn=self.args.dnn, fp16=self.args.half)
+            self.model = model
+            stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
+            imgsz = check_img_size(self.args.img_size, s=stride)
+            if engine:
+                self.args.batch_size = model.batch_size
+            else:
+                self.device = model.device
+                if not (pt or jit):
+                    self.args.batch_size = 1  # export.py models default to batch-size 1
+                    self.logger.info(
+                        f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
+
+            if self.args.data.endswith(".yaml"):
+                data = check_dataset_yaml(self.args.data)
+            else:
+                data = check_dataset(self.args.data)
+            self.dataloader = self.get_dataloader(data.get("val") or data.set("test"), self.args.batch_size)
 
         model.eval()
 
@@ -101,6 +126,9 @@ class BaseValidator:
         return stats | trainer.label_loss_items(loss.cpu() / len(self.dataloader), prefix="val") \
                 if self.training else stats
 
+    def get_dataloader(self, dataset_path, batch_size):
+        raise Exception("get_dataloder function not implemented for this validator")
+
     def preprocess(self, batch):
         return batch