save

onnx_ir.save(model, path, format=None, external_data=None, size_threshold_bytes=256, callback=None)

将 ONNX 模型保存到文件。

调用后模型保持不变。如果任何现有的外部张量引用了提供的 external_data 路径，则在外部数据被覆盖后，它将失效。要获得有效模型，请使用 load() 加载新保存的模型，或提供一个模型中任何张量当前未引用的不同外部数据路径。

提示

可以通过传递回调函数来实现一个简单的进度条，如下所示

import onnx_ir as ir
import tqdm

with tqdm.tqdm() as pbar:
    total_set = False

    def callback(tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo) -> None:
        nonlocal total_set
        if not total_set:
            pbar.total = metadata.total
            total_set = True

        pbar.update()
        pbar.set_description(f"Saving {tensor.name} ({tensor.dtype}, {tensor.shape}) at offset {metadata.offset}")

    ir.save(
        ...,
        callback=callback,
    )
参数:

  • model (Model) – 要保存的模型。

  • path (str | PathLike) – 保存模型的路径。例如“model.onnx”。

  • format (str | None) – 文件的格式（例如 protobuf、textproto、json 等）。如果为 None，则从文件扩展名推断格式。

  • external_data (str | PathLike | None) – 保存外部数据的相对路径。指定时，模型中的所有初始化器都将转换为外部数据并保存到指定目录。如果为 None，则所有张量都将 unmodified 保存。也就是说，如果模型中的张量已经是外部的，它将以相同的外部信息保存；如果张量不是外部的，它将以 ONNX Proto 消息的形式序列化。

  • size_threshold_bytes (int) – 如果张量大小（以字节为单位）大于此阈值，则保存到外部数据。仅当设置了 external_data 时才有效。

  • callback (Callable[[TensorProtocol, CallbackInfo], None] | None) – 为每个保存到外部数据以进行调试或日志记录的张量调用的回调函数。

引发:

ValueError – 如果外部数据路径是绝对路径。

返回类型:

无