ONNX模型的修改可以通过转成JSON文件再将JSON文件转回ONNX。
以一个Conv算子构成的模型为例:
ONNX转JSON
使用MessageToJson进行转换,如下将一个conv算子构成的ONNX模型转成JSON
import onnx
from google.protobuf.json_format import MessageToJson, Parse
onnx_model = onnx.load("Conv.onnx")
message = MessageToJson(onnx_model)
with open("conv.json", "w") as fo:
fo.write(message)
转换之后模型由JSON表示:
{
"irVersion": "7",
"producerName": "onnx-example",
"graph": {
"node": [
{
"input": [
"X",
"W",
"B"
],
"output": [
"Y"
],
"opType": "Conv",
"attribute": [
{
"name": "strides",
"ints": [
"2",
"2"
],
"type": "INTS"
}
]
}
],
"name": "test_conv_mode",
"initializer": [
{
"dims": [
"2",
"2",
"3",
"3"
],
"dataType": 1,
"floatData": [
-0.6021352410316467,
...
],
"name": "W"
},
{
"dims": [
"2"
],
"dataType": 1,
"floatData": [
1.0,
2.0
],
"name": "B"
}
],
"input": [
{
"name": "X",
"type": {
"tensorType": {
"elemType": 1,
"shape": {
"dim": [
{
"dimValue": "1"
},
{
"dimValue": "2"
},
{
"dimValue": "4"
},
{
"dimValue": "4"
}
]
}
}
}
}
],
"output": [
{
"name": "Y",
"type": {
"tensorType": {
"elemType": 1,
"shape": {
"dim": [
{
"dimValue": "1"
},
{
"dimValue": "2"
},
{
"dimValue": "2"
},
{
"dimValue": "2"
}
]
}
}
}
}
]
},
"opsetImport": [
{
"version": "12"
}
]
}
JSON转ONNX
在上述json中增加padding属性:
使用Parse模块解析成ONNX:
with open("conv.json", "r") as fi:
onnx_json = json.loads(fi.read())
onnx_str = json.dumps(onnx_json)
convert_model = Parse(onnx_str, onnx.ModelProto())
onnx.save(convert_model, "Conv_1.onnx")
存储为onnx格式之后再可视化可以看到Conv算子增加了padding属性。