pytorch绘制卷积神经网络框架图

在PyTorch中，可以使用torchviz库来绘制卷积神经网络的框架图。下面是一个简单的示例代码，展示了如何使用torchviz来可视化卷积神经网络的结构：

import torch
import torch.nn as nn
from torchviz import make_dot

# 定义一个简单的卷积神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 50, 5)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.max_pool2d(x, 2, 2)
        x = nn.functional.relu(self.conv2(x))
        x = nn.functional.max_pool2d(x, 2, 2)
        x = x.view(-1, 4*4*50)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 创建一个随机输入
input = torch.randn(1, 1, 28, 28)

# 创建网络实例
net = Net()

# 将网络实例和输入传递给make_dot函数
output = net(input)
dot = make_dot(output, params=dict(net.named_parameters()))

# 保存框架图为PDF文件
dot.format = 'pdf'
dot.render(filename='convnet')

运行上述代码后，将会生成一个名为convnet.pdf的PDF文件，其中包含了卷积神经网络的框架图。