Pip debug fixes (#139)

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

ultralytics/nn/tasks.py

@@ -292,7 +292,8 @@ def attempt_load_weights(weights, device=None, inplace=True, fuse=False):
         ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float()  # FP32 model
 
         # Model compatibility updates
-        ckpt.args = {k: v for k, v in args.items() if k in DEFAULT_CONFIG_KEYS}
+        ckpt.args = {k: v for k, v in args.items() if k in DEFAULT_CONFIG_KEYS}  # attach args to model
+        ckpt.pt_path = weights  # attach *.pt file path to model
         if not hasattr(ckpt, 'stride'):
             ckpt.stride = torch.tensor([32.])
 

ultralytics/yolo/engine/exporter.py

@@ -156,7 +156,7 @@ class Exporter:
         jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags  # export booleans
 
         # Load PyTorch model
-        self.device = select_device(self.args.device)
+        self.device = select_device(self.args.device or 'cpu')
         if self.args.half:
             if self.device.type == 'cpu' or not coreml:
                 LOGGER.info('half=True only compatible with GPU or CoreML export, i.e. use device=0 or format=coreml')
@@ -172,7 +172,9 @@ class Exporter:
 
         # Input
         im = torch.zeros(self.args.batch_size, 3, *self.imgsz).to(self.device)
-        file = Path(getattr(model, 'yaml_file', None) or Path(model.yaml['yaml_file']).name)
+        file = Path(getattr(model, 'pt_path', None) or model.yaml['yaml_file'])
+        if file.suffix == '.yaml':
+            file = Path(file.name)
 
         # Update model
         model = deepcopy(model)

ultralytics/yolo/engine/model.py

@@ -213,9 +213,8 @@ class YOLO:
 
     @smart_inference_mode()
     def __call__(self, imgs):
-        if not self.model:
-            LOGGER.info("model not initialized!")
-        return self.model(imgs)
+        device = next(self.model.parameters()).device  # get model device
+        return self.model(imgs.to(device))
 
     def forward(self, imgs):
         return self.__call__(imgs)

ultralytics/yolo/engine/trainer.py

@@ -81,12 +81,11 @@ class BaseTrainer:
             overrides = {}
         self.args = get_config(config, overrides)
         self.check_resume()
-        init_seeds(self.args.seed + 1 + RANK, deterministic=self.args.deterministic)
-
         self.console = LOGGER
         self.validator = None
         self.model = None
         self.callbacks = defaultdict(list)
+        init_seeds(self.args.seed + 1 + RANK, deterministic=self.args.deterministic)
 
         # Dirs
         project = self.args.project or f"runs/{self.args.task}"

ultralytics/yolo/utils/__init__.py

@@ -62,6 +62,7 @@ HELP_MSG = \
 pd.options.display.max_columns = 10
 cv2.setNumThreads(0)  # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
 os.environ['NUMEXPR_MAX_THREADS'] = str(NUM_THREADS)  # NumExpr max threads
+os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'  # for deterministic training
 
 # Default config dictionary
 with open(DEFAULT_CONFIG, errors='ignore') as f:

ultralytics/yolo/v8/detect/train.py

@@ -189,10 +189,11 @@ class Loss:
 def train(cfg):
     cfg.model = cfg.model or "yolov8n.yaml"
     cfg.data = cfg.data or "coco128.yaml"  # or yolo.ClassificationDataset("mnist")
-    # cfg.imgsz = 160
-    # cfg.epochs = 5
-    trainer = DetectionTrainer(cfg)
-    trainer.train()
+    # trainer = DetectionTrainer(cfg)
+    # trainer.train()
+    from ultralytics import YOLO
+    model = YOLO(cfg.model)
+    model.train(**cfg)
 
 
 if __name__ == "__main__":

ultralytics/yolo/v8/segment/train.py

@@ -176,8 +176,11 @@ class SegLoss:
 def train(cfg):
     cfg.model = cfg.model or "yolov8n-seg.yaml"
     cfg.data = cfg.data or "coco128-seg.yaml"  # or yolo.ClassificationDataset("mnist")
-    trainer = SegmentationTrainer(cfg)
-    trainer.train()
+    # trainer = SegmentationTrainer(cfg)
+    # trainer.train()
+    from ultralytics import YOLO
+    model = YOLO(cfg.model)
+    model.train(**cfg)
 
 
 if __name__ == "__main__":