对pytorch里面的网络模型结构可视化

这里对构建好的模型进行可视化，对每一层的网络及其参数进行输出

net是我自己构建的模型

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d (32, 64, 5,padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10),
        )
    
    def forward(self,x):
        x = self.model(x)
        return x      

net = Net()

需要用到torchviz这个库里面的make_dot，自己构造一个测试的数据，一个batch为64的3通道32*32的tensor，这个尺寸需要满足自己设计的模型的输入数据维度和大小

from torchviz import make_dot
x = torch.randn(64, 3, 32, 32).requires_grad_(True)  # 定义一个网络的输入值
y = net(x) 

pic_1 = make_dot(y, params=dict(list(net.named_parameters()) + [('x', x)]))
pic_1.format = "png"  #定义pic为模型可视化后的输出,这里输出为png格式
 # 指定文件生成的文件夹pic里面
pic_1.directory = "pic"
 # 生成文件
pic_1.view()

最后输出的模型图片如下