Override fixes and general updates (#129)

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

ultralytics/yolo/engine/exporter.py

@@ -137,8 +137,6 @@ class Exporter:
         """
         if overrides is None:
             overrides = {}
-        if 'batch_size' not in overrides:
-            overrides['batch_size'] = 1  # set default export batch size
         self.args = get_config(config, overrides)
         project = self.args.project or f"runs/{self.args.task}"
         name = self.args.name or "exp"  # hardcode mode as export doesn't require it
@@ -166,6 +164,8 @@ class Exporter:
             assert not self.args.dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic'
 
         # Checks
+        if self.args.batch_size == 16:
+            self.args.batch_size = 1  # TODO: resolve batch_size 16 default in config.yaml
         self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2)  # check image size
         if self.args.optimize:
             assert self.device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'

ultralytics/yolo/engine/model.py

@@ -1,6 +1,7 @@
-import torch
 from pathlib import Path
 
+import torch
+
 from ultralytics import yolo  # noqa
 from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights
 from ultralytics.yolo.configs import get_config

ultralytics/yolo/engine/trainer.py

@@ -107,6 +107,8 @@ class BaseTrainer:
         self.device = utils.torch_utils.select_device(self.args.device, self.batch_size)
         self.amp = self.device.type != 'cpu'
         self.scaler = amp.GradScaler(enabled=self.amp)
+        if self.device.type == 'cpu':
+            self.args.workers = 0  # faster CPU training as time dominated by inference, not dataloading
 
         # Model and Dataloaders.
         self.model = self.args.model