下面是针对pytorch查看网络参数显存占用量等操作的完整攻略。
为了查看神经网络的参数总量,我们可以使用 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 个可训练的参数。
在训练神经网络时,我们需要时刻关注神经网络的显存占用量,以免显存溢出或训练速度过慢。下面我们介绍两种查看神经网络显存占用量的方法。
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 的显存。
本文链接:http://task.lmcjl.com/news/16601.html