Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Ultralytics Code Refactor https://ultralytics.com/actions (#13191)

Browse Source 
* Refactor code for speed and clarity

* Auto-format by https://ultralytics.com/actions

---------

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Glenn JocherandUltralyticsAssistant authored and GitHub committed 2024-07-16 01:20:57 +02:00
1 parent 41e26a7f83
commit 12be49963c
6 files changed
+216 -213
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

export.py
+84 -102
View File
@@ -97,15 +97,15 @@ class iOSModel(torch.nn.Module):
Args:
model (torch.nn.Module): The PyTorch model to be adapted for iOS compatibility.
im (torch.Tensor): An input tensor representing a batch of images with shape (batch, channel, height, width).
im (torch.Tensor): An input tensor representing a batch of images with shape (B, C, H, W).
Returns:
None: This method does not return any value.
Notes:
This initializer configures normalization based on the input image dimensions, which is critical for ensuring the
model's compatibility and proper functionality on iOS devices. The normalization step involves dividing by the image
width if the image is square; otherwise, additional conditions might apply (trimmed for brevity).
This initializer configures normalization based on the input image dimensions, which is critical for
ensuring the model's compatibility and proper functionality on iOS devices. The normalization step
involves dividing by the image width if the image is square; otherwise, additional conditions might apply.
"""
super().__init__()
b, c, h, w = im.shape # batch, channel, height, width
@@ -120,14 +120,15 @@ class iOSModel(torch.nn.Module):
def forward(self, x):
"""
Runs a forward pass on the input tensor, returning class confidences and normalized coordinates.
Run a forward pass on the input tensor, returning class confidences and normalized coordinates.
Args:
x (torch.Tensor): Input tensor containing the image data.
x (torch.Tensor): Input tensor containing the image data with shape (batch, channels, height, width).
Returns:
torch.Tensor: Concatenated tensor with normalized coordinates (xywh), confidence scores (conf),
and class probabilities (cls).
and class probabilities (cls), having shape (N, 4 + 1 + C), where N is the number of predictions,
and C is the number of classes.
Examples:
```python
@@ -144,9 +145,9 @@ def export_formats():
Returns a DataFrame of supported YOLOv5 model export formats and their properties.
Returns:
pd.DataFrame: A pandas DataFrame containing supported export formats and their properties. The DataFrame includes
columns for format name, CLI argument suffix, file extension or directory name, and boolean flags indicating if the
export format supports training and detection.
pandas.DataFrame: A DataFrame containing supported export formats and their properties. The DataFrame
includes columns for format name, CLI argument suffix, file extension or directory name, and boolean flags
indicating if the export format supports training and detection.
Examples:
```python
@@ -156,7 +157,7 @@ def export_formats():
Notes:
The DataFrame contains the following columns:
- Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.)
- Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.).
- Include Argument: The argument to use with the export script to include this format.
- File Suffix: File extension or directory name associated with the format.
- Supports Training: Whether the format supports training.
@@ -181,7 +182,7 @@ def export_formats():
def try_export(inner_func):
"""
Logs success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export.
Log success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export.
Args:
inner_func (Callable): The model export function to be wrapped by the decorator.
@@ -202,7 +203,8 @@ def try_export(inner_func):
```
Notes:
For additional requirements and model export formats, refer to the [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics).
For additional requirements and model export formats, refer to the
[Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics).
"""
inner_args = get_default_args(inner_func)
@@ -224,7 +226,7 @@ def try_export(inner_func):
@try_export
def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
"""
Exports a YOLOv5 model to the TorchScript format.
Export a YOLOv5 model to the TorchScript format.
Args:
model (torch.nn.Module): The YOLOv5 model to be exported.
@@ -235,8 +237,8 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:"
Returns:
(str | None, torch.jit.ScriptModule | None): A tuple containing the file path of the exported model
(as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements
of the tuple will be None.
(as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements
of the tuple will be None.
Notes:
- This function uses tracing to create the TorchScript model.
@@ -303,6 +305,26 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX
```
pip install onnx onnx-simplifier onnxruntime onnxruntime-gpu
```
Example:
```python
from pathlib import Path
import torch
from models.experimental import attempt_load
from utils.torch_utils import select_device
# Load model
weights = 'yolov5s.pt'
device = select_device('')
model = attempt_load(weights, map_location=device)
# Example input tensor
im = torch.zeros(1, 3, 640, 640).to(device)
# Export model
file_path = Path('yolov5s.onnx')
export_onnx(model, im, file_path, opset=12, dynamic=True, simplify=True)
```
"""
check_requirements("onnx>=1.12.0")
import onnx
@@ -361,10 +383,10 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX
@try_export
def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")):
"""
Exports a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization.
Export a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization.
Args:
file (Path): The path to the output file where the OpenVINO model will be saved.
file (Path): Path to the output file where the OpenVINO model will be saved.
metadata (dict): Dictionary including model metadata such as names and strides.
half (bool): If True, export the model with FP16 precision.
int8 (bool): If True, export the model with INT8 quantization.
@@ -372,8 +394,8 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:
prefix (str): Prefix string for logging purposes (default is "OpenVINO:").
Returns:
(str, openvino.runtime.Model | None): Returns the OpenVINO model file path and openvino.runtime.Model object if
export is successful; otherwise, returns None.
(str, openvino.runtime.Model | None): The OpenVINO model file path and openvino.runtime.Model object if export is
successful; otherwise, None.
Notes:
- Requires `openvino-dev` package version 2023.0 or higher. Install with:
@@ -451,7 +473,7 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:
@try_export
def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
"""
Export a YOLOv5 model to PaddlePaddle format using X2Paddle, saving the converted model and metadata.
Export a YOLOv5 PyTorch model to PaddlePaddle format using X2Paddle, saving the converted model and metadata.
Args:
model (torch.nn.Module): The YOLOv5 model to be exported.
@@ -554,12 +576,12 @@ def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")):
@try_export
def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr("TensorRT:")):
"""
Exports a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
Export a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
Args:
model (torch.nn.Module): YOLOv5 model to be exported.
im (torch.Tensor): Input tensor of shape (B, C, H, W).
file (Path): Path to save the exported model.
file (pathlib.Path): Path to save the exported model.
half (bool): Set to True to export with FP16 precision.
dynamic (bool): Set to True to enable dynamic input shapes.
simplify (bool): Set to True to simplify the model during export.
@@ -568,7 +590,7 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
prefix (str): Log message prefix.
Returns:
(Path, None): Tuple containing the path to the exported model and None.
(pathlib.Path, None): Tuple containing the path to the exported model and None.
Raises:
AssertionError: If executed on CPU instead of GPU.
@@ -666,7 +688,7 @@ def export_saved_model(
prefix=colorstr("TensorFlow SavedModel:"),
):
"""
Exports a YOLOv5 model to the TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression
Export a YOLOv5 model to the TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression
(NMS).
Args:
@@ -791,7 +813,7 @@ def export_tflite(
):
# YOLOv5 TensorFlow Lite export
"""
Exports a YOLOv5 model to TensorFlow Lite format with optional INT8 quantization and NMS support.
Export a YOLOv5 model to TensorFlow Lite format with optional INT8 quantization and NMS support.
Args:
keras_model (tf.keras.Model): The Keras model to be exported.
@@ -805,7 +827,8 @@ def export_tflite(
prefix (str): Prefix for log messages.
Returns:
(str | None, tflite.Model | None): The file path of the exported TFLite model and the TFLite model instance, or None if export failed.
(str | None, tflite.Model | None): The file path of the exported TFLite model and the TFLite model instance, or None
if the export failed.
Example:
```python
@@ -863,8 +886,6 @@ def export_tflite(
@try_export
def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
"""
Clear and concise summary line describing the function's purpose:
Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler.
Args:
@@ -927,28 +948,30 @@ def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
@try_export
def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
"""
Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization.
Convert a YOLOv5 model to TensorFlow.js format with optional uint8 quantization.
Args:
file (Path): Path to the input model file.
int8 (bool): If True, applies uint8 quantization.
prefix (str): Prefix for logging information (default: colorstr("TensorFlow.js:")).
file (Path): Path to the YOLOv5 model file to be converted, typically having a ".pt" or ".onnx" extension.
int8 (bool): If True, applies uint8 quantization during the conversion process.
prefix (str): Optional prefix for logging messages, default is 'TensorFlow.js:' with color formatting.
Returns:
(str, None): The output directory path as a string and None.
(str, None): Tuple containing the output directory path as a string and None.
Notes:
This function requires `tensorflowjs` to be installed. You can install it using:
```shell
pip install tensorflowjs
```
- This function requires the `tensorflowjs` package. Install it using:
```shell
pip install tensorflowjs
```
- The converted TensorFlow.js model will be saved in a directory with the "_web_model" suffix appended to the original file name.
- The conversion involves running shell commands that invoke the TensorFlow.js converter tool.
Example usage:
Example:
```python
export_tfjs(Path('yolov5s.onnx'), int8=False)
from pathlib import Path
file = Path('yolov5.onnx')
export_tfjs(file, int8=False)
```
The TensorFlow.js converted model is saved in the directory specified by `file` with "_web_model" suffix.
"""
check_requirements("tensorflowjs")
import tensorflowjs as tfjs
@@ -987,11 +1010,13 @@ def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
def add_tflite_metadata(file, metadata, num_outputs):
"""
Adds metadata to a TensorFlow Lite (TFLite) model file, supporting multiple outputs, based on TensorFlow guidelines.
Adds metadata to a TensorFlow Lite (TFLite) model file, supporting multiple outputs according to TensorFlow
guidelines.
Args:
file (str): Path to the TFLite model file to which metadata will be added.
metadata (dict): Metadata information to be added to the model, structured as required by the TFLite metadata schema.
Common keys include "name", "description", "version", "author", and "license".
num_outputs (int): Number of output tensors the model has, used to configure the metadata properly.
Returns:
@@ -1047,20 +1072,20 @@ def add_tflite_metadata(file, metadata, num_outputs):
def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")):
"""
Converts a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output
shapes and saving the model.
Convert a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output
shapes, and saving the model.
Args:
model (torch.nn.Module): The YOLOv5 PyTorch model.
im (torch.Tensor): Input tensor example with shape [N, C, H, W], where N is the batch size, C is the number of
channels, H is the height, and W is the width.
model (torch.nn.Module): The YOLOv5 PyTorch model to be converted.
im (torch.Tensor): Example input tensor with shape (N, C, H, W), where N is the batch size, C is the number of channels,
H is the height, and W is the width.
file (Path): Path to save the converted CoreML model.
names (dict[int, str]): Dictionary mapping class indices to class names.
y (torch.Tensor): Output tensor from the PyTorch model's forward pass.
prefix (str): Custom prefix for logging messages.
Returns:
Path: Path to the saved CoreML model (.mlmodel).
(Path): Path to the saved CoreML model (.mlmodel).
Raises:
AssertionError: If the number of class names does not match the number of classes in the model.
@@ -1075,18 +1100,14 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:
from pathlib import Path
import torch
# Load YOLOv5 model and an example input tensor
model = torch.load("yolov5s.pt")
im = torch.zeros(1, 3, 640, 640) # Example input tensor
model = torch.load('yolov5s.pt') # Load YOLOv5 model
im = torch.zeros((1, 3, 640, 640)) # Example input tensor
# Define class names
names = {0: "person", 1: "bicycle", 2: "car", ...}
names = {0: "person", 1: "bicycle", 2: "car", ...} # Define class names
# Perform forward pass to get model output
y = model(im)
y = model(im) # Perform forward pass to get model output
# Convert to CoreML
output_file = Path("yolov5s.mlmodel")
output_file = Path('yolov5s.mlmodel') # Convert to CoreML
pipeline_coreml(model, im, output_file, names, y)
```
"""
@@ -1246,8 +1267,6 @@ def run(
conf_thres=0.25, # TF.js NMS: confidence threshold
):
"""
Clear and concise summary line describing the function's purpose:
Exports a YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow.
Args:
@@ -1424,14 +1443,14 @@ def run(
def parse_opt(known=False):
"""
Parses command-line arguments for YOLOv5 model export configurations.
Parse command-line options for YOLOv5 model export configurations.
Args:
known (bool): If True, `argparse.ArgumentParser.parse_known_args` is used to parse command-line arguments; otherwise,
`argparse.ArgumentParser.parse_args` is used. Defaults to False.
known (bool): If True, uses `argparse.ArgumentParser.parse_known_args`; otherwise, uses `argparse.ArgumentParser.parse_args`.
Default is False.
Returns:
argparse.Namespace: An object containing parsed command-line arguments.
argparse.Namespace: Object containing parsed command-line arguments.
Example:
```python
@@ -1475,44 +1494,7 @@ def parse_opt(known=False):
def main(opt):
"""
Exports the YOLOv5 model to specified formats, including ONNX, TensorRT, CoreML, and TensorFlow.
Args:
opt (argparse.Namespace): Parsed command-line arguments containing the export configurations.
- data (str): Path to the dataset YAML configuration file (e.g., 'data/coco128.yaml').
- weights (list[str] | str): Paths to the pretrained model weights file(s) (e.g., 'yolov5s.pt').
- imgsz (list[int]): Image size as a list [height, width].
- batch_size (int): Batch size for exporting the model.
- device (str): Device to run the export on, such as '0' for GPU, or 'cpu' for CPU.
- half (bool): Flag to export the model with FP16 half-precision.
- inplace (bool): Set the YOLOv5 Detect() module inplace mode to True.
- keras (bool): Flag to use Keras for TensorFlow SavedModel export.
- optimize (bool): Optimize TorchScript model for mobile deployment.
- int8 (bool): Apply INT8 quantization for CoreML or TensorFlow models.
- per_tensor (bool): Apply per-tensor quantization for TensorFlow models.
- dynamic (bool): Enable dynamic axes for ONNX, TensorFlow, or TensorRT exports.
- simplify (bool): Simplify ONNX model during export.
- opset (int): ONNX opset version.
- verbose (bool): Enable verbose logging for TensorRT export.
- workspace (int): TensorRT workspace size in GB.
- nms (bool): Add non-maximum suppression (NMS) to the TensorFlow model.
- agnostic_nms (bool): Add class-agnostic NMS to the TensorFlow model.
- topk_per_class (int): Top-K boxes per class to keep for TensorFlow.js NMS.
- topk_all (int): Top-K boxes for all classes to keep for TensorFlow.js NMS.
- iou_thres (float): IoU threshold for NMS.
- conf_thres (float): Confidence threshold for NMS.
- include (list[str]): List of formats to include in export (e.g., ['torchscript', 'onnx']).
Returns:
list[str]: List of paths to the exported model files.
Example:
```python
opt = parse_opt()
main(opt)
```
"""
"""Run(**vars(opt)) # Execute the run function with parsed options."""
for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
run(**vars(opt))