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

README and Docs updates with A100 TensorRT times (#270)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
2023-01-11 21:54:41 +01:00
committed by GitHub
parent 216cf2ddb6
commit e18ae9d8e1
10 changed files with 250 additions and 241 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
ultralytics/nn/autobackend.py
View File

@ -22,32 +22,31 @@ class AutoBackend(nn.Module):
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
MultiBackend class for python inference on various platforms using Ultralytics YOLO.
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
weights (str): The path to the weights file. Default: 'yolov8n.pt'
device (torch.device): The device to run the model on.
dnn (bool): Use OpenCV's DNN module for inference if True, defaults to False.
data (dict): Additional data, optional
fp16 (bool): If True, use half precision. Default: False
fuse (bool): Whether to fuse the model or not. Default: 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 |
Supported formats and their usage:
Platform | Weights Format
-----------------------|------------------
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)
@ -234,15 +233,16 @@ class AutoBackend(nn.Module):
def forward(self, im, augment=False, visualize=False):
"""
Runs inference on the given model
Runs inference on the YOLOv8 MultiBackend 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
im (torch.tensor): The image tensor to perform inference on.
augment (bool): whether to perform data augmentation during inference, defaults to False
visualize (bool): whether to visualize the output predictions, defaults to False
Returns:
(tuple): Tuple containing the raw output tensor, and the processed output for visualization (if visualize=True)
"""
# YOLOv5 MultiBackend inference
b, ch, h, w = im.shape # batch, channel, height, width
if self.fp16 and im.dtype != torch.float16:
im = im.half() # to FP16
@ -325,19 +325,25 @@ class AutoBackend(nn.Module):
def from_numpy(self, x):
"""
`from_numpy` converts a numpy array to a tensor
Convert a numpy array to a tensor.
Args:
x: the numpy array to convert
"""
Args:
x (numpy.ndarray): The array to be converted.
Returns:
(torch.tensor): The converted tensor
"""
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
Warm up the model by running one forward pass with a dummy input.
Args:
imgsz: the size of the image you want to run inference on.
imgsz (tuple): The shape of the dummy input tensor in the format (batch_size, channels, height, width)
Returns:
(None): This method runs the forward pass and don't return any value
"""
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):

51
ultralytics/nn/tasks.py
View File

@ -17,35 +17,36 @@ from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, initialize_wei
class BaseModel(nn.Module):
'''
The BaseModel class is a base class for all the models in the Ultralytics YOLO family.
'''
"""
The BaseModel class serves as a base class for all the models in the Ultralytics YOLO family.
"""
def forward(self, x, profile=False, visualize=False):
"""
> `forward` is a wrapper for `_forward_once` that runs the model on a single scale
Forward pass of the model on a single scale.
Wrapper for `_forward_once` method.
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
x (torch.tensor): The input image tensor
profile (bool): Whether to profile the model, defaults to False
visualize (bool): Whether to return the intermediate feature maps, defaults to False
Returns:
The output of the network.
(torch.tensor): 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
Perform a forward pass through 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
x (torch.tensor): The input tensor to the model
profile (bool): Print the computation time of each layer if True, defaults to False.
visualize (bool): Save the feature maps of the model if True, defaults to False
Returns:
The last layer of the model.
(torch.tensor): The last output of the model.
"""
y, dt = [], [] # outputs
for m in self.model:
@ -62,13 +63,15 @@ class BaseModel(nn.Module):
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
Profile the computation time and FLOPs of a single layer of the model on a given input. Appends the results to the provided list.
Args:
m: the model
x: the input image
dt: list of time taken for each layer
m (nn.Module): The layer to be profiled.
x (torch.Tensor): The input data to the layer.
dt (list): A list to store the computation time of the layer.
Returns:
None
"""
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
@ -84,10 +87,10 @@ class BaseModel(nn.Module):
def fuse(self):
"""
> It takes a model and fuses the Conv2d() and BatchNorm2d() layers into a single layer
Fuse the `Conv2d()` and `BatchNorm2d()` layers of the model into a single layer, in order to improve the computation efficiency.
Returns:
The model is being returned.
(nn.Module): The fused model is returned.
"""
LOGGER.info('Fusing layers... ')
for m in self.model.modules():
@ -103,8 +106,8 @@ class BaseModel(nn.Module):
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
verbose (bool): if True, prints out the model information. Defaults to False
imgsz (int): the size of the image that the model will be trained on. Defaults to 640
"""
model_info(self, verbose, imgsz)
@ -129,10 +132,10 @@ class BaseModel(nn.Module):
def load(self, weights):
"""
> This function loads the weights of the model from a file
This function loads the weights of the model from a file
Args:
weights: The weights to load into the model.
weights (str): 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!")