Ivadomed: ONNX model

[andreanne-lemay]

Making deep learning models accessible in SCT creates a new need for cpu optimized models since most users don't have GPUs. Exporting our torch models in ONNX format leads to reduced memory usage and faster inference. ONNX is already integrated in PyTorch (documentation here) which makes it very easy to use.

To save model:

torch.onnx.export(model, model_input, "model.onnx")

For inference:

import onnxruntime

numpy_img = nib.load(filename).get_fdata()

ort_session = onnxruntime.InferenceSession("model.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: numpy_img}
ort_outs = ort_session.run(None, ort_inputs)

Full example here.

There is also an option to have dynamic shapes using the option dynamic_axis. Fixed size input is not required (unlike what I thought at first @jcohenadad ).

I tried ONNX models on 2 models: t2star_sc (2D Unet with images of 160x160) and a custom 3D unet (images of size 512x256x16).

ONNX has a faster inference time and uses less memory than torch models when used on CPU. The model size is the same in .onnx or .pt format. Convolutions in PyTorch are not optimized for cpu, which results in a big memory improvement when using ONNX on 3D model (conv3d requires a lot of memory). The improvement is not as noted for the 2D model, but the inference time and memory consumption are still better.

Also, the warnings induced by outdated models disappear with ONNX format ( issue #225 ).

I suggest to save our "production" models in this format (models stored on osf). We could keep .pt models for development purposes since .onnx format is not very friendly for training.

[charleygros]

Wow!!! @andreanne-lemay this is gold!

[olix86]

@andreanne-lemay can we close this since #234 is merged? :)

[andreanne-lemay]

@andreanne-lemay can we close this since #234 is merged? :)

I still have a PR on SCT related to ONNX that is not merged yet (https://github.com/neuropoly/spinalcordtoolbox/pull/2735). I was waiting to merge this PR before closing the issue :)

[jcohenadad]

Merged! https://github.com/neuropoly/spinalcordtoolbox/pull/2735