Arlo/tensor rt

[sheridana]

Description

• Adds unragging class for keras layers (useful if unragging is needed prior to model saving)
• Adds WIP functionality for converting keras models to TensorRT models for faster inference when available.
• Todo: clean up, add automatic system checks

Types of changes

• Bugfix
• New feature
• Refactor / Code style update (no logical changes)
• Build / CI changes
• Documentation Update
• Other (explain)

Does this address any currently open issues?

1112

Outside contributors checklist

• Review the guidelines for contributing to this repository
• Read and sign the CLA and add yourself to the authors list
• Make sure you are making a pull request against the develop branch (not main). Also you should start your branch off develop
• Add tests that prove your fix is effective or that your feature works
• Add necessary documentation (if appropriate)

Thank you for contributing to SLEAP!

❤️

[codecov]

Codecov Report

Merging #1138 (24d067a) into develop (b37b34f) will decrease coverage by 0.10%.
The diff coverage is 32.00%.

@@             Coverage Diff             @@
##           develop    #1138      +/-   ##
===========================================
- Coverage    70.05%   69.95%   -0.10%     
===========================================
  Files          131      131              
  Lines        22872    22916      +44     
===========================================
+ Hits         16022    16032      +10     
- Misses        6850     6884      +34     

Impacted Files	Coverage Δ
sleap/nn/data/utils.py	88.88% <28.57%> (-8.99%)	⬇️
sleap/nn/inference.py	78.41% <32.55%> (-1.54%)	⬇️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[sheridana]

Example usage and benchmarking: https://github.com/sheridana/sleap_tensor_rt

Todos:

• We probably want to use real data rather than a tracing batch to get full optimization
• Figure out handling the changed dict keys after unragging in a more principled way
• Get it to work for other predictors
• Not sure best way to add tests for this since it is very system & hardware dependent...
• Maybe add automatic check for precision availability? E.g if FP16 isn't available revert to FP32 immediately. This is already handled internally but is pretty verbose. Tricky to handle this because it's checked deep under the hood, maybe here
• Add support for Int8?