Skip to content

Introduce Options in ONNX Export

Anthony Correia requested to merge anthonyc/onnx_export_with_options into main

This MR introduces options in ONNX exports.

Options options is a set of strings passed to

  • in python: model.to_onnx(path, options=<options>)
  • in the export_model_to_onnx.py script: --options <options>

For instance, options can be {"use_fp16"}. In this case, the model will be converted to fp16 before being saved to ONNX. The latter is also implemented in this MR. Another application is the option use_trt_scatter which allows to use the TensorRT custom scatter_add operation.

In order for the embedding model to benefit from this options parameter, it was defined in ModelBase.

Merge request reports