search

Deci-AI / super-gradients

Easily train or fine-tune SOTA computer vision models with one open source training library. The home of Yolo-NAS.
https://www.supergradients.com
Apache License 2.0
4.59k stars 510 forks source link

Onnx model supporting batch prediction #1964

Closed eyv249 closed 7 months ago

eyv249 commented 7 months ago

💡 Your Question

Hi, I am exporting my model as such: best_model = models.get( model_name=Models.PP_YOLOE_L, num_classes=8, checkpoint_path="yolox/pp_long_run_filtered_classes/RUN_20240412_081806_035870/average_model.pth" ) export_result = best_model.export("pp_yolo.onnx", input_image_shape=(640,640), confidence_threshold = 0.75)

I need to be able to feed the model batches of images, however the model is restricted to input on this format: "Model expects input image of shape [1, 3, 640, 640]", where I would need [N, 3, 640, 640].

Any help? Thanks!

Versions

No response

shaydeci commented 7 months ago

@eyv249 you can control the batch size of the exported model by passing it explicitly through the batch_size argument.

eyv249 commented 7 months ago

Thank you @shaydeci! Can I set this batch_size to be dynamic or does it have to be fixed?

shaydeci commented 7 months ago

@eyv249 We do not support dynamic batch size at the moment, however contributions are always welcome (: