pytorch自定义网络模型结构图如何展示

PyTorch自定义网络模型结构图可以通过使用PyTorch内置的`torchsummary`模块来生成。这个模块可以帮助我们快速地展示模型的参数数量、每一层的输出形状等重要信息。

以下是一个简单的示例，展示了如何使用`torchsummary`模块来生成自定义网络模型的结构图：

import torch
import torch.nn as nn
from torchsummary import summary

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(64 * 8 * 8, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = nn.functional.relu(x)
        x = self.pool(x)
        x = x.view(-1, 64 * 8 * 8)
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.fc2(x)
        return x

model = MyModel()
summary(model, input_size=(3, 32, 32))

运行以上代码，就可以得到如下的输出：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 32, 32, 32]             896
         MaxPool2d-2           [-1, 32, 16, 16]               0
            Conv2d-3           [-1, 64, 16, 16]          18,496
         MaxPool2d-4             [-1, 64, 8, 8]               0
            Linear-5                  [-1, 128]         524,416
            Linear-6                   [-1, 10]           1,290
================================================================
Total params: 545,098
Trainable params: 545,098
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.01
Forward/backward pass size (MB): 0.75
Params size (MB): 2.08
Estimated Total Size (MB): 2.85
----------------------------------------------------------------

可以看到，`summary`函数生成了一个包含每一层输出形状、参数数量等信息的表格，以及估计的模型大小。