PyTorch转ONNX报错：aten::dict算子在opset12中不支持（DeeplabV3）

报错原因 #

没错，就是模型输出为字典形式导致的问题。ONNX opset 12不支持导出返回字典的PyTorch模型，因为aten::dict算子不在该opset的支持范围内。

解决方法 #

你需要对原模型做一层包装，让它直接返回你需要的张量（也就是原模型输出字典里的out字段），而不是返回字典。具体修改如下：

1. 定义包装模型类：

class WrappedModel(torch.nn.Module):
    def __init__(self, original_model):
        super(WrappedModel, self).__init__()
        self.original_model = original_model
    
    def forward(self, x):
        outputs = self.original_model(x)
        # 只返回需要的'out'张量
        return outputs['out']

2. 修改模型加载与导出代码：

import onnx
import torch
import onnxruntime
import numpy as np

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

batch_size = 1
# 加载原模型
original_model = torch.jit.load('../saved_weights/model.pt')
# 包装模型
model = WrappedModel(original_model).to(device)
model.eval()  # 切换到eval模式，避免训练层影响导出结果

x = torch.randn(1, 3, 224, 224, requires_grad=True).to(device)
torch_out = model(x)

torch.onnx.export(model, x, "../saved_weights/model.onnx",
              export_params=True,
              opset_version=12,
              do_constant_folding=True,
              input_names=['input'],
              output_names=['out'],
              dynamic_axes={ 'input': {0 : 'batch_size'},
                            'out': {0 : 'batch_size'}})  # 注意dynamic_axes的键要和output_names对应

补充说明 #

• 导出前必须切换到eval()模式，避免Dropout、BatchNorm等训练专属层的行为干扰导出结果。
• 如果需要保留多个输出张量，可以让包装模型返回元组（比如return outputs['out'], outputs['aux']），同时修改output_names为对应的列表（比如output_names=['out', 'aux']），ONNX支持导出多输出张量的模型。

内容的提问来源于stack exchange，提问作者HuckleberryFinn