Pytorch报错：AssertionError: Only output_size=[1, 1] is supported Vgg19网络结构可视化

在做vgg16、vgg19模型可视化时遇到的错误。
错误原因：AdaptiveAvgPool2d(output_size=(7,7))
同样类型错误：
https://discuss.pytorch.org/t/when-using-torch-onnx-export-errors-occurs-assertionerror-only-output-size-1-1-is-supported/38528
错误主要是由输出大小不是[1,1]的AdaptiveAvgPool2d引起的。

修改方法：将AdaptiveAvgPool2d改为AvgPool2d即可，不影响对可视化结果的理解。

import os
import torch
from torch import nn
import tensorwatch as tw
from torch import nn
from torchvision import models

os.environ["PATH"] += r'C:\Program Files (x86)\Graphviz2.38\bin' # 安装graphviz的路径
model = models.vgg19(pretrained=True)


class VGG19(nn.Module):
    def __init__(self):
        super(VGG19, self).__init__()
        self.features = model.features
        self.avg_pool = nn.AvgPool2d(1,1) # 对于224*224的输入，features输出大小为[1,512,7,7],这里主要为了不影响输出大小
        self.classifier = model.classifier # 输入通道数为512*7*7

    def forward(self, x):
        x = self.features(x)
        x = self.avg_pool(x)
        x = x.view(1, -1)
        x = self.classifier(x)
        return x


vgg19 = VGG19()
dummy_input = torch.rand(1, 3, 224, 224)
print(vgg19(dummy_input).shape)
tw.draw_model(vgg19, dummy_input).save('vgg19')

得到的VGG19可视化结果为：