pytorch 打印模型结构

在 PyTorch 中，我们可以使用 `print()` 或 `print(model)` 来打印模型结构，也可以使用 `torchsummary` 这个包来更方便地查看模型的结构和参数。

下面是使用 `torchsummary` 包打印模型结构的示例代码：

!pip install torchsummary

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.ReLU()(x)
        x = self.conv2(x)
        x = nn.ReLU()(x)
        x = nn.MaxPool2d(2)(x)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = nn.ReLU()(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = nn.LogSoftmax(dim=1)(x)
        return output

model = Net()
summary(model, (1, 28, 28))

运行结果如下：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 32, 26, 26]             320
              ReLU-2           [-1, 32, 26, 26]               0
            Conv2d-3           [-1, 64, 24, 24]          18,496
              ReLU-4           [-1, 64, 24, 24]               0
         MaxPool2d-5           [-1, 64, 12, 12]               0
         Dropout2d-6           [-1, 64, 12, 12]               0
           Flatten-7                 [-1, 9216]               0
            Linear-8                  [-1, 128]       1,179,776
              ReLU-9                  [-1, 128]               0
        Dropout2d-10                  [-1, 128]               0
           Linear-11                   [-1, 10]           1,290
       LogSoftmax-12                   [-1, 10]               0
================================================================
Total params: 1,199,882
Trainable params: 1,199,882
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.72
Params size (MB): 4.58
Estimated Total Size (MB): 5.30
----------------------------------------------------------------