onnx模型修改：去掉Dropout层

news2024/12/23 17:55:28

文章目录

- 尝试1：强行设置dropout层train mode为False
- 尝试2：找到onnx模型中的dropout, train mode设置为False
- 尝试3：直接删除dropout层，连接其输入输出
- 结语

最近训练模型使用了tinyvit，性能挺强的：

但是导出onnx时，会提示dropout层的train mode被设置为True了。

UserWarning: ONNX export mode is set to TrainingMode.EVAL, but operator 'dropout' is set to train=True. Exporting with train=True.

这个警告如果只是使用onnxruntime去推理的话，可以不用处理，但是如果使用openvino则会在转换模型时失败。因为导出的onnx中出现了Dropout层，一般的推理框架是不支持推理的时候用dropout的。

尝试1：强行设置dropout层train mode为False

for m in torch_model.modules():
        if isinstance(m, torch.nn.Dropout):
            m.training = False

问题依旧

尝试2：找到onnx模型中的dropout, train mode设置为False

做这个尝试的本意是先设置为False, 再用onnx-simplify去优化一把，理论上会把dropout层去掉。


# 遍历模型的所有Dropout节点, 找到所有的training mode节点名称
training_mode_inputs=[]
for node in model.graph.node:
    if node.op_type == 'Dropout':
        # 获取Dropout节点的training_mode输入（假设是最后一个输入）
        training_mode_input = node.input[-1]
        # 检查这个输入是否指向之前找到的值为True的常量节点
        training_mode_inputs.append(training_mode_input)

# 遍历所有初始化器
for initializer in model.graph.initializer:
    # 检查初始化器是否是我们要找的training_mode输入
    if initializer.name in training_mode_inputs:
        # 假设这个初始化器是一个布尔值，我们将其修改为False
        # 注意：ONNX中的布尔值是以int64类型存储的，0表示False，1表示True
        # initializer.data_type = onnx.TensorProto.INT64
        initializer.int64_data[:] = [0]  # 修改为False
from onnx import helper
new_initializers = []

for initializer in model.graph.initializer:
    if initializer.name in training_mode_inputs:
        # 创建一个新的TensorProto对象，值为False
        new_initializer = helper.make_tensor(
            name=initializer.name,  # 保持原来的名称
            data_type=onnx.TensorProto.BOOL,
            dims=initializer.dims,  # 保持原来的维度
            vals=[0]  # 设置值为False（在ONNX中用0表示）
        )
        new_initializers.append(new_initializer)
    else:
        new_initializers.append(initializer)

# 替换原来的初始化器列表
# Clear existing initializers
model.graph.ClearField('initializer')
# Add the new initializers
model.graph.initializer.extend(new_initializers)

理想很丰满，现实很骨感···并没有发生什么变化

尝试3：直接删除dropout层，连接其输入输出

dropout层在推理的时候也没什么用，直接删除，然后连接上原dropout的输入输出层就好了

import onnx
from onnx import helper

# 加载模型
onnx_model = onnx.load(model_path)
graph = onnx_model.graph

# 找到 Dropout 层
nodes_to_remove = [node for node in graph.node if node.op_type == 'Dropout']

# 删除 Dropout 层并重新连接
for node in nodes_to_remove:
    input_name = node.input[0]
    output_name = node.output[0]
    
    # 找到所有使用 Dropout 输出作为输入的节点
    for next_node in graph.node:
        for i, input_name in enumerate(next_node.input):
            if input_name == node.output[0]:
                next_node.input[i] = node.input[0]
    
    # 从图中移除 Dropout 节点
    graph.node.remove(node)

# 保存修改后的模型
# check if the model is valid
onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, 'tinyvit_11m_sim_replace.onnx')

成功了，模型的dropout层都被删除了。