如何用PlotNeuralNet可视化自定义PyTorch CNN架构Net20？

我正在撰写论文，需要展示分析所用的PyTorch CNN架构（Net20）的可视化效果。找到PlotNeuralNet工具，它能生成LaTeX代码用于绘制神经网络，适合报告和演示，但不清楚如何准确定义自己的特定架构。

PlotNeuralNet示例代码 #

import sys
sys.path.append('../')
from pycore.tikzeng import *

# define your arch

arch = \
[
    to_head( '..' ),
    to_cor(),
    to_begin(),
    to_Conv("conv1", 512, 64, offset="(0,0,0)", to="(0,0,0)", height=64, depth=64, width=2 ),
    to_Pool("pool1", offset="(0,0,0)", to="(conv1-east)"),
    to_Conv("conv2", 128, 64, offset="(1,0,0)", to="(pool1-east)", height=32, depth=32, width=2 ),
    to_connection( "pool1", "conv2"),
    to_Pool("pool2", offset="(0,0,0)", to="(conv2-east)", height=28, depth=28, width=1),
    to_SoftMax("soft1", 10 ,"(3,0,0)", "(pool1-east)", caption="SOFT"  ),
    to_connection("pool2", "soft1"),  
    to_Sum("sum1", offset="(1.5,0,0)", to="(soft1-east)", radius=2.5, opacity=0.6),
    to_connection("soft1", "sum1"),
    to_end()
]

def main():
    namefile = str(sys.argv[0]).split('.')[0]
    to_generate(arch, namefile + '.tex' )

if __name__ == '__main__':
    main()

待可视化的Net20模型代码 #

class Net20(nn.Module): 
    """ CNN for 20-day Image 
    This particular model should have:
    - 3 blocks
    - 64 layers in first block, multiply by 2 each subsequent block
    - filter size (5,3)
    - vertical stride = 3 (but only in first layer)
    - vertical dilation = 2 (but only in first layer)
    - Leaky Relu activation function
    - max pooling (2,1) at the end of each block
    """ 
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Sequential(
            Conv2dSame(1, 64, kernel_size=(5,3), stride=(3,1), dilation=(2,1)),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(negative_slope=0.01, inplace=True),
            nn.MaxPool2d((2, 1), ceil_mode=True)
        )
        self.layer2 = nn.Sequential(
            Conv2dSame(64, 128, kernel_size=(5,3)),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(negative_slope=0.01, inplace=True),
            nn.MaxPool2d((2, 1), ceil_mode=True)
        )
        self.layer3 = nn.Sequential(
            Conv2dSame(128, 256, kernel_size=(5,3)),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(negative_slope=0.01, inplace=True),
            nn.MaxPool2d((2, 1), ceil_mode=True)
        )
        self.fc1 = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(46080, 1), 
        )

    def forward(self, x):
        x = x.reshape(-1,1,64,60)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = x.reshape(-1,46080)
        x = self.fc1(x)
        return x

针对Net20的PlotNeuralNet配置代码 #

下面是适配Net20架构的配置代码，标注了每个模块的对应关系：

import sys
sys.path.append('../')
from pycore.tikzeng import *

# 定义Net20架构
arch = \
[
    to_head('..'),
    to_cor(),
    to_begin(),
    # 输入层：对应reshape后的(1,64,60)张量
    to_input("input", height=64, depth=60, width=1, caption="Input (1×64×60)"),
    
    # 第一个卷积块
    to_Conv("conv1", 60, 64, offset="(1,0,0)", to="(input-east)", 
            height=10, depth=60, width=2,
            caption="Conv2d(1→64, (5,3), stride=(3,1), dilation=(2,1))"),
    to_BatchNorm("bn1", 60, 64, offset="(0.5,0,0)", to="(conv1-east)", height=10, depth=60, width=2),
    to_LeakyRelu("relu1", 60, 64, offset="(0.5,0,0)", to="(bn1-east)", height=10, depth=60, width=2),
    to_Pool("pool1", offset="(0.5,0,0)", to="(relu1-east)", height=5, depth=60, width=2, 
            caption="MaxPool2d((2,1))"),
    to_connection("input", "conv1"),
    to_connection("conv1", "bn1"),
    to_connection("bn1", "relu1"),
    to_connection("relu1", "pool1"),
    
    # 第二个卷积块
    to_Conv("conv2", 60, 128, offset="(1,0,0)", to="(pool1-east)", 
            height=5, depth=60, width=4, caption="Conv2d(64→128, (5,3))"),
    to_BatchNorm("bn2", 60, 128, offset="(0.5,0,0)", to="(conv2-east)", height=5, depth=60, width=4),
    to_LeakyRelu("relu2", 60, 128, offset="(0.5,0,0)", to="(bn2-east)", height=5, depth=60, width=4),
    to_Pool("pool2", offset="(0.5,0,0)", to="(relu2-east)", height=3, depth=60, width=4, 
            caption="MaxPool2d((2,1))"),
    to_connection("pool1", "conv2"),
    to_connection("conv2", "bn2"),
    to_connection("bn2", "relu2"),
    to_connection("relu2", "pool2"),
    
    # 第三个卷积块
    to_Conv("conv3", 60, 256, offset="(1,0,0)", to="(pool2-east)", 
            height=3, depth=60, width=8, caption="Conv2d(128→256, (5,3))"),
    to_BatchNorm("bn3", 60, 256, offset="(0.5,0,0)", to="(conv3-east)", height=3, depth=60, width=8),
    to_LeakyRelu("relu3", 60, 256, offset="(0.5,0,0)", to="(bn3-east)", height=3, depth=60, width=8),
    to_Pool("pool3", offset="(0.5,0,0)", to="(relu3-east)", height=2, depth=60, width=8, 
            caption="MaxPool2d((2,1))"),
    to_connection("pool2", "conv3"),
    to_connection("conv3", "bn3"),
    to_connection("bn3", "relu3"),
    to_connection("relu3", "pool3"),
    
    # 全连接层
    to_Flatten("flatten", offset="(1,0,0)", to="(pool3-east)", caption="Flatten → 46080"),
    to_Dropout("dropout", 1, offset="(0.5,0,0)", to="(flatten-east)", width=1, height=1, depth=1, caption="Dropout(p=0.5)"),
    to_FullyConnected("fc1", 1, offset="(0.5,0,0)", to="(dropout-east)", width=1, height=1, depth=1, caption="Linear(46080→1)"),
    to_connection("pool3", "flatten"),
    to_connection("flatten", "dropout"),
    to_connection("dropout", "fc1"),
    
    to_end()
]

def main():
    namefile = str(sys.argv[0]).split('.')[0]
    to_generate(arch, namefile + '.tex' )

if __name__ == '__main__':
    main()

配置说明 #

• 输入层：对应模型中x.reshape(-1,1,64,60)的张量，设置height=64（垂直维度）、depth=60（水平维度）、width=1（通道数）。
• 卷积层：to_Conv的参数中，第一个数字是特征图的水平维度（保持60不变），第二个是通道数；height根据卷积和池化后的尺寸估算，width按通道数比例设置（64对应2，128对应4，256对应8）。
• 批量归一化、LeakyReLU：使用to_BatchNorm和to_LeakyRelu模块，尺寸与前一层卷积输出匹配。
• 池化层：to_Pool对应MaxPool2d((2,1))，height减半，depth不变。
• 全连接部分：用to_Flatten表示展平操作，to_Dropout和to_FullyConnected对应模型中的全连接层。

替代可视化方法 #

如果PlotNeuralNet配置太繁琐，还可以用以下工具：

• torchinfo：直接打印模型的详细结构和参数，包括输入输出尺寸，适合快速查看架构。示例代码：

  from torchinfo import summary
  model = Net20()
  summary(model, input_size=(1, 1, 64, 60))  # 对应输入张量形状
  

• Netron：可视化工具，支持PyTorch模型，可直接加载.pt或.pth模型文件，生成交互式架构图，适合直观展示。

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