mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 07:06:36 +00:00
[AZ-171] Enable dynamic batch size for ONNX, TensorRT, and CoreML exports
Made-with: Cursor
This commit is contained in:
+4
-1
@@ -37,6 +37,7 @@ def export_onnx(model_path, batch_size=None):
|
||||
format="onnx",
|
||||
imgsz=constants.config.export.onnx_imgsz,
|
||||
batch=batch_size,
|
||||
dynamic=True,
|
||||
simplify=True,
|
||||
nms=True,
|
||||
)
|
||||
@@ -47,13 +48,15 @@ def export_coreml(model_path):
|
||||
model.export(
|
||||
format="coreml",
|
||||
imgsz=constants.config.export.onnx_imgsz,
|
||||
dynamic=True,
|
||||
)
|
||||
|
||||
|
||||
def export_tensorrt(model_path):
|
||||
YOLO(model_path).export(
|
||||
format='engine',
|
||||
batch=4,
|
||||
batch=8,
|
||||
dynamic=True,
|
||||
half=True,
|
||||
simplify=True,
|
||||
nms=True
|
||||
|
||||
Reference in New Issue
Block a user