外部数据

加载带有外部数据的 ONNX 模型

  • [默认] 如果外部数据与模型在同一目录下,只需使用 onnx.load()

import onnx

onnx_model = onnx.load("path/to/the/model.onnx")
  • 如果外部数据在另一个目录下,在使用 onnx.load() 后,使用 load_external_data_for_model() 来指定目录路径并加载

import onnx
from onnx.external_data_helper import load_external_data_for_model

onnx_model = onnx.load("path/to/the/model.onnx", load_external_data=False)
load_external_data_for_model(onnx_model, "data/directory/path/")
# Then the onnx_model has loaded the external data from the specific directory

将 ONNX 模型转换为外部数据

import onnx
from onnx.external_data_helper import convert_model_to_external_data

onnx_model = ... # Your model in memory as ModelProto
convert_model_to_external_data(onnx_model, all_tensors_to_one_file=True, location="filename", size_threshold=1024, convert_attribute=False)
# Must be followed by save_model to save the converted model to a specific path
onnx.save_model(onnx_model, "path/to/save/the/model.onnx")
# Then the onnx_model has converted raw data as external data and saved to specific directory

将 ONNX 模型转换为外部数据并保存

import onnx

onnx_model = ... # Your model in memory as ModelProto
onnx.save_model(onnx_model, "path/to/save/the/model.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="filename", size_threshold=1024, convert_attribute=False)
# Then the onnx_model has converted raw data as external data and saved to specific directory

带有外部数据的模型的 onnx.checker

带有外部数据的模型(小于 2GB)

当前检查器支持检查带有外部数据的模型。请向检查器指定已加载的 onnx 模型或模型路径。

大于 2GB 的大型模型

但是,对于大于 2GB 的模型,请使用 onnx.checker 的模型路径,并且外部数据需要位于同一目录下。

import onnx

onnx.checker.check_model("path/to/the/model.onnx")
# onnx.checker.check_model(loaded_onnx_model) will fail if given >2GB model

TensorProto: data_location 和 external_data 字段

TensorProto 消息类型中有两个与外部数据相关的字段。

data_location 字段

data_location 字段存储此张量的数据位置。值必须是以下之一:

  • DEFAULT - 数据存储在 protobuf 消息内部。数据存储在 raw_data(如果设置)中,否则存储在特定类型字段中。

  • EXTERNAL - 数据存储在由 external_data 字段描述的外部位置。

如果未设置,则其行为与值为 DEFAULT 时相同。

external_data 字段

external_data 字段存储描述数据位置的字符串键值对

已识别的键为

  • "location" (必需) - 相对于 ONNX protobuf 模型存储的文件系统目录的文件路径。禁止使用上级目录路径组件(例如...),解析时应将其去除。

  • "offset" (可选) - 存储数据开始的字节位置。以字符串形式存储的整数。偏移量值应为页面大小(通常为 4kb)的倍数,以支持 mmap。在 Windows 上,偏移量值应为 VirtualAlloc 分配粒度(通常为 64kb)的倍数,以启用内存映射

  • "length" (可选) - 包含数据字节数。以字符串形式存储的整数。

  • "checksum" (可选) - ‘location’ 键下指定文件的 SHA1 摘要。

ONNX 文件加载后,所有 external_data 字段都可以通过一个附加键 ("basepath") 进行更新,该键存储 ONNX 模型文件加载的目录路径。

外部数据文件

存储在外部数据文件中的数据将采用与当前 ONNX 实现中 raw_data 字段使用的相同二进制字节字符串格式。

参考 https://github.com/onnx/onnx/pull/678