pytorch查看网络参数显存占用量等操作

下面是针对pytorch查看网络参数显存占用量等操作的完整攻略。

1. 查看网络参数总量

为了查看神经网络的参数总量，我们可以使用 torchsummary 库中的 summary 函数。该函数可以打印出我们定义的模型结构及其参数量等相关信息。

首先，我们需要在命令行中使用 pip 安装 torchsummary 库：

pip install torchsummary

然后，我们可以在pytorch代码中导入该库，使用以下代码查看神经网络的参数总量：

import torch
from torchsummary import summary

# 定义神经网络模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 6, 5)
        self.pool = torch.nn.MaxPool2d(2, 2)
        self.conv2 = torch.nn.Conv2d(6, 16, 5)
        self.fc1 = torch.nn.Linear(16 * 5 * 5, 120)
        self.fc2 = torch.nn.Linear(120, 84)
        self.fc3 = torch.nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 创建模型实例并查看参数总量
model = Net()
summary(model, (3, 32, 32))

其中，我们首先定义了一个神经网络模型 Net，然后在 main 函数中创建了一个 Net 的实例 model。最后，使用 summary 函数查看神经网络模型的参数总量。

在这个示例中，我们的神经网络有 61,706 个可训练的参数。

2. 查看网络的显存占用量

在训练神经网络时，我们需要时刻关注神经网络的显存占用量，以免显存溢出或训练速度过慢。下面我们介绍两种查看神经网络显存占用量的方法。

方法一：pytorch自带工具

pytorch自带诊断工具 torch.cuda.memory_summary 可以用来查看系统显存状态和分配情况，使用方法如下：

import torch

# 定义张量
a = torch.randn(1024, 1024).cuda()

# 打印显存分配情况
print(torch.cuda.memory_summary())

运行上述代码后，我们可以得到类似下面的输出：

|    GPU     |    累计使用 |
|:---------:|:-------------:|
|  Quadro P4000  |  3554 MB      |

这表示当前显存使用了 3554 MB，其中包括了我们之前创建的 1024*1024 的张量占用的显存。

方法二：torch.autograd.profiler

torch.autograd.profiler 是 pytorch 自带的性能评测工具，可以用于查看神经网络的显存使用情况以及计算时间、函数调用顺序等。

import torch

# 定义神经网络模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 6, 5)
        self.pool = torch.nn.MaxPool2d(2, 2)
        self.conv2 = torch.nn.Conv2d(6, 16, 5)
        self.fc1 = torch.nn.Linear(16 * 5 * 5, 120)
        self.fc2 = torch.nn.Linear(120, 84)
        self.fc3 = torch.nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 创建模型实例
model = Net()

# 创建输入张量，将张量移动到GPU上
inputs = torch.randn(1, 3, 32, 32).cuda()

# 使用Profiler查看网络的显存占用情况
with torch.autograd.profiler.profile(enabled=True, use_cuda=True) as profile:
    model(inputs)

print(profile)

运行上述代码后，我们可以得到类似下面的输出：

-----------------------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  
Name                                     CPU time        CUDA time        Calls           CPU total       CUDA total      Input size     
-----------------------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  
_convolution                            0.200us          1.757ms          2               0.000us          3.513ms          [1, 6, 28, 28]  
_max_pool2d                             0.025us          0.120us          2               0.050us          0.240us          [1, 6, 28, 28]  
_convolution                            0.032us          0.272ms          2               0.000us          0.544ms          [1, 16, 10, 10] 
_max_pool2d                             0.004us          0.052us          2               0.008us          0.104us          [1, 16, 10, 10] 
_view                                   0.065us          0.120us          2               0.130us          0.240us          [1, 400]        
_linear                                 0.064us          0.073ms          3               0.000us          0.219ms          [1, 120]        
_relu                                   0.025us          0.002ms          3               0.000us          0.005ms          [1, 120]        
_linear                                 0.020us          0.022ms          3               0.000us          0.067ms          [1, 84]         
_relu                                   0.011us          0.001ms          3               0.000us          0.004ms          [1, 84]         
_linear                                 0.012us          0.002ms          1               0.012us          0.002ms          [1, 10]         
-----------------------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  
Self CPU time total: 0.518ms
Self CUDA time total: 4.059ms

其中，CUDA time 就是每个操作在 GPU 上使用的显存总量，可以通过求和得到模型整体的显存占用情况。在这个示例中，我们的模型总共占用了 4.059 MB 的显存。