diff --git a/tests/tests.py b/tests/tests.py index 989d296..e7a22fe 100644 --- a/tests/tests.py +++ b/tests/tests.py @@ -3,56 +3,50 @@ import torch from ultralytics import YOLO from ultralytics.yolo.utils import ROOT +MODEL = ROOT / 'weights/yolov8n.pt' +CFG = 'yolov8n.yaml' + def test_model_forward(): - model = YOLO("yolov8n.yaml") + model = YOLO(CFG) img = torch.rand(1, 3, 320, 320) model.forward(img) model(img) def test_model_info(): - model = YOLO("yolov8n.yaml") + model = YOLO(CFG) model.info() - model = YOLO("yolov8n.pt") + model = YOLO(MODEL) model.info(verbose=True) def test_model_fuse(): - model = YOLO("yolov8n.yaml") + model = YOLO(CFG) model.fuse() - model = YOLO("yolov8n.pt") + model = YOLO(MODEL) model.fuse() def test_predict_dir(): - model = YOLO("yolov8n.pt") + model = YOLO(MODEL) model.predict(source=ROOT / "assets") def test_val(): - model = YOLO("yolov8n.pt") + model = YOLO(MODEL) model.val(data="coco128.yaml", imgsz=32) -def test_train_resume(): - model = YOLO("yolov8n.yaml") - model.train(epochs=1, imgsz=32, data="coco128.yaml") - try: - model.resume(task="detect") - except AssertionError: - print("Successfully caught resume assert!") - - def test_train_scratch(): - model = YOLO("yolov8n.yaml") + model = YOLO(CFG) model.train(data="coco128.yaml", epochs=1, imgsz=32) img = torch.rand(1, 3, 320, 320) model(img) def test_train_pretrained(): - model = YOLO("yolov8n.pt") + model = YOLO(MODEL) model.train(data="coco128.yaml", epochs=1, imgsz=32) img = torch.rand(1, 3, 320, 320) model(img) @@ -77,27 +71,27 @@ def test_export_torchscript(): from ultralytics.yolo.engine.exporter import export_formats print(export_formats()) - model = YOLO("yolov8n.yaml") + model = YOLO(MODEL) model.export(format='torchscript') def test_export_onnx(): - model = YOLO("yolov8n.yaml") + model = YOLO(MODEL) model.export(format='onnx') def test_export_openvino(): - model = YOLO("yolov8n.yaml") + model = YOLO(MODEL) model.export(format='openvino') def test_export_coreml(): - model = YOLO("yolov8n.yaml") + model = YOLO(MODEL) model.export(format='coreml') def test_export_paddle(): - model = YOLO("yolov8n.yaml") + model = YOLO(MODEL) model.export(format='paddle') diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index 38063a6..f1b4fcb 100644 --- a/ultralytics/nn/tasks.py +++ b/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.]) diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py index 857a46c..5765a9a 100644 --- a/ultralytics/yolo/engine/exporter.py +++ b/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) diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 31ce376..e1b3281 100644 --- a/ultralytics/yolo/engine/model.py +++ b/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) diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index 4006ca9..7b8248b 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/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}" diff --git a/ultralytics/yolo/utils/__init__.py b/ultralytics/yolo/utils/__init__.py index 37d12e2..be0b749 100644 --- a/ultralytics/yolo/utils/__init__.py +++ b/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: diff --git a/ultralytics/yolo/v8/detect/train.py b/ultralytics/yolo/v8/detect/train.py index 19a357c..f4aa037 100644 --- a/ultralytics/yolo/v8/detect/train.py +++ b/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__": diff --git a/ultralytics/yolo/v8/segment/train.py b/ultralytics/yolo/v8/segment/train.py index 529cd29..6445355 100644 --- a/ultralytics/yolo/v8/segment/train.py +++ b/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__":