ONNX模型动态维度修改的C++实现方法

概述

在使用ONNX模型时，经常会遇到模型输入或输出包含动态维度的情况。这些动态维度通常表示为维度参数(如"w"、"h"等)，而不是具体的数值。本文介绍如何在C++环境下修改ONNX模型中的动态维度，将其转换为固定数值。

问题背景

ONNX模型中的张量形状可以包含两种类型的维度：

1. 固定维度：直接指定维度大小的整数值
2. 动态维度：使用字符串参数表示，如"batch_size"、"height"、"width"等

在实际部署中，我们通常需要将这些动态维度转换为固定值，以便模型能够处理特定尺寸的输入数据。

解决方案

核心思路

通过ONNX的C++接口，我们可以：

1. 遍历模型的所有输入
2. 检查每个输入的维度信息
3. 将动态维度参数替换为具体的数值

实现步骤

1. 准备维度映射表：创建一个映射表，将维度参数名称映射到具体的数值

   std::map<std::string, int> valSubstitute;
   valSubstitute["w"] = 640;
   valSubstitute["h"] = 480;
   

2. 遍历模型输入：使用ONNX的GraphProto接口遍历所有输入

   for (int i = 0; i < graph_proto->input_size(); i++) {
       auto node = graph_proto->input(i);
       // 处理每个输入的维度
   }
   

3. 处理每个维度：对于每个维度，检查是否为动态维度，并进行替换

   for (int j = 0; j < node.type().tensor_type().shape().dim_size(); j++) {
       auto dim = node.type().tensor_type().shape().dim(j);
       if (dim.has_dim_param()) {
           // 如果是动态维度，检查是否需要替换
           if (valSubstitute.find(dim.dim_param()) != valSubstitute.end()) {
               dim.set_dim_value(valSubstitute[dim.dim_param()]);
               dim.clear_dim_param();
           }
       }
   }
   

注意事项

1. 直接修改问题：直接修改GraphProto中的维度可能不会生效，因为ONNX使用protobuf消息，需要正确处理消息字段

2. 正确修改方法：应该先获取可修改的维度引用，然后进行修改

   auto* mutable_dim = node.mutable_type()->mutable_tensor_type()->mutable_shape()->mutable_dim(j);
   if (mutable_dim->has_dim_param()) {
       // 进行修改操作
   }
   

3. 模型克隆：更安全的方法是先克隆整个模型，然后在克隆体上进行修改

完整示例

// 克隆原始模型
onnx::ModelProto model_copy;
model_copy.CopyFrom(*original_model);

// 准备维度映射
std::map<std::string, int> dim_mapping = {{"w", 640}, {"h", 480}};

// 获取可修改的graph
auto* graph = model_copy.mutable_graph();

// 遍历所有输入
for (int i = 0; i < graph->input_size(); ++i) {
    auto* input = graph->mutable_input(i);
    if (input->has_type() && input->type().has_tensor_type()) {
        auto* shape = input->mutable_type()->mutable_tensor_type()->mutable_shape();
        for (int j = 0; j < shape->dim_size(); ++j) {
            auto* dim = shape->mutable_dim(j);
            if (dim->has_dim_param()) {
                const auto& param_name = dim->dim_param();
                if (dim_mapping.count(param_name)) {
                    dim->set_dim_value(dim_mapping.at(param_name));
                    dim->clear_dim_param();
                }
            }
        }
    }
}

总结

在C++中修改ONNX模型的动态维度需要注意protobuf消息的处理方式。通过正确使用mutable_前缀的方法获取可修改的消息引用，可以确保维度修改操作生效。对于生产环境，建议采用模型克隆后再修改的方式，以避免意外修改原始模型。这种方法适用于需要将动态尺寸模型转换为固定尺寸模型的场景，便于后续的优化和部署。