Export, detect and validation with TensorRT engine file (#5699)
* Export and detect with TensorRT engine file * Resolve `isort` * Make validation works with TensorRT engine * feat: update export docstring * feat: change suffix from *.trt to *.engine * feat: get rid of pycuda * feat: make compatiable with val.py * feat: support detect with fp16 engine * Add Lite to Edge TPU string * Remove *.trt comment * Revert to standard success logger.info string * Fix Deprecation Warning ``` export.py:310: DeprecationWarning: Use build_serialized_network instead. with builder.build_engine(network, config) as engine, open(f, 'wb') as t: ``` * Revert deprecation warning fix @imyhxy it seems we can't apply the deprecation warning fix because then export fails, so I'm reverting my previous change here. * Update export.py * Update export.py * Update common.py * export onnx to file before building TensorRT engine file * feat: triger ONNX export failed early * feat: load ONNX model from file Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
1 parent
f17c86b7f0
commit
7a39803476
4 files changed
+90
-11
No files matched your search
@@ -77,11 +77,11 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
||||
# Load model
|
||||
device = select_device(device)
|
||||
model = DetectMultiBackend(weights, device=device, dnn=dnn)
|
||||
stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
|
||||
stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine
|
||||
imgsz = check_img_size(imgsz, s=stride) # check image size
|
||||
|
||||
# Half
|
||||
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
||||
half &= (pt or engine) and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
||||
if pt:
|
||||
model.model.half() if half else model.model.float()
|
||||
|
||||
|
||||
Reference in new issue
Block a user