调用torchsummary查看DeepLabV3预训练模型时报OrderedDict无size属性问题问询

PyTorch官方提供的图像分割类模型（包括你用到的deeplabv3_resnet50）返回值为collections.OrderedDict类型，默认包含out（主分割输出）、aux（辅助分支输出）两个键对应的张量。而你调用的summary函数默认会尝试对模型输出调用.size()方法获取张量维度，OrderedDict本身无该属性，因此触发报错。

此外你原有代码存在笔误：模型调用路径多写了一级segmentation，正确路径为torchvision.models.segmentation.deeplabv3_resnet50，多余层级会额外触发模块不存在的报错。

方案1：封装模型仅返回主输出张量 #

对原始模型做轻量封装，调用后直接返回主输出张量，即可适配torchsummary的检测逻辑，示例代码如下：

import torch
import torchvision
from torchsummary import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载原始预训练模型
raw_model = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=True).eval().to(device)

# 模型封装
class Deeplabv3Wrapper(torch.nn.Module):
    def __init__(self, raw_model):
        super().__init__()
        self.raw_model = raw_model
    def forward(self, x):
        # 仅返回主分割输出张量
        return self.raw_model(x)['out']

model = Deeplabv3Wrapper(raw_model).to(device)
summary(model, (3,256,256))

方案2：使用支持OrderedDict输出的summary工具 #

不需要修改模型结构，换用原生支持解析分割模型输出的torchinfo库即可直接调用：

1. 安装依赖：
   pip install torchinfo
2. 调用示例：

import torch
import torchvision
from torchinfo import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=True).eval().to(device)
# input_size需指定batch维度，此处batch size设为1
summary(model, input_size=(1,3,256,256))

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