使用torchsummary进行模型评估测试报错
问题:使用from torchvision.models.vgg import VGG引用VGG-16。作为FCN前置网络,在使用torchsummary进行模型评估测试报错
//作为FCN前置网络,在使用torchsummary进行模型评估测试报错
Traceback (most recent call last):
File "F:/ggsddu/graduationProject/tongue_segmentation/FCN.py", line 163, in <module>
summary(model=fcn_model, input_size=(3, 224, 224), batch_size=-1, device="cpu")
File "D:\anaconda\envs\pytorch1.2\lib\site-packages\torchsummary\torchsummary.py", line 72, in summary
model(*x)
File "D:\anaconda\envs\pytorch1.2\lib\site-packages\torch\nn\modules\module.py", line 493, in __call__
result = self.forward(*input, **kwargs)
File "F:/ggsddu/graduationProject/tongue_segmentation/FCN.py", line 28, in forward
output = self.pretrained_net(x)
File "D:\anaconda\envs\pytorch1.2\lib\site-packages\torch\nn\modules\module.py", line 495, in __call__
hook_result = hook(self, input, result)
File "D:\anaconda\envs\pytorch1.2\lib\site-packages\torchsummary\torchsummary.py", line 26, in hook
summary[m_key]["output_shape"] = list(output.size())
AttributeError: 'dict' object has no attribute 'size'
进程已结束,退出代码为 1
解决方案:
打开torchsummary.py,找到23行if isinstance(output, (list, tuple)):
在之前加入以下代码将dict中value转化为list类型
if isinstance(output, dict):
output = list(output.values()
再运行成功输出模型