外部数据

加载具有外部数据的 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 字段存储了该张量数据的存储位置。值必须是以下之一:

  • MESSAGE - 数据存储在 protobuf 消息内的类型特定字段中。

  • RAW - 数据存储在 raw_data 字段中。

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

  • value 未设置 - 旧值。假定数据存储在 raw_data(如果已设置)中,否则存储在消息中。

external_data 字段

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

已识别的键有:

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

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

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

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

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

外部数据文件

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

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