mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

AutoShape() models as DetectMultiBackend() instances (#5845)

Browse Source 
* Update AutoShape()

* autodownload ONNX

* Cleanup

* Finish updates

* Add Usage

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* fix device

* Update hubconf.py

* Update common.py

* smart param selection

* autodownload all formats

* autopad only pytorch models

* new_shape edits

* stride tensor fix

* Cleanup
This commit is contained in:
Glenn Jocher 2021-12-04 15:00:07 +01:00 committed by GitHub
parent d885799c71
commit 7bf04d9bbf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 35 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
export.py
View File

@ -411,7 +411,7 @@ def parse_opt():
parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization') parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes') parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
parser.add_argument('--opset', type=int, default=13, help='ONNX: opset version') parser.add_argument('--opset', type=int, default=14, help='ONNX: opset version')
parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log') parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)') parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep') parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')

14
hubconf.py
View File

@ -5,6 +5,7 @@ PyTorch Hub models https://pytorch.org/hub/ultralytics_yolov5/
Usage: Usage:
import torch import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
model = torch.hub.load('ultralytics/yolov5:master', 'custom', 'path/to/yolov5s.onnx') # file from branch
""" """
import torch import torch
@ -27,26 +28,25 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
""" """
from pathlib import Path from pathlib import Path
from models.common import AutoShape from models.common import AutoShape, DetectMultiBackend
from models.experimental import attempt_load
from models.yolo import Model from models.yolo import Model
from utils.downloads import attempt_download from utils.downloads import attempt_download
from utils.general import check_requirements, intersect_dicts, set_logging from utils.general import check_requirements, intersect_dicts, set_logging
from utils.torch_utils import select_device from utils.torch_utils import select_device
file = Path(__file__).resolve()
check_requirements(exclude=('tensorboard', 'thop', 'opencv-python')) check_requirements(exclude=('tensorboard', 'thop', 'opencv-python'))
set_logging(verbose=verbose) set_logging(verbose=verbose)
save_dir = Path('') if str(name).endswith('.pt') else file.parent name = Path(name)
path = (save_dir / name).with_suffix('.pt') # checkpoint path path = name.with_suffix('.pt') if name.suffix == '' else name # checkpoint path
try: try:
device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device) device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)
if pretrained and channels == 3 and classes == 80: if pretrained and channels == 3 and classes == 80:
model = attempt_load(path, map_location=device) # download/load FP32 model model = DetectMultiBackend(path, device=device) # download/load FP32 model
# model = models.experimental.attempt_load(path, map_location=device) # download/load FP32 model
else: else:
cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path cfg = list((Path(__file__).parent / 'models').rglob(f'{path.name}.yaml'))[0] # model.yaml path
model = Model(cfg, channels, classes) # create model model = Model(cfg, channels, classes) # create model
if pretrained: if pretrained:
ckpt = torch.load(attempt_download(path), map_location=device) # load ckpt = torch.load(attempt_download(path), map_location=device) # load

40
models/common.py
View File

@ -276,7 +276,7 @@ class Concat(nn.Module):
class DetectMultiBackend(nn.Module): class DetectMultiBackend(nn.Module):
# YOLOv5 MultiBackend class for python inference on various backends # YOLOv5 MultiBackend class for python inference on various backends
def __init__(self, weights='yolov5s.pt', device=None, dnn=True): def __init__(self, weights='yolov5s.pt', device=None, dnn=False):
# Usage: # Usage:
# PyTorch: weights = *.pt # PyTorch: weights = *.pt
# TorchScript: *.torchscript # TorchScript: *.torchscript
@ -287,6 +287,8 @@ class DetectMultiBackend(nn.Module):
# ONNX Runtime: *.onnx # ONNX Runtime: *.onnx
# OpenCV DNN: *.onnx with dnn=True # OpenCV DNN: *.onnx with dnn=True
# TensorRT: *.engine # TensorRT: *.engine
from models.experimental import attempt_download, attempt_load # scoped to avoid circular import
super().__init__() super().__init__()
w = str(weights[0] if isinstance(weights, list) else weights) w = str(weights[0] if isinstance(weights, list) else weights)
suffix = Path(w).suffix.lower() suffix = Path(w).suffix.lower()
@ -294,6 +296,7 @@ class DetectMultiBackend(nn.Module):
check_suffix(w, suffixes) # check weights have acceptable suffix check_suffix(w, suffixes) # check weights have acceptable suffix
pt, jit, onnx, engine, tflite, pb, saved_model, coreml = (suffix == x for x in suffixes) # backend booleans pt, jit, onnx, engine, tflite, pb, saved_model, coreml = (suffix == x for x in suffixes) # backend booleans
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
attempt_download(w) # download if not local
if jit: # TorchScript if jit: # TorchScript
LOGGER.info(f'Loading {w} for TorchScript inference...') LOGGER.info(f'Loading {w} for TorchScript inference...')
@ -303,11 +306,12 @@ class DetectMultiBackend(nn.Module):
d = json.loads(extra_files['config.txt']) # extra_files dict d = json.loads(extra_files['config.txt']) # extra_files dict
stride, names = int(d['stride']), d['names'] stride, names = int(d['stride']), d['names']
elif pt: # PyTorch elif pt: # PyTorch
from models.experimental import attempt_load # scoped to avoid circular import
model = attempt_load(weights, map_location=device) model = attempt_load(weights, map_location=device)
stride = int(model.stride.max()) # model stride stride = int(model.stride.max()) # model stride
names = model.module.names if hasattr(model, 'module') else model.names # get class names names = model.module.names if hasattr(model, 'module') else model.names # get class names
self.model = model # explicitly assign for to(), cpu(), cuda(), half()
elif coreml: # CoreML elif coreml: # CoreML
LOGGER.info(f'Loading {w} for CoreML inference...')
import coremltools as ct import coremltools as ct
model = ct.models.MLModel(w) model = ct.models.MLModel(w)
elif dnn: # ONNX OpenCV DNN elif dnn: # ONNX OpenCV DNN
@ -316,7 +320,7 @@ class DetectMultiBackend(nn.Module):
net = cv2.dnn.readNetFromONNX(w) net = cv2.dnn.readNetFromONNX(w)
elif onnx: # ONNX Runtime elif onnx: # ONNX Runtime
LOGGER.info(f'Loading {w} for ONNX Runtime inference...') LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
check_requirements(('onnx', 'onnxruntime-gpu' if torch.has_cuda else 'onnxruntime')) check_requirements(('onnx', 'onnxruntime-gpu' if torch.cuda.is_available() else 'onnxruntime'))
import onnxruntime import onnxruntime
session = onnxruntime.InferenceSession(w, None) session = onnxruntime.InferenceSession(w, None)
elif engine: # TensorRT elif engine: # TensorRT
@ -376,7 +380,7 @@ class DetectMultiBackend(nn.Module):
if self.pt: # PyTorch if self.pt: # PyTorch
y = self.model(im) if self.jit else self.model(im, augment=augment, visualize=visualize) y = self.model(im) if self.jit else self.model(im, augment=augment, visualize=visualize)
return y if val else y[0] return y if val else y[0]
elif self.coreml: # CoreML *.mlmodel elif self.coreml: # CoreML
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3) im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
im = Image.fromarray((im[0] * 255).astype('uint8')) im = Image.fromarray((im[0] * 255).astype('uint8'))
# im = im.resize((192, 320), Image.ANTIALIAS) # im = im.resize((192, 320), Image.ANTIALIAS)
@ -433,24 +437,28 @@ class AutoShape(nn.Module):
# YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS # YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
conf = 0.25 # NMS confidence threshold conf = 0.25 # NMS confidence threshold
iou = 0.45 # NMS IoU threshold iou = 0.45 # NMS IoU threshold
classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs agnostic = False # NMS class-agnostic
multi_label = False # NMS multiple labels per box multi_label = False # NMS multiple labels per box
classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
max_det = 1000 # maximum number of detections per image max_det = 1000 # maximum number of detections per image
def __init__(self, model): def __init__(self, model):
super().__init__() super().__init__()
LOGGER.info('Adding AutoShape... ') LOGGER.info('Adding AutoShape... ')
copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=()) # copy attributes copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=()) # copy attributes
self.dmb = isinstance(model, DetectMultiBackend) # DetectMultiBackend() instance
self.pt = not self.dmb or model.pt # PyTorch model
self.model = model.eval() self.model = model.eval()
def _apply(self, fn): def _apply(self, fn):
# Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
self = super()._apply(fn) self = super()._apply(fn)
m = self.model.model[-1] # Detect() if self.pt:
m.stride = fn(m.stride) m = self.model.model.model[-1] if self.dmb else self.model.model[-1] # Detect()
m.grid = list(map(fn, m.grid)) m.stride = fn(m.stride)
if isinstance(m.anchor_grid, list): m.grid = list(map(fn, m.grid))
m.anchor_grid = list(map(fn, m.anchor_grid)) if isinstance(m.anchor_grid, list):
m.anchor_grid = list(map(fn, m.anchor_grid))
return self return self
@torch.no_grad() @torch.no_grad()
@ -465,7 +473,7 @@ class AutoShape(nn.Module):
# multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
t = [time_sync()] t = [time_sync()]
p = next(self.model.parameters()) # for device and type p = next(self.model.parameters()) if self.pt else torch.zeros(1) # for device and type
if isinstance(imgs, torch.Tensor): # torch if isinstance(imgs, torch.Tensor): # torch
with amp.autocast(enabled=p.device.type != 'cpu'): with amp.autocast(enabled=p.device.type != 'cpu'):
return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
@ -489,8 +497,8 @@ class AutoShape(nn.Module):
g = (size / max(s)) # gain g = (size / max(s)) # gain
shape1.append([y * g for y in s]) shape1.append([y * g for y in s])
imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape shape1 = [make_divisible(x, self.stride) for x in np.stack(shape1, 0).max(0)] # inference shape
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad x = [letterbox(im, new_shape=shape1 if self.pt else size, auto=False)[0] for im in imgs] # pad
x = np.stack(x, 0) if n > 1 else x[0][None] # stack x = np.stack(x, 0) if n > 1 else x[0][None] # stack
x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32 x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
@ -498,12 +506,12 @@ class AutoShape(nn.Module):
with amp.autocast(enabled=p.device.type != 'cpu'): with amp.autocast(enabled=p.device.type != 'cpu'):
# Inference # Inference
y = self.model(x, augment, profile)[0] # forward y = self.model(x, augment, profile) # forward
t.append(time_sync()) t.append(time_sync())
# Post-process # Post-process
y = non_max_suppression(y, self.conf, iou_thres=self.iou, classes=self.classes, y = non_max_suppression(y if self.dmb else y[0], self.conf, iou_thres=self.iou, classes=self.classes,
multi_label=self.multi_label, max_det=self.max_det) # NMS agnostic=self.agnostic, multi_label=self.multi_label, max_det=self.max_det) # NMS
for i in range(n): for i in range(n):
scale_coords(shape1, y[i][:, :4], shape0[i]) scale_coords(shape1, y[i][:, :4], shape0[i])

4
utils/general.py
View File

@ -455,7 +455,9 @@ def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1):
def make_divisible(x, divisor): def make_divisible(x, divisor):
# Returns x evenly divisible by divisor # Returns nearest x divisible by divisor
if isinstance(divisor, torch.Tensor):
divisor = int(divisor.max()) # to int
return math.ceil(x / divisor) * divisor return math.ceil(x / divisor) * divisor