onnx-safetensors

ONNX扩展，用于保存到和从safetensors加载 🤗。

安装

pip install --upgrade onnx-safetensors

用法

将张量加载到ONNX模型中

import onnx_safetensors

# Provide your ONNX model here
model: onnx.ModelProto
tensor_file = "model.safetensors"
# Apply weights from the safetensors file to the model
onnx_safetensors.load_file(model, tensor_file)

将权重保存到safetensors文件

import onnx_safetensors

# Provide your ONNX model here
model: onnx.ModelProto
tensor_file = "model.safetensors"
# Save weights from to the safetensors file
onnx_safetensors.save_file(model, tensor_file, convert_attributes=True)

# Save weights from to the safetensors file and clear the raw_data fields of the ONNX model to reduce its size
# model will be updated inplace
onnx_safetensors.save_file(model, tensor_file, convert_attributes=True, strip_data=True)