如何用Java修改ONNX模型名称以适配Cumulocity Zementis上传

核心原理 #

ONNX模型本质是基于Protocol Buffer（ProtoBuf）定义的序列化文件，因此可以直接解析并修改其ProtoBuf结构中的元数据字段，绕开ONNX Runtime库中不可变的OnnxModelMetaData限制，实现模型名称的修改。

具体实现步骤 #

• 引入ONNX ProtoBuf依赖：在Java项目中通过Maven或Gradle引入对应版本的onnx-proto包，版本需与目标ONNX模型的版本匹配（如模型为ONNX 1.13则使用1.13.0版本的依赖）
• 读取解析模型：将本地ONNX模型文件读取为字节数组，解析成ONNX官方定义的ModelProto对象
• 修改模型名称：遍历ModelProto的metadata_props字段，找到键为name的条目，修改其值为原名称追加版本后缀；若未找到该条目，则直接添加新的name键值对，同时可兜底修改graph.name字段（部分平台可能依赖该字段）
• 重新序列化保存：将修改后的ModelProto对象序列化为字节数组，写入新的ONNX文件

代码示例 #

import org.onnx.proto.OnnxProto3.ModelProto;
import org.onnx.proto.OnnxProto3.StringStringEntryProto;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;

public class OnnxModelRenamer {
    public static void renameOnnxModel(String inputPath, String outputPath, String versionSuffix) throws IOException {
        // 读取原模型
        FileInputStream fis = new FileInputStream(inputPath);
        ModelProto model = ModelProto.parseFrom(fis);
        fis.close();

        ModelProto.Builder modelBuilder = model.toBuilder();
        modelBuilder.clearMetadataProps();
        boolean nameUpdated = false;

        // 遍历元数据修改名称
        for (StringStringEntryProto entry : model.getMetadataPropsList()) {
            if (entry.getKey().equals("name")) {
                String newName = entry.getValue() + "_v" + versionSuffix;
                modelBuilder.addMetadataProps(StringStringEntryProto.newBuilder()
                        .setKey("name")
                        .setValue(newName)
                        .build());
                nameUpdated = true;
            } else {
                modelBuilder.addMetadataProps(entry);
            }
        }

        // 无name字段时添加并修改图名称
        if (!nameUpdated) {
            String baseName = model.getGraph().getName();
            String newName = baseName + "_v" + versionSuffix;
            modelBuilder.addMetadataProps(StringStringEntryProto.newBuilder()
                    .setKey("name")
                    .setValue(newName)
                    .build());
            modelBuilder.getGraphBuilder().setName(newName);
        }

        // 保存修改后的模型
        ModelProto updatedModel = modelBuilder.build();
        FileOutputStream fos = new FileOutputStream(outputPath);
        updatedModel.writeTo(fos);
        fos.close();
    }

    public static void main(String[] args) throws IOException {
        renameOnnxModel("original_model.onnx", "renamed_model_v1.onnx", "1");
    }
}

注意事项 #

• 修改完成后需验证模型可用性，可通过ONNX Runtime加载测试，确保模型推理功能不受影响
• 若平台对模型的其他元数据字段有校验，需确保修改过程中不破坏原有字段的完整性

内容的提问来源于stack exchange，提问作者Waldek Buś