在使用“torch.onnx.export”时，是否有方法包含模型描述？

Question

问题描述
我们希望我们的onnx模型有某种形式的描述，理想情况下还有一些其他元数据，包括我们的内部版本号。目前，我们使用pytorch lightning进行培训，并使用onnxruntime进行推理

下面是一个最小的可执行示例，通过以下方式分配模型描述：

1. 使用torch.onnx.export导出
2. 用onnx.load加载
3. 设置model.doc_string
4. 使用onnx.save导出
5. 用onnxruntime.InferenceSession加载

问题

Using onnx seems unnecessary, is there a way to include a model description when using torch.onnx.export?

可复制示例

>>> import torch
>>> import torchvision
>>> from onnxruntime import InferenceSession
>>>
>>> # get a sample model
>>> dummy_input = torch.randn(10, 3, 224, 224)
>>> _model = torchvision.models.alexnet(pretrained=True)
>>> input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
>>> output_names = [ "output1" ]
>>>
>>> # export onnx
>>> torch.onnx.export(_model, dummy_input, "alexnet.onnx", verbose=False, input_names=input_names, output_names=output_names,strip_doc_string=False)
>>>
>>> # read in the exported onnx for infererence
>>> sess = InferenceSession('alexnet.onnx')
>>> meta = sess.get_modelmeta()
>>>
>>> # review onnx metadata
>>> meta.description
''

这将下载一个演示模型，当我们打印meta.description时，我们可以看到它是空白的

我可以通过使用onnx加载模型并再次保存来设置此描述

>>> import onnx
>>> model = onnx.load('alexnet.onnx')
>>> model.doc_string = 'my_description'
>>> onnx.save(model, 'alexnet2.onnx')
>>> sess = InferenceSession('alexnet2.onnx')
>>> meta = sess.get_modelmeta()
>>>
>>> # review onnx metadata
>>> meta.description
'my_description'