xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Docs]: Link buttons, add autobackend, BaseModel and ops (#130)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ayush Chaurasia
2023-01-02 20:42:30 +05:30
committed by GitHub
parent af6e3c536b
commit 8996c5c6cf
10 changed files with 562 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
ultralytics/nn/autobackend.py
View File

@ -17,22 +17,36 @@ from ultralytics.yolo.utils.ops import xywh2xyxy
class AutoBackend(nn.Module):
# YOLOv5 MultiBackend class for python inference on various backends
def __init__(self, weights='yolov8n.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):
# Usage:
# PyTorch: weights = *.pt
# TorchScript: *.torchscript
# ONNX Runtime: *.onnx
# ONNX OpenCV DNN: *.onnx --dnn
# OpenVINO: *.xml
# CoreML: *.mlmodel
# TensorRT: *.engine
# TensorFlow SavedModel: *_saved_model
# TensorFlow GraphDef: *.pb
# TensorFlow Lite: *.tflite
# TensorFlow Edge TPU: *_edgetpu.tflite
# PaddlePaddle: *_paddle_model
def __init__(self, weights='yolov8n.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):
"""
Ultralytics YOLO MultiBackend class for python inference on various backends
Args:
weights: the path to the weights file. Defaults to yolov8n.pt
device: The device to run the model on.
dnn: If you want to use OpenCV's DNN module to run the inference, set this to True. Defaults to
False
data: a dictionary containing the following keys:
fp16: If true, will use half precision. Defaults to False
fuse: whether to fuse the model or not. Defaults to True
Supported format and their usage:
| Platform | weights |
|-----------------------|------------------|
| PyTorch | *.pt |
| TorchScript | *.torchscript |
| ONNX Runtime | *.onnx |
| ONNX OpenCV DNN | *.onnx --dnn |
| OpenVINO | *.xml |
| CoreML | *.mlmodel |
| TensorRT | *.engine |
| TensorFlow SavedModel | *_saved_model |
| TensorFlow GraphDef | *.pb |
| TensorFlow Lite | *.tflite |
| TensorFlow Edge TPU | *_edgetpu.tflite |
| PaddlePaddle | *_paddle_model |
"""
super().__init__()
w = str(weights[0] if isinstance(weights, list) else weights)
nn_module = isinstance(weights, torch.nn.Module)
@ -215,6 +229,15 @@ class AutoBackend(nn.Module):
self.__dict__.update(locals()) # assign all variables to self
def forward(self, im, augment=False, visualize=False):
"""
Runs inference on the given model
Args:
im: the image tensor
augment: whether to augment the image. Defaults to False
visualize: if True, then the network will output the feature maps of the last convolutional layer.
Defaults to False
"""
# YOLOv5 MultiBackend inference
b, ch, h, w = im.shape # batch, channel, height, width
if self.fp16 and im.dtype != torch.float16:
@ -297,10 +320,21 @@ class AutoBackend(nn.Module):
return self.from_numpy(y)
def from_numpy(self, x):
"""
`from_numpy` converts a numpy array to a tensor
Args:
x: the numpy array to convert
"""
return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x
def warmup(self, imgsz=(1, 3, 640, 640)):
# Warmup model by running inference once
"""
Warmup model by running inference once
Args:
imgsz: the size of the image you want to run inference on.
"""
warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton, self.nn_module
if any(warmup_types) and (self.device.type != 'cpu' or self.triton):
im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input
@ -309,6 +343,12 @@ class AutoBackend(nn.Module):
@staticmethod
def _model_type(p='path/to/model.pt'):
"""
This function takes a path to a model file and returns the model type
Args:
p: path to the model file. Defaults to path/to/model.pt
"""
# Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx
# types = [pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle]
from ultralytics.yolo.engine.exporter import export_formats
@ -323,6 +363,12 @@ class AutoBackend(nn.Module):
@staticmethod
def _load_metadata(f=Path('path/to/meta.yaml')):
"""
> Loads the metadata from a yaml file
Args:
f: The path to the metadata file.
"""
from ultralytics.yolo.utils.files import yaml_load
# Load metadata from meta.yaml if it exists

72
ultralytics/nn/tasks.py
View File

@ -17,11 +17,36 @@ from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, initialize_wei
class BaseModel(nn.Module):
# YOLOv5 base model
'''
The BaseModel class is a base class for all the models in the Ultralytics YOLO family.
'''
def forward(self, x, profile=False, visualize=False):
return self._forward_once(x, profile, visualize) # single-scale inference, train
"""
> `forward` is a wrapper for `_forward_once` that runs the model on a single scale
Args:
x: the input image
profile: whether to profile the model. Defaults to False
visualize: if True, will return the intermediate feature maps. Defaults to False
Returns:
The output of the network.
"""
return self._forward_once(x, profile, visualize)
def _forward_once(self, x, profile=False, visualize=False):
"""
> Forward pass of the network
Args:
x: input to the model
profile: if True, the time taken for each layer will be printed. Defaults to False
visualize: If True, it will save the feature maps of the model. Defaults to False
Returns:
The last layer of the model.
"""
y, dt = [], [] # outputs
for m in self.model:
if m.f != -1: # if not from previous layer
@ -36,6 +61,15 @@ class BaseModel(nn.Module):
return x
def _profile_one_layer(self, m, x, dt):
"""
It takes a model, an input, and a list of times, and it profiles the model on the input, appending
the time to the list
Args:
m: the model
x: the input image
dt: list of time taken for each layer
"""
c = m == self.model[-1] # is final layer, copy input as inplace fix
o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPs
t = time_sync()
@ -48,7 +82,13 @@ class BaseModel(nn.Module):
if c:
LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s} Total")
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
def fuse(self):
"""
> It takes a model and fuses the Conv2d() and BatchNorm2d() layers into a single layer
Returns:
The model is being returned.
"""
LOGGER.info('Fusing layers... ')
for m in self.model.modules():
if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
@ -58,11 +98,27 @@ class BaseModel(nn.Module):
self.info()
return self
def info(self, verbose=False, imgsz=640): # print model information
def info(self, verbose=False, imgsz=640):
"""
Prints model information
Args:
verbose: if True, prints out the model information. Defaults to False
imgsz: the size of the image that the model will be trained on. Defaults to 640
"""
model_info(self, verbose, imgsz)
def _apply(self, fn):
# Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
"""
`_apply()` is a function that applies a function to all the tensors in the model that are not
parameters or registered buffers
Args:
fn: the function to apply to the model
Returns:
A model that is a Detect() object.
"""
self = super()._apply(fn)
m = self.model[-1] # Detect()
if isinstance(m, (Detect, Segment)):
@ -72,6 +128,12 @@ class BaseModel(nn.Module):
return self
def load(self, weights):
"""
> This function loads the weights of the model from a file
Args:
weights: The weights to load into the model.
"""
# Force all tasks to implement this function
raise NotImplementedError("This function needs to be implemented by derived classes!")