怎么预估pytorch的显存占用?

参考来源:

显存占用计算

初衷:

如何让网络接收一张超大分辨图片的图片,如果直接输入会cuda of meomory,只有使用多卡模型并行的方式,就是将网络拆开放入多个GPU,但每个子网络的显存不能超过最大显存,这就需要估计各个子网络的显存占用.

环境:

ubuntu 20.04
python 3.8.19
pytorch 1.10.0 py3.8_cuda11.3_cudnn8.2.0_0

实验:

实验一:论证显存和总参数量的关系

使用nvidia-smi -l 监测显存真实占用
使用torchsummary得到模型参数总量,包括输入图片,模型参数量,中间层特征图大小,前向传播和后向传播量,好像还缺了一个优化器?

"""
使用torchinfo进行显存推断,和torchsummary区别不大
"""
from torchvision.models.resnet import ResNet, Bottleneck
from torchsummary import summary

class SingleGPUResNet50(ResNet):
    def __init__(self, *args, **kwargs):
        super(SingleGPUResNet50, self).__init__(
            Bottleneck, [3, 4, 6, 3], num_classes=1000, *args, **kwargs)
        # 将整个模型移到单个GPU上
        self.to('cuda:0')

# 定义 ResNet-50 模型
model = SingleGPUResNet50()

# 获取模型摘要
summary(model, input_size=(3, 448, 448), batch_size=1, device='cuda') # 384.62M,2296M
# summary(model, input_size=(3, 2240, 2240)) # 28800.64M,12682M
# summary(model, input_size=(3, 1120, 1120)) # 7267.44M,4668M
# y=0.3671x+2088.14,112,170.3M,2150M,2250M 
# 448,1245.95M,2545.53M,2560M

前三行注释含义:前者会预估总参数量,后者为实际显存占用,通过前三对数据进行线性关系推理即y=0.3671x+2088.14,后两行注释为预测和真实的差别.

实验二:pytoch框架占用显存估计
使用一个很简单的网络,但显存占用仍然为1986M

"""
使用简单例子计算显存占用情况
"""
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# 定义一个简单的神经网络
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(40, 20)
        # self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(20, 10)

    def forward(self, x):
        x = self.fc1(x)
        # x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

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

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 创建随机数据模拟输入
data = torch.randn(1, 40)  # 假设我们有一个64个样本的小批量
labels = torch.randint(0, 10, (1,))  # 生成64个样本的随机标签

# 将数据移至GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
data = data.to(device)
labels = labels.to(device)

# 训练模型
for epoch in range(10000):  # 进行10个训练周期
    optimizer.zero_grad()   # 清除旧的梯度
    output = model(data)  # 前向传播
    loss = criterion(output, labels)  # 计算损失
    loss.backward()  # 反向传播
    optimizer.step()  # 更新权重

    # 打印显存占用信息
    print(f"Epoch {epoch+1}, Loss: {loss.item()}")
    # if torch.cuda.is_available():
    #     print(f"Current GPU memory usage: {torch.cuda.memory_allocated(device) / 1024 ** 3:.2f} GB")

# 注意：如果是在CPU上运行，`torch.cuda.memory_allocated` 将不会工作。

# 1986M

结语

这只是很浅显的例子,希望大家一起来讨论.因为我没有在相关搜索中找到具体的实例进行比较.