* 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

No files matched your search

+59 -43
View File
@@ -24,10 +24,11 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
classes (int, optional): Number of classes the model is expected to detect. Defaults to 80.
autoshape (bool, optional): If True, applies the YOLOv5 .autoshape() wrapper for various input formats. Defaults to True.
verbose (bool, optional): If True, prints detailed information during the model creation/loading process. Defaults to True.
device (str | torch.device | None, optional): Device to use for model parameters (e.g., 'cpu', 'cuda'). If None, selects the best available device. Defaults to None.
device (str | torch.device | None, optional): Device to use for model parameters (e.g., 'cpu', 'cuda'). If None, selects
the best available device. Defaults to None.
Returns:
DetectMultiBackend | AutoShape: The loaded YOLOv5 model, potentially wrapped with AutoShape if specified.
(DetectMultiBackend | AutoShape): The loaded YOLOv5 model, potentially wrapped with AutoShape if specified.
Examples:
```python
@@ -108,8 +109,8 @@ def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None):
Args:
path (str): Path to the custom model file (e.g., 'path/to/model.pt').
autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model if True, enabling compatibility with various input types
(default is True).
autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model if True, enabling compatibility with various input
types (default is True).
_verbose (bool): If True, prints all informational messages to the screen; otherwise, operates silently
(default is True).
device (str | torch.device | None): Device to load the model on, e.g., 'cpu', 'cuda', torch.device('cuda:0'), etc.
@@ -174,32 +175,42 @@ def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
"""
Creates YOLOv5-small model with options for pretraining, input channels, class count, autoshaping, verbosity, and
device.
Create a YOLOv5-small (yolov5s) model with options for pretraining, input channels, class count, autoshaping,
verbosity, and device configuration.
Args:
pretrained (bool, optional): Flag to load pretrained weights into the model. Defaults to True.
channels (int, optional): Number of input channels. Defaults to 3.
classes (int, optional): Number of model classes. Defaults to 80.
autoshape (bool, optional): Whether to apply YOLOv5 .autoshape() wrapper to the model for preprocessed inputs.
autoshape (bool, optional): Whether to wrap the model with YOLOv5's .autoshape() for handling various input formats.
Defaults to True.
_verbose (bool, optional): Flag to print detailed information on model loading. Defaults to True.
device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda'. If None,
auto-select the best available device. Defaults to None.
_verbose (bool, optional): Flag to print detailed information regarding model loading. Defaults to True.
device (str | torch.device | None, optional): Device to use for model computation, can be 'cpu', 'cuda', or
torch.device instances. If None, automatically selects the best available device. Defaults to None.
Returns:
torch.nn.Module: The YOLOv5-small model loaded with specified configurations and optionally pretrained weights.
torch.nn.Module: The YOLOv5-small model configured and loaded according to the specified parameters.
Usage:
Example:
```python
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # official model
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s') # from branch
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt') # custom/local model
model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local') # local repo
# Load the official YOLOv5-small model with pretrained weights
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
# Load the YOLOv5-small model from a specific branch
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s')
# Load a custom YOLOv5-small model from a local checkpoint
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt')
# Load a local YOLOv5-small model specifying source as local repository
model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local')
```
For more information, visit https://pytorch.org/hub/ultralytics_yolov5.
Notes:
For more details on model loading and customization, visit
the [YOLOv5 PyTorch Hub Documentation](https://pytorch.org/hub/ultralytics_yolov5).
"""
return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device)
@@ -213,22 +224,26 @@ def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
pretrained (bool, optional): Whether to load pretrained weights into the model. Default is True.
channels (int, optional): Number of input channels. Default is 3.
classes (int, optional): Number of model classes. Default is 80.
autoshape (bool, optional): Apply YOLOv5 .autoshape() wrapper to the model for handling various input formats. Default is True.
autoshape (bool, optional): Apply YOLOv5 .autoshape() wrapper to the model for handling various input formats.
Default is True.
_verbose (bool, optional): Whether to print detailed information to the screen. Default is True.
device (str | torch.device | None, optional): Device specification to use for model parameters (e.g., 'cpu', 'cuda'). Default is None.
device (str | torch.device | None, optional): Device specification to use for model parameters (e.g., 'cpu', 'cuda').
Default is None.
Returns:
torch.nn.Module: The instantiated YOLOv5-medium model.
Usage Example:
```python
import torch
```python
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5m') # Load YOLOv5-medium from Ultralytics repository
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5m') # Load from the master branch
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5m.pt') # Load a custom/local YOLOv5-medium model
model = torch.hub.load('.', 'custom', 'yolov5m.pt', source='local') # Load from a local repository
```
model = torch.hub.load('ultralytics/yolov5', 'yolov5m') # Load YOLOv5-medium from Ultralytics repository
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5m') # Load from the master branch
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5m.pt') # Load a custom/local YOLOv5-medium model
model = torch.hub.load('.', 'custom', 'yolov5m.pt', source='local') # Load from a local repository
```
For more information, visit https://pytorch.org/hub/ultralytics_yolov5.
"""
return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device)
@@ -251,7 +266,7 @@ def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
YOLOv5 model (torch.nn.Module): The YOLOv5-large model instantiated with specified configurations and possibly
pretrained weights.
Example:
Examples:
```python
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5l')
@@ -266,21 +281,22 @@ def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
"""
Instantiates the YOLOv5-xlarge model with customizable pretraining, channel count, class count, autoshaping,
verbosity, and device.
Perform object detection using the YOLOv5-xlarge model with options for pretraining, input channels, class count,
autoshaping, verbosity, and device specification.
Args:
pretrained (bool): If True, loads pretrained weights into the model. Defaults to True.
channels (int): Number of input channels. Defaults to 3.
classes (int): Number of model classes. Defaults to 80.
autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model for easier image handling. Defaults to
channels (int): Number of input channels for the model. Defaults to 3.
classes (int): Number of model classes for object detection. Defaults to 80.
autoshape (bool): If True, applies the YOLOv5 .autoshape() wrapper for handling different input formats. Defaults to
True.
_verbose (bool): If True, prints detailed information to the screen. Defaults to True.
device (str | torch.device | None): Device for model parameters, e.g., 'cpu', 'cuda:0', or a torch.device object.
_verbose (bool): If True, prints detailed information during model loading. Defaults to True.
device (str | torch.device | None): Device specification for computing the model, e.g., 'cpu', 'cuda:0', torch.device('cuda').
Defaults to None.
Returns:
torch.nn.Module: The instantiated YOLOv5-xlarge model.
torch.nn.Module: The YOLOv5-xlarge model loaded with the specified parameters, optionally with pretrained weights and
autoshaping applied.
Example:
```python
@@ -288,7 +304,7 @@ def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
model = torch.hub.load('ultralytics/yolov5', 'yolov5x')
```
For more details and usage, refer to the official YOLOv5 PyTorch Hub models documentation:
For additional details, refer to the official YOLOv5 PyTorch Hub models documentation:
https://pytorch.org/hub/ultralytics_yolov5
"""
return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device)
@@ -308,7 +324,7 @@ def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
Default is None.
Returns:
torch.nn.Module: YOLOv5 model loaded with the specified configurations.
torch.nn.Module: YOLOv5-nano-P6 model loaded with the specified configurations.
Example:
```python
@@ -362,7 +378,7 @@ def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
"""
Creates YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity, and
Create YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity, and
device.
Args:
@@ -389,15 +405,15 @@ def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
```
Notes:
- The model can be loaded with pre-trained weights for better performance on specific tasks.
- The autoshape feature simplifies input handling by allowing various popular data formats.
- The model can be loaded with pre-trained weights for better performance on specific tasks.
- The autoshape feature simplifies input handling by allowing various popular data formats.
"""
return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
"""
Instantiates the YOLOv5-large-P6 model with customizable pretraining, channel and class counts, autoshaping,
Instantiate the YOLOv5-large-P6 model with options for pretraining, channel and class counts, autoshaping,
verbosity, and device selection.
Args:
@@ -407,12 +423,12 @@ def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
autoshape (bool, optional): If True, apply YOLOv5 .autoshape() wrapper to the model for input flexibility. Default is True.
_verbose (bool, optional): If True, print all information to the screen. Default is True.
device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda', or torch.device.
If None, automatically selects the best available device. Default is None.
If None, automatically selects the best available device. Default is None.
Returns:
torch.nn.Module: The instantiated YOLOv5-large-P6 model.
Usage:
Example:
```python
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5l6') # official model