From 701fba47700dee7a6bd8ff4d219eca91adc9d3be Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 16 Mar 2023 15:42:44 +0100 Subject: [PATCH] `ultralytics 8.0.54` TFLite export improvements and fixes (#1447) Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- README.md | 2 +- README.zh-CN.md | 2 +- docs/modes/export.md | 6 ++-- docs/modes/track.md | 10 +++--- docs/modes/train.md | 14 ++++++-- docs/modes/val.md | 4 +-- docs/quickstart.md | 10 +++--- docs/tasks/classify.md | 26 +++++++------- docs/tasks/detect.md | 26 +++++++------- docs/tasks/keypoints.md | 26 +++++++------- docs/tasks/segment.md | 26 +++++++------- docs/usage/callbacks.md | 2 +- docs/usage/cli.md | 4 +-- docs/usage/python.md | 12 +++---- requirements.txt | 2 +- tests/test_python.py | 4 +-- ultralytics/__init__.py | 2 +- ultralytics/nn/modules.py | 6 ++-- ultralytics/nn/tasks.py | 34 +++++++----------- ultralytics/yolo/cfg/__init__.py | 11 +++--- ultralytics/yolo/engine/exporter.py | 2 +- ultralytics/yolo/engine/model.py | 30 +++++++++------- ultralytics/yolo/engine/results.py | 18 ++++++---- ultralytics/yolo/utils/__init__.py | 6 ++-- ultralytics/yolo/utils/autobatch.py | 35 ++++++++++++++----- .../yolo/utils/callbacks/tensorboard.py | 15 ++++---- ultralytics/yolo/utils/checks.py | 22 ++++++------ ultralytics/yolo/v8/detect/predict.py | 2 +- ultralytics/yolo/v8/detect/train.py | 3 +- ultralytics/yolo/v8/segment/predict.py | 2 +- 30 files changed, 198 insertions(+), 166 deletions(-) diff --git a/README.md b/README.md index 782a58e..86a4232 100644 --- a/README.md +++ b/README.md @@ -79,7 +79,7 @@ pip install ultralytics YOLOv8 may be used directly in the Command Line Interface (CLI) with a `yolo` command: ```bash -yolo predict model=yolov8n.pt source="https://ultralytics.com/images/bus.jpg" +yolo predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg' ``` `yolo` can be used for a variety of tasks and modes and accepts additional arguments, i.e. `imgsz=640`. See the YOLOv8 diff --git a/README.zh-CN.md b/README.zh-CN.md index 6f06f4e..9c6d554 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -71,7 +71,7 @@ pip install ultralytics YOLOv8 可以直接在命令行界面（CLI）中使用 `yolo` 命令运行： ```bash -yolo predict model=yolov8n.pt source="https://ultralytics.com/images/bus.jpg" +yolo predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg' ``` `yolo`可以用于各种任务和模式，并接受额外的参数，例如 `imgsz=640`。参见 YOLOv8 [文档](https://docs.ultralytics.com) diff --git a/docs/modes/export.md b/docs/modes/export.md index 3c4deed..f454466 100644 --- a/docs/modes/export.md +++ b/docs/modes/export.md @@ -22,11 +22,11 @@ export arguments. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom trained + model = YOLO('yolov8n.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom trained # Export the model - model.export(format="onnx") + model.export(format='onnx') ``` === "CLI" diff --git a/docs/modes/track.md b/docs/modes/track.md index 0fe3c3e..8058f38 100644 --- a/docs/modes/track.md +++ b/docs/modes/track.md @@ -26,9 +26,9 @@ Use a trained YOLOv8n/YOLOv8n-seg model to run tracker on video streams. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.pt") # load an official detection model - model = YOLO("yolov8n-seg.pt") # load an official segmentation model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n.pt') # load an official detection model + model = YOLO('yolov8n-seg.pt') # load an official segmentation model + model = YOLO('path/to/best.pt') # load a custom model # Track with the model results = model.track(source="https://youtu.be/Zgi9g1ksQHc", show=True) @@ -60,7 +60,7 @@ to [predict page](https://docs.ultralytics.com/modes/predict/). ```python from ultralytics import YOLO - model = YOLO("yolov8n.pt") + model = YOLO('yolov8n.pt') results = model.track(source="https://youtu.be/Zgi9g1ksQHc", conf=0.3, iou=0.5, show=True) ``` === "CLI" @@ -82,7 +82,7 @@ any configurations(expect the `tracker_type`) you need to. ```python from ultralytics import YOLO - model = YOLO("yolov8n.pt") + model = YOLO('yolov8n.pt') results = model.track(source="https://youtu.be/Zgi9g1ksQHc", tracker='custom_tracker.yaml') ``` === "CLI" diff --git a/docs/modes/train.md b/docs/modes/train.md index dc1f8a6..8716250 100644 --- a/docs/modes/train.md +++ b/docs/modes/train.md @@ -21,16 +21,24 @@ training arguments. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.yaml") # build a new model from scratch - model = YOLO("yolov8n.pt") # load a pretrained model (recommended for training) + model = YOLO('yolov8n.yaml') # build a new model from YAML + model = YOLO('yolov8n.pt') # load a pretrained model (recommended for training) + model = YOLO('yolov8n.yaml').load('yolov8n.pt') # build from YAML and transfer weights # Train the model - model.train(data="coco128.yaml", epochs=100, imgsz=640) + model.train(data='coco128.yaml', epochs=100, imgsz=640) ``` === "CLI" ```bash + # Build a new model from YAML and start training from scratch + yolo detect train data=coco128.yaml model=yolov8n.yaml epochs=100 imgsz=640 + + # Start training from a pretrained *.pt model yolo detect train data=coco128.yaml model=yolov8n.pt epochs=100 imgsz=640 + + # Build a new model from YAML, transfer pretrained weights to it and start training + yolo detect train data=coco128.yaml model=yolov8n.yaml pretrained=yolov8n.pt epochs=100 imgsz=640 ``` ## Arguments diff --git a/docs/modes/val.md b/docs/modes/val.md index be4175b..da5e7f1 100644 --- a/docs/modes/val.md +++ b/docs/modes/val.md @@ -21,8 +21,8 @@ training `data` and arguments as model attributes. See Arguments section below f from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom model # Validate the model metrics = model.val() # no arguments needed, dataset and settings remembered diff --git a/docs/quickstart.md b/docs/quickstart.md index 40777b1..3eb4443 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -60,14 +60,14 @@ classification into their Python projects using YOLOv8. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.yaml") # build a new model from scratch - model = YOLO("yolov8n.pt") # load a pretrained model (recommended for training) + model = YOLO('yolov8n.yaml') # build a new model from scratch + model = YOLO('yolov8n.pt') # load a pretrained model (recommended for training) # Use the model - results = model.train(data="coco128.yaml", epochs=3) # train the model + results = model.train(data='coco128.yaml', epochs=3) # train the model results = model.val() # evaluate model performance on the validation set - results = model("https://ultralytics.com/images/bus.jpg") # predict on an image - success = model.export(format="onnx") # export the model to ONNX format + results = model('https://ultralytics.com/images/bus.jpg') # predict on an image + success = model.export(format='onnx') # export the model to ONNX format ``` [Python Guide](usage/python.md){.md-button .md-button--primary} diff --git a/docs/tasks/classify.md b/docs/tasks/classify.md index f0d518c..9c25a4b 100644 --- a/docs/tasks/classify.md +++ b/docs/tasks/classify.md @@ -26,11 +26,11 @@ see the [Configuration](../usage/cfg.md) page. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n-cls.yaml") # build a new model from scratch - model = YOLO("yolov8n-cls.pt") # load a pretrained model (recommended for training) + model = YOLO('yolov8n-cls.yaml') # build a new model from scratch + model = YOLO('yolov8n-cls.pt') # load a pretrained model (recommended for training) # Train the model - model.train(data="mnist160", epochs=100, imgsz=64) + model.train(data='mnist160', epochs=100, imgsz=64) ``` === "CLI" @@ -51,8 +51,8 @@ it's training `data` and arguments as model attributes. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n-cls.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n-cls.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom model # Validate the model metrics = model.val() # no arguments needed, dataset and settings remembered @@ -78,17 +78,17 @@ Use a trained YOLOv8n-cls model to run predictions on images. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n-cls.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n-cls.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom model # Predict with the model - results = model("https://ultralytics.com/images/bus.jpg") # predict on an image + results = model('https://ultralytics.com/images/bus.jpg') # predict on an image ``` === "CLI" ```bash - yolo classify predict model=yolov8n-cls.pt source="https://ultralytics.com/images/bus.jpg" # predict with official model - yolo classify predict model=path/to/best.pt source="https://ultralytics.com/images/bus.jpg" # predict with custom model + yolo classify predict model=yolov8n-cls.pt source='https://ultralytics.com/images/bus.jpg' # predict with official model + yolo classify predict model=path/to/best.pt source='https://ultralytics.com/images/bus.jpg' # predict with custom model ``` Read more details of `predict` in our [Predict](https://docs.ultralytics.com/modes/predict/) page. @@ -105,11 +105,11 @@ Export a YOLOv8n-cls model to a different format like ONNX, CoreML, etc. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n-cls.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom trained + model = YOLO('yolov8n-cls.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom trained # Export the model - model.export(format="onnx") + model.export(format='onnx') ``` === "CLI" diff --git a/docs/tasks/detect.md b/docs/tasks/detect.md index 442c2f3..cf13b23 100644 --- a/docs/tasks/detect.md +++ b/docs/tasks/detect.md @@ -26,11 +26,11 @@ the [Configuration](../usage/cfg.md) page. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.yaml") # build a new model from scratch - model = YOLO("yolov8n.pt") # load a pretrained model (recommended for training) + model = YOLO('yolov8n.yaml') # build a new model from scratch + model = YOLO('yolov8n.pt') # load a pretrained model (recommended for training) # Train the model - model.train(data="coco128.yaml", epochs=100, imgsz=640) + model.train(data='coco128.yaml', epochs=100, imgsz=640) ``` === "CLI" @@ -51,8 +51,8 @@ training `data` and arguments as model attributes. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom model # Validate the model metrics = model.val() # no arguments needed, dataset and settings remembered @@ -80,17 +80,17 @@ Use a trained YOLOv8n model to run predictions on images. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom model # Predict with the model - results = model("https://ultralytics.com/images/bus.jpg") # predict on an image + results = model('https://ultralytics.com/images/bus.jpg') # predict on an image ``` === "CLI" ```bash - yolo detect predict model=yolov8n.pt source="https://ultralytics.com/images/bus.jpg" # predict with official model - yolo detect predict model=path/to/best.pt source="https://ultralytics.com/images/bus.jpg" # predict with custom model + yolo detect predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg' # predict with official model + yolo detect predict model=path/to/best.pt source='https://ultralytics.com/images/bus.jpg' # predict with custom model ``` Read more details of `predict` in our [Predict](https://docs.ultralytics.com/modes/predict/) page. @@ -107,11 +107,11 @@ Export a YOLOv8n model to a different format like ONNX, CoreML, etc. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom trained + model = YOLO('yolov8n.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom trained # Export the model - model.export(format="onnx") + model.export(format='onnx') ``` === "CLI" diff --git a/docs/tasks/keypoints.md b/docs/tasks/keypoints.md index 14cd9f5..60a688f 100644 --- a/docs/tasks/keypoints.md +++ b/docs/tasks/keypoints.md @@ -28,11 +28,11 @@ train an OpenPose model on a custom dataset, see the OpenPose Training page. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.yaml") # build a new model from scratch - model = YOLO("yolov8n.pt") # load a pretrained model (recommended for training) + model = YOLO('yolov8n.yaml') # build a new model from scratch + model = YOLO('yolov8n.pt') # load a pretrained model (recommended for training) # Train the model - model.train(data="coco128.yaml", epochs=100, imgsz=640) + model.train(data='coco128.yaml', epochs=100, imgsz=640) ``` === "CLI" @@ -53,8 +53,8 @@ training `data` and arguments as model attributes. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom model # Validate the model metrics = model.val() # no arguments needed, dataset and settings remembered @@ -82,17 +82,17 @@ Use a trained YOLOv8n model to run predictions on images. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom model # Predict with the model - results = model("https://ultralytics.com/images/bus.jpg") # predict on an image + results = model('https://ultralytics.com/images/bus.jpg') # predict on an image ``` === "CLI" ```bash - yolo detect predict model=yolov8n.pt source="https://ultralytics.com/images/bus.jpg" # predict with official model - yolo detect predict model=path/to/best.pt source="https://ultralytics.com/images/bus.jpg" # predict with custom model + yolo detect predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg' # predict with official model + yolo detect predict model=path/to/best.pt source='https://ultralytics.com/images/bus.jpg' # predict with custom model ``` Read more details of `predict` in our [Predict](https://docs.ultralytics.com/modes/predict/) page. @@ -109,11 +109,11 @@ Export a YOLOv8n model to a different format like ONNX, CoreML, etc. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom trained + model = YOLO('yolov8n.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom trained # Export the model - model.export(format="onnx") + model.export(format='onnx') ``` === "CLI" diff --git a/docs/tasks/segment.md b/docs/tasks/segment.md index 47a03d5..f89c3fe 100644 --- a/docs/tasks/segment.md +++ b/docs/tasks/segment.md @@ -26,11 +26,11 @@ arguments see the [Configuration](../usage/cfg.md) page. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n-seg.yaml") # build a new model from scratch - model = YOLO("yolov8n-seg.pt") # load a pretrained model (recommended for training) + model = YOLO('yolov8n-seg.yaml') # build a new model from scratch + model = YOLO('yolov8n-seg.pt') # load a pretrained model (recommended for training) # Train the model - model.train(data="coco128-seg.yaml", epochs=100, imgsz=640) + model.train(data='coco128-seg.yaml', epochs=100, imgsz=640) ``` === "CLI" @@ -51,8 +51,8 @@ retains it's training `data` and arguments as model attributes. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n-seg.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n-seg.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom model # Validate the model metrics = model.val() # no arguments needed, dataset and settings remembered @@ -84,17 +84,17 @@ Use a trained YOLOv8n-seg model to run predictions on images. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n-seg.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom model + model = YOLO('yolov8n-seg.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom model # Predict with the model - results = model("https://ultralytics.com/images/bus.jpg") # predict on an image + results = model('https://ultralytics.com/images/bus.jpg') # predict on an image ``` === "CLI" ```bash - yolo segment predict model=yolov8n-seg.pt source="https://ultralytics.com/images/bus.jpg" # predict with official model - yolo segment predict model=path/to/best.pt source="https://ultralytics.com/images/bus.jpg" # predict with custom model + yolo segment predict model=yolov8n-seg.pt source='https://ultralytics.com/images/bus.jpg' # predict with official model + yolo segment predict model=path/to/best.pt source='https://ultralytics.com/images/bus.jpg' # predict with custom model ``` Read more details of `predict` in our [Predict](https://docs.ultralytics.com/modes/predict/) page. @@ -111,11 +111,11 @@ Export a YOLOv8n-seg model to a different format like ONNX, CoreML, etc. from ultralytics import YOLO # Load a model - model = YOLO("yolov8n-seg.pt") # load an official model - model = YOLO("path/to/best.pt") # load a custom trained + model = YOLO('yolov8n-seg.pt') # load an official model + model = YOLO('path/to/best.pt') # load a custom trained # Export the model - model.export(format="onnx") + model.export(format='onnx') ``` === "CLI" diff --git a/docs/usage/callbacks.md b/docs/usage/callbacks.md index 4a3834f..b505df2 100644 --- a/docs/usage/callbacks.md +++ b/docs/usage/callbacks.md @@ -17,7 +17,7 @@ def on_predict_batch_end(predictor): im0s = im0s if isinstance(im0s, list) else [im0s] predictor.results = zip(predictor.results, im0s) -model = YOLO(f"yolov8n.pt") +model = YOLO(f'yolov8n.pt') model.add_callback("on_predict_batch_end", on_predict_batch_end) for (result, frame) in model.track/predict(): pass diff --git a/docs/usage/cli.md b/docs/usage/cli.md index 0c32246..20ada7c 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -59,8 +59,8 @@ Use a trained YOLOv8n model to run predictions on images. !!! example "" ```bash - yolo detect predict model=yolov8n.pt source="https://ultralytics.com/images/bus.jpg" # predict with official model - yolo detect predict model=path/to/best.pt source="https://ultralytics.com/images/bus.jpg" # predict with custom model + yolo detect predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg' # predict with official model + yolo detect predict model=path/to/best.pt source='https://ultralytics.com/images/bus.jpg' # predict with custom model ``` ## Export diff --git a/docs/usage/python.md b/docs/usage/python.md index 10ddef9..60ef051 100644 --- a/docs/usage/python.md +++ b/docs/usage/python.md @@ -6,7 +6,7 @@ The simplest way of simply using YOLOv8 directly in a Python environment. ```python from ultralytics import YOLO - model = YOLO("yolov8n.pt") # pass any model type + model = YOLO('yolov8n.pt') # pass any model type model.train(epochs=5) ``` @@ -14,8 +14,8 @@ The simplest way of simply using YOLOv8 directly in a Python environment. ```python from ultralytics import YOLO - model = YOLO("yolov8n.yaml") - model.train(data="coco128.yaml", epochs=5) + model = YOLO('yolov8n.yaml') + model.train(data='coco128.yaml', epochs=5) ``` === "Resume" @@ -31,8 +31,8 @@ The simplest way of simply using YOLOv8 directly in a Python environment. ```python from ultralytics import YOLO - model = YOLO("yolov8n.yaml") - model.train(data="coco128.yaml", epochs=5) + model = YOLO('yolov8n.yaml') + model.train(data='coco128.yaml', epochs=5) model.val() # It'll automatically evaluate the data you trained. ``` @@ -44,7 +44,7 @@ The simplest way of simply using YOLOv8 directly in a Python environment. # It'll use the data yaml file in model.pt if you don't set data. model.val() # or you can set the data you want to val - model.val(data="coco128.yaml") + model.val(data='coco128.yaml') ``` !!! example "Predict" diff --git a/requirements.txt b/requirements.txt index 98480e6..099fdcb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ # Base ---------------------------------------- matplotlib>=3.2.2 -numpy>=1.18.5 +numpy>=1.21.6 opencv-python>=4.6.0 Pillow>=7.1.2 PyYAML>=5.3.1 diff --git a/tests/test_python.py b/tests/test_python.py index 3ec0bf7..22446a9 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -207,9 +207,9 @@ def test_predict_callback_and_setup(): def test_result(): model = YOLO('yolov8n-seg.pt') res = model([SOURCE, SOURCE]) - res[0].cpu().numpy() res[0].plot(show_conf=False) - print(res[0].path) + res[0] = res[0].cpu().numpy() + print(res[0].path, res[0].masks.masks) model = YOLO('yolov8n.pt') res = model(SOURCE) diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index dfddfe9..aceaf5d 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -1,6 +1,6 @@ # Ultralytics YOLO 🚀, GPL-3.0 license -__version__ = '8.0.53' +__version__ = '8.0.54' from ultralytics.yolo.engine.model import YOLO from ultralytics.yolo.utils.checks import check_yolo as checks diff --git a/ultralytics/nn/modules.py b/ultralytics/nn/modules.py index fcc599f..ddf0085 100644 --- a/ultralytics/nn/modules.py +++ b/ultralytics/nn/modules.py @@ -411,12 +411,12 @@ class Detect(nn.Module): self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5)) self.shape = shape - if self.export and self.format == 'edgetpu': # FlexSplitV ops issue - x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2) + x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2) + if self.export and self.format in ('saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs'): # avoid TF FlexSplitV ops box = x_cat[:, :self.reg_max * 4] cls = x_cat[:, self.reg_max * 4:] else: - box, cls = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).split((self.reg_max * 4, self.nc), 1) + box, cls = x_cat.split((self.reg_max * 4, self.nc), 1) dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides y = torch.cat((dbox, cls.sigmoid()), 1) return y if self.export else (y, x) diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index 65ab0b1..d7af6e9 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -11,8 +11,8 @@ import torch.nn as nn from ultralytics.nn.modules import (C1, C2, C3, C3TR, SPP, SPPF, Bottleneck, BottleneckCSP, C2f, C3Ghost, C3x, Classify, Concat, Conv, ConvTranspose, Detect, DWConv, DWConvTranspose2d, Ensemble, Focus, GhostBottleneck, GhostConv, Segment) -from ultralytics.yolo.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, colorstr, emojis, yaml_load -from ultralytics.yolo.utils.checks import check_requirements, check_yaml +from ultralytics.yolo.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, colorstr, emojis, yaml_load +from ultralytics.yolo.utils.checks import check_requirements, check_suffix, check_yaml from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, fuse_deconv_and_bn, initialize_weights, intersect_dicts, make_divisible, model_info, scale_img, time_sync) @@ -151,15 +151,19 @@ class BaseModel(nn.Module): m.strides = fn(m.strides) return self - def load(self, weights): - """ - This function loads the weights of the model from a file + def load(self, weights, verbose=True): + """Load the weights into the model. Args: - weights (str): The weights to load into the model. + weights (dict) or (torch.nn.Module): The pre-trained weights to be loaded. + verbose (bool, optional): Whether to log the transfer progress. Defaults to True. """ - # Force all tasks to implement this function - raise NotImplementedError('This function needs to be implemented by derived classes!') + model = weights['model'] if isinstance(weights, dict) else weights # torchvision models are not dicts + csd = model.float().state_dict() # checkpoint state_dict as FP32 + csd = intersect_dicts(csd, self.state_dict()) # intersect + self.load_state_dict(csd, strict=False) # load + if verbose: + LOGGER.info(f'Transferred {len(csd)}/{len(self.model.state_dict())} items from pretrained weights') class DetectionModel(BaseModel): @@ -234,13 +238,6 @@ class DetectionModel(BaseModel): y[-1] = y[-1][..., i:] # small return y - def load(self, weights, verbose=True): - csd = weights.float().state_dict() # checkpoint state_dict as FP32 - csd = intersect_dicts(csd, self.state_dict()) # intersect - self.load_state_dict(csd, strict=False) # load - if verbose and RANK == -1: - LOGGER.info(f'Transferred {len(csd)}/{len(self.model.state_dict())} items from pretrained weights') - class SegmentationModel(DetectionModel): # YOLOv8 segmentation model @@ -293,12 +290,6 @@ class ClassificationModel(BaseModel): self.names = {i: f'{i}' for i in range(self.yaml['nc'])} # default names dict self.info() - def load(self, weights): - model = weights['model'] if isinstance(weights, dict) else weights # torchvision models are not dicts - csd = model.float().state_dict() - csd = intersect_dicts(csd, self.state_dict()) # intersect - self.load_state_dict(csd, strict=False) # load - @staticmethod def reshape_outputs(model, nc): # Update a TorchVision classification model to class count 'n' if required @@ -338,6 +329,7 @@ def torch_safe_load(weight): """ from ultralytics.yolo.utils.downloads import attempt_download_asset + check_suffix(file=weight, suffix='.pt') file = attempt_download_asset(weight) # search online if missing locally try: return torch.load(file, map_location='cpu'), file # load diff --git a/ultralytics/yolo/cfg/__init__.py b/ultralytics/yolo/cfg/__init__.py index c1d3af8..e61a30b 100644 --- a/ultralytics/yolo/cfg/__init__.py +++ b/ultralytics/yolo/cfg/__init__.py @@ -54,11 +54,10 @@ CFG_FRACTION_KEYS = ('dropout', 'iou', 'lr0', 'lrf', 'momentum', 'weight_decay', 'fliplr', 'mosaic', 'mixup', 'copy_paste', 'conf', 'iou') # fractional floats limited to 0.0 - 1.0 CFG_INT_KEYS = ('epochs', 'patience', 'batch', 'workers', 'seed', 'close_mosaic', 'mask_ratio', 'max_det', 'vid_stride', 'line_thickness', 'workspace', 'nbs', 'save_period') -CFG_BOOL_KEYS = ('save', 'exist_ok', 'pretrained', 'verbose', 'deterministic', 'single_cls', 'image_weights', 'rect', - 'cos_lr', 'overlap_mask', 'val', 'save_json', 'save_hybrid', 'half', 'dnn', 'plots', 'show', - 'save_txt', 'save_conf', 'save_crop', 'hide_labels', 'hide_conf', 'visualize', 'augment', - 'agnostic_nms', 'retina_masks', 'boxes', 'keras', 'optimize', 'int8', 'dynamic', 'simplify', 'nms', - 'v5loader') +CFG_BOOL_KEYS = ('save', 'exist_ok', 'verbose', 'deterministic', 'single_cls', 'image_weights', 'rect', 'cos_lr', + 'overlap_mask', 'val', 'save_json', 'save_hybrid', 'half', 'dnn', 'plots', 'show', 'save_txt', + 'save_conf', 'save_crop', 'hide_labels', 'hide_conf', 'visualize', 'augment', 'agnostic_nms', + 'retina_masks', 'boxes', 'keras', 'optimize', 'int8', 'dynamic', 'simplify', 'nms', 'v5loader') # Define valid tasks and modes MODES = 'train', 'val', 'predict', 'export', 'track', 'benchmark' @@ -290,6 +289,8 @@ def entrypoint(debug=''): from ultralytics.yolo.engine.model import YOLO overrides['model'] = model model = YOLO(model, task=task) + if isinstance(overrides.get('pretrained'), str): + model.load(overrides['pretrained']) # Task Update if task != model.task: diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py index a920c3d..1943c84 100644 --- a/ultralytics/yolo/engine/exporter.py +++ b/ultralytics/yolo/engine/exporter.py @@ -188,7 +188,7 @@ class Exporter: m.dynamic = self.args.dynamic m.export = True m.format = self.args.format - elif isinstance(m, C2f) and not edgetpu: + elif isinstance(m, C2f) and not any((saved_model, pb, tflite, edgetpu, tfjs)): # EdgeTPU does not support FlexSplitV while split provides cleaner ONNX graph m.forward = m.forward_split diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index f9f73b5..df46056 100644 --- a/ultralytics/yolo/engine/model.py +++ b/ultralytics/yolo/engine/model.py @@ -8,8 +8,8 @@ from ultralytics.nn.tasks import (ClassificationModel, DetectionModel, Segmentat guess_model_task, nn) from ultralytics.yolo.cfg import get_cfg from ultralytics.yolo.engine.exporter import Exporter -from ultralytics.yolo.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, ONLINE, RANK, ROOT, - callbacks, is_git_dir, is_pip_package, yaml_load) +from ultralytics.yolo.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, ROOT, callbacks, + is_git_dir, yaml_load) from ultralytics.yolo.utils.checks import check_file, check_imgsz, check_pip_update_available, check_yaml from ultralytics.yolo.utils.downloads import GITHUB_ASSET_STEMS from ultralytics.yolo.utils.torch_utils import smart_inference_mode @@ -153,16 +153,10 @@ class YOLO: f"'yolo export model=yolov8n.pt', but exported formats like ONNX, TensorRT etc. only " f"support 'predict' and 'val' modes, i.e. 'yolo predict model=yolov8n.onnx'.") - def _check_pip_update(self): - """ - Inform user of ultralytics package update availability - """ - if ONLINE and is_pip_package(): - check_pip_update_available() - - def reset(self): + @smart_inference_mode() + def reset_weights(self): """ - Resets the model modules. + Resets the model modules parameters to randomly initialized values, losing all training information. """ self._check_is_pytorch_model() for m in self.model.modules(): @@ -170,6 +164,18 @@ class YOLO: m.reset_parameters() for p in self.model.parameters(): p.requires_grad = True + return self + + @smart_inference_mode() + def load(self, weights='yolov8n.pt'): + """ + Transfers parameters with matching names and shapes from 'weights' to model. + """ + self._check_is_pytorch_model() + if isinstance(weights, (str, Path)): + weights, self.ckpt = attempt_load_one_weight(weights) + self.model.load(weights) + return self def info(self, verbose=False): """ @@ -299,7 +305,7 @@ class YOLO: **kwargs (Any): Any number of arguments representing the training configuration. """ self._check_is_pytorch_model() - self._check_pip_update() + check_pip_update_available() overrides = self.overrides.copy() overrides.update(kwargs) if kwargs.get('cfg'): diff --git a/ultralytics/yolo/engine/results.py b/ultralytics/yolo/engine/results.py index 8f80a36..02a6839 100644 --- a/ultralytics/yolo/engine/results.py +++ b/ultralytics/yolo/engine/results.py @@ -48,7 +48,7 @@ class Results: self.probs = probs if probs is not None else None self.names = names self.path = path - self._keys = [k for k in ('boxes', 'masks', 'probs') if getattr(self, k) is not None] + self._keys = ('boxes', 'masks', 'probs') def pandas(self): pass @@ -56,7 +56,7 @@ class Results: def __getitem__(self, idx): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for k in self._keys: + for k in self.keys: setattr(r, k, getattr(self, k)[idx]) return r @@ -70,30 +70,30 @@ class Results: def cpu(self): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for k in self._keys: + for k in self.keys: setattr(r, k, getattr(self, k).cpu()) return r def numpy(self): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for k in self._keys: + for k in self.keys: setattr(r, k, getattr(self, k).numpy()) return r def cuda(self): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for k in self._keys: + for k in self.keys: setattr(r, k, getattr(self, k).cuda()) return r def to(self, *args, **kwargs): r = Results(orig_img=self.orig_img, path=self.path, names=self.names) - for k in self._keys: + for k in self.keys: setattr(r, k, getattr(self, k).to(*args, **kwargs)) return r def __len__(self): - for k in self._keys: + for k in self.keys: return len(getattr(self, k)) def __str__(self): @@ -107,6 +107,10 @@ class Results: name = self.__class__.__name__ raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.__doc__}") + @property + def keys(self): + return [k for k in self._keys if getattr(self, k) is not None] + def plot(self, show_conf=True, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'): """ Plots the detection results on an input RGB image. Accepts a numpy array (cv2) or a PIL Image. diff --git a/ultralytics/yolo/utils/__init__.py b/ultralytics/yolo/utils/__init__.py index 4bcfdde..a821aa6 100644 --- a/ultralytics/yolo/utils/__init__.py +++ b/ultralytics/yolo/utils/__init__.py @@ -46,14 +46,14 @@ HELP_MSG = \ from ultralytics import YOLO # Load a model - model = YOLO("yolov8n.yaml") # build a new model from scratch + model = YOLO('yolov8n.yaml') # build a new model from scratch model = YOLO("yolov8n.pt") # load a pretrained model (recommended for training) # Use the model results = model.train(data="coco128.yaml", epochs=3) # train the model results = model.val() # evaluate model performance on the validation set - results = model("https://ultralytics.com/images/bus.jpg") # predict on an image - success = model.export(format="onnx") # export the model to ONNX format + results = model('https://ultralytics.com/images/bus.jpg') # predict on an image + success = model.export(format='onnx') # export the model to ONNX format 3. Use the command line interface (CLI): diff --git a/ultralytics/yolo/utils/autobatch.py b/ultralytics/yolo/utils/autobatch.py index d528864..d730a98 100644 --- a/ultralytics/yolo/utils/autobatch.py +++ b/ultralytics/yolo/utils/autobatch.py @@ -1,6 +1,6 @@ # Ultralytics YOLO 🚀, GPL-3.0 license """ -AutoBatch utils +Functions for estimating the best YOLO batch size to use a fraction of the available CUDA memory in PyTorch. """ from copy import deepcopy @@ -13,18 +13,35 @@ from ultralytics.yolo.utils.torch_utils import profile def check_train_batch_size(model, imgsz=640, amp=True): - # Check YOLOv5 training batch size + """ + Check YOLO training batch size using the autobatch() function. + + Args: + model (torch.nn.Module): YOLO model to check batch size for. + imgsz (int): Image size used for training. + amp (bool): If True, use automatic mixed precision (AMP) for training. + + Returns: + int: Optimal batch size computed using the autobatch() function. + """ + with torch.cuda.amp.autocast(amp): return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size -def autobatch(model, imgsz=640, fraction=0.7, batch_size=16): - # Automatically estimate best YOLOv5 batch size to use `fraction` of available CUDA memory - # Usage: - # import torch - # from utils.autobatch import autobatch - # model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False) - # print(autobatch(model)) +def autobatch(model, imgsz=640, fraction=0.67, batch_size=16): + """ + Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory. + + Args: + model: YOLO model to compute batch size for. + imgsz (int, optional): The image size used as input for the YOLO model. Defaults to 640. + fraction (float, optional): The fraction of available CUDA memory to use. Defaults to 0.67. + batch_size (int, optional): The default batch size to use if an error is detected. Defaults to 16. + + Returns: + int: The optimal batch size. + """ # Check device prefix = colorstr('AutoBatch: ') diff --git a/ultralytics/yolo/utils/callbacks/tensorboard.py b/ultralytics/yolo/utils/callbacks/tensorboard.py index 76835de..07d8347 100644 --- a/ultralytics/yolo/utils/callbacks/tensorboard.py +++ b/ultralytics/yolo/utils/callbacks/tensorboard.py @@ -1,5 +1,5 @@ # Ultralytics YOLO 🚀, GPL-3.0 license -from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING +from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING, colorstr try: from torch.utils.tensorboard import SummaryWriter @@ -18,11 +18,14 @@ def _log_scalars(scalars, step=0): def on_pretrain_routine_start(trainer): - global writer - try: - writer = SummaryWriter(str(trainer.save_dir)) - except Exception as e: - LOGGER.warning(f'WARNING ⚠️ TensorBoard not initialized correctly, not logging this run. {e}') + if SummaryWriter: + try: + global writer + writer = SummaryWriter(str(trainer.save_dir)) + prefix = colorstr('TensorBoard: ') + LOGGER.info(f"{prefix}Start with 'tensorboard --logdir {trainer.save_dir}', view at http://localhost:6006/") + except Exception as e: + LOGGER.warning(f'WARNING ⚠️ TensorBoard not initialized correctly, not logging this run. {e}') def on_batch_end(trainer): diff --git a/ultralytics/yolo/utils/checks.py b/ultralytics/yolo/utils/checks.py index 891bac8..06aaf72 100644 --- a/ultralytics/yolo/utils/checks.py +++ b/ultralytics/yolo/utils/checks.py @@ -20,8 +20,8 @@ import requests import torch from matplotlib import font_manager -from ultralytics.yolo.utils import (AUTOINSTALL, LOGGER, ROOT, USER_CONFIG_DIR, TryExcept, colorstr, downloads, emojis, - is_colab, is_docker, is_jupyter, is_online) +from ultralytics.yolo.utils import (AUTOINSTALL, LOGGER, ONLINE, ROOT, USER_CONFIG_DIR, TryExcept, colorstr, downloads, + emojis, is_colab, is_docker, is_jupyter, is_online, is_pip_package) def is_ascii(s) -> bool: @@ -141,12 +141,14 @@ def check_pip_update_available(): Returns: bool: True if an update is available, False otherwise. """ - from ultralytics import __version__ - latest = check_latest_pypi_version() - if pkg.parse_version(__version__) < pkg.parse_version(latest): # update is available - LOGGER.info(f'New https://pypi.org/project/ultralytics/{latest} available 😃 ' - f"Update with 'pip install -U ultralytics'") - return True + if ONLINE and is_pip_package(): + with contextlib.suppress(ConnectionError): + from ultralytics import __version__ + latest = check_latest_pypi_version() + if pkg.parse_version(__version__) < pkg.parse_version(latest): # update is available + LOGGER.info(f'New https://pypi.org/project/ultralytics/{latest} available 😃 ' + f"Update with 'pip install -U ultralytics'") + return True return False @@ -235,11 +237,11 @@ def check_suffix(file='yolov8n.pt', suffix='.pt', msg=''): # Check file(s) for acceptable suffix if file and suffix: if isinstance(suffix, str): - suffix = [suffix] + suffix = (suffix, ) for f in file if isinstance(file, (list, tuple)) else [file]: s = Path(f).suffix.lower() # file suffix if len(s): - assert s in suffix, f'{msg}{f} acceptable suffix is {suffix}' + assert s in suffix, f'{msg}{f} acceptable suffix is {suffix}, not {s}' def check_yolov5u_filename(file: str, verbose: bool = True): diff --git a/ultralytics/yolo/v8/detect/predict.py b/ultralytics/yolo/v8/detect/predict.py index 98210f6..7c38674 100644 --- a/ultralytics/yolo/v8/detect/predict.py +++ b/ultralytics/yolo/v8/detect/predict.py @@ -76,7 +76,7 @@ class DetectionPredictor(BasePredictor): if self.args.save_crop: save_one_box(d.xyxy, imc, - file=self.save_dir / 'crops' / self.model.model.names[c] / f'{self.data_path.stem}.jpg', + file=self.save_dir / 'crops' / self.model.names[c] / f'{self.data_path.stem}.jpg', BGR=True) return log_string diff --git a/ultralytics/yolo/v8/detect/train.py b/ultralytics/yolo/v8/detect/train.py index a3e8f21..d2584db 100644 --- a/ultralytics/yolo/v8/detect/train.py +++ b/ultralytics/yolo/v8/detect/train.py @@ -58,10 +58,9 @@ class DetectionTrainer(BaseTrainer): # TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc def get_model(self, cfg=None, weights=None, verbose=True): - model = DetectionModel(cfg, ch=3, nc=self.data['nc'], verbose=verbose and RANK == -1) + model = DetectionModel(cfg, nc=self.data['nc'], verbose=verbose and RANK == -1) if weights: model.load(weights) - return model def get_validator(self): diff --git a/ultralytics/yolo/v8/segment/predict.py b/ultralytics/yolo/v8/segment/predict.py index ef5a8d8..82fc47f 100644 --- a/ultralytics/yolo/v8/segment/predict.py +++ b/ultralytics/yolo/v8/segment/predict.py @@ -90,7 +90,7 @@ class SegmentationPredictor(DetectionPredictor): if self.args.save_crop: save_one_box(d.xyxy, imc, - file=self.save_dir / 'crops' / self.model.model.names[c] / f'{self.data_path.stem}.jpg', + file=self.save_dir / 'crops' / self.model.names[c] / f'{self.data_path.stem}.jpg', BGR=True) return log_string