3 ms·

So I've been looking into ONNX to speedup inference, is there some killer feature I should look at JAX for?

by amrb 4y ago

kkielhofner 4y ago
Generally from what I've seen the biggest inference speedup win with ONNX is to get the model to ONNX then to TRT (TensorRT) - assuming Nvidia hardware. Once in "TRT land" you can play with fp32, fp16, int8 (with calibration), etc. That said (again generally) ONNX does tend to perform better when compared to native (TF savedmodel, pytorch torchscript, whatever). With TensorFlow and Pytorch there are also ways to export/compile directly to TRT but from an ops standpoint I don't prefer this. Certain inference serving solutions like Nvidia Triton Inference Server will even take an ONNX model and then do TRT compilation (with cache!) on the actual inference hardware dynamically at model load time. This is really nice because you can deploy a standard ONNX model across instances and varying GPU hardware and always get TRT optimized and compatible with Compute Capability, etc. Really handy and basically comes down to a few lines of config in the model configuration. I'm not terribly familiar with JAX but I have to imagine there's ONNX export or straight to TRT export somewhere.
chas 4y ago
Jax is a great tool, but it’s really best for training and experimentation. The transformations outlined in this post (amongst others) make it easy to turn simple and straightforward code into high performance parallel code. While this is changing, inferences hasn’t been a historical area of emphasis for the project, so it wouldn’t be my first choice if that was your primary goal.
- sdenton4 4y ago
  Yeah inference is a real weakness. I've used jax2tf to get tf saved models and tflite models, but it was a real journey... There's some effort going into systems for saving and restoring the computation graph for Jax programs, which will help a lot. I'm surprised it didn't happen sooner, as it seems like quite a natural fit with the jax architecture.
UncleOxidant 4y ago
About a year ago I was tasked with comparing ONNX runtime implementations of certain components like convolution with some of our own in-house implementations. There was just no comparison. ONNX runtime has some pretty crazy fast implementations. Lots of magic in there. Concluded that we weren't going to be able to beat those without a lot of effort and expertise that we didn't have in our team.

So I've been looking into ONNX to speedup inference, is there some killer feature I should look at JAX for?

by amrb 4y ago

So I've been looking into ONNX to speedup inference, is there some killer feature I should look at JAX for?

Generally from what I've seen the biggest inference speedup win with ONNX is to get the model to ONNX then to TRT (TensorRT) - assuming Nvidia hardware. Once in "TRT land" you can play with fp32, fp16, int8 (with calibration), etc. That said (again generally) ONNX does tend to perform better when compared to native (TF savedmodel, pytorch torchscript, whatever). With TensorFlow and Pytorch there are also ways to export/compile directly to TRT but from an ops standpoint I don't prefer this. Certain inference serving solutions like Nvidia Triton Inference Server will even take an ONNX model and then do TRT compilation (with cache!) on the actual inference hardware dynamically at model load time. This is really nice because you can deploy a standard ONNX model across instances and varying GPU hardware and always get TRT optimized and compatible with Compute Capability, etc. Really handy and basically comes down to a few lines of config in the model configuration. I'm not terribly familiar with JAX but I have to imagine there's ONNX export or straight to TRT export somewhere.

chas 4y ago

Jax is a great tool, but it’s really best for training and experimentation. The transformations outlined in this post (amongst others) make it easy to turn simple and straightforward code into high performance parallel code. While this is changing, inferences hasn’t been a historical area of emphasis for the project, so it wouldn’t be my first choice if that was your primary goal.

sdenton4 4y ago
Yeah inference is a real weakness. I've used jax2tf to get tf saved models and tflite models, but it was a real journey... There's some effort going into systems for saving and restoring the computation graph for Jax programs, which will help a lot. I'm surprised it didn't happen sooner, as it seems like quite a natural fit with the jax architecture.

UncleOxidant 4y ago

About a year ago I was tasked with comparing ONNX runtime implementations of certain components like convolution with some of our own in-house implementations. There was just no comparison. ONNX runtime has some pretty crazy fast implementations. Lots of magic in there. Concluded that we weren't going to be able to beat those without a lot of effort and expertise that we didn't have in our team.