Upgrade onnxsim to v0.4.1 (#8632)
* upgrade onnxsim to v0.4.1 Signed-off-by: daquexian <daquexian566@gmail.com> * Update export.py * Update export.py * Update export.py * Update export.py * Update export.py Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
1 parent
fbe67e4653
commit
92e47b85d9
2 files changed
+5
-6
No files matched your search
@@ -152,13 +152,12 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
|
||||
# Simplify
|
||||
if simplify:
|
||||
try:
|
||||
check_requirements(('onnx-simplifier',))
|
||||
cuda = torch.cuda.is_available()
|
||||
check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
|
||||
import onnxsim
|
||||
|
||||
LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
|
||||
model_onnx, check = onnxsim.simplify(model_onnx,
|
||||
dynamic_input_shape=dynamic,
|
||||
input_shapes={'images': list(im.shape)} if dynamic else None)
|
||||
model_onnx, check = onnxsim.simplify(model_onnx)
|
||||
assert check, 'assert check failed'
|
||||
onnx.save(model_onnx, f)
|
||||
except Exception as e:
|
||||
@@ -493,7 +492,7 @@ def run(
|
||||
imgsz *= 2 if len(imgsz) == 1 else 1 # expand
|
||||
assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
|
||||
if optimize:
|
||||
assert device.type != 'cuda', '--optimize not compatible with cuda devices, i.e. use --device cpu'
|
||||
assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'
|
||||
|
||||
# Input
|
||||
gs = int(max(model.stride)) # grid size (max stride)
|
||||
|
||||
Reference in new issue
Block a user