Skip to content

Allow using ONNX compiled model in sklearn interface #179

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 21 commits into
base: main
Choose a base branch
from
Draft

Conversation

LeoGrin
Copy link
Collaborator

@LeoGrin LeoGrin commented Feb 11, 2025

Fix #97

Downside compared to PyTorch

We don't do the memory estimation + reduction part, so it's easier to get memory error with large datasets (TODO)

(It seems that if if I export the onnx model on gpu I get less memory error afterward but not sure)

Quick speed analysis:

For 1Kx10 on a T4 GPU node (seems to be quite hardware dependent):

  • first fit: 5s for ONNX (we initialize the session), 0.4s for PyTorch
  • first predict: ~2s for ONNX, 6s for Pytorch
  • 2nd fit: 0.05s for ONNX, 0.15s for PyTorch
  • 2nd predict: 0.6-1.2s for ONNX (depending on where you export it, probably some more gains here?), 0.22s for PyTorch

Seems like for smaller datasets ONNX becomes faster even for 2nd predict.

Performance analysis on real datasets

Seemed good but need to rerun.

TODOs

  • Make memory reduction work on ONNX
  • Allow to set the onnx session from outside to work with sklearn cloning
  • Check more carefully that we haven't degraded performance (e.g on categoricals)
  • Can we make ONNX model faster by exporting it with the right dummy input? How much speed vary by hardware?

@LeoGrin LeoGrin marked this pull request as draft February 11, 2025 13:41
@PriorLabs PriorLabs deleted a comment from CLAassistant Feb 12, 2025
@LeoGrin LeoGrin mentioned this pull request Feb 20, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Feature: Support ONNX conversion for models
1 participant