save_model

onnx_safetensors.save_model(model, model_path, /, *, external_data=None, size_threshold=0, max_shard_size=None)[source]

Save an ONNX model to a file with external data in a safetensors file.

Added in version 1.3.0: Added the function.

Parameters:

  • model (TModel) – ONNX model to save.

  • model_path (str | PathLike) – Path to the ONNX model file. E.g. “model.onnx”.

  • external_data (str | PathLike | None) – Path to the safetensors file relative to the ONNX model file. E.g. “model.safetensors”. If not provided, it will be derived from the model_path by replacing the extension with “.safetensors”.

  • size_threshold (int) – Minimum size in bytes for a tensor to be saved. Default is 0, which saves all tensors.

  • max_shard_size (int | str | None) – Maximum size in bytes for a checkpoint before being sharded. If expressed as a string, needs to be digits followed by a unit (like “5GB” or “100MB”). If None, no sharding is performed. When sharding is enabled, multiple safetensors files will be created with names like “model-00001-of-00003.safetensors”, and an index file “model.safetensors.index.json” will be created to map tensors to their respective shard files.

Raises:

ValueError – If external_data does not end with “.safetensors”.

Return type:

None