PyTorch内存与参数管理笔记

获取PyTorch模型的总参数量：

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

total_params = sum(p.numel() for p in model.parameters())
print(f"Total number of parameters: {total_params}")

查看每层的参数

获取每个层的参数量：

for name, parameter in model.named_parameters():
    print(f"{name}: {parameter.numel()}")

查看当前GPU内存使用

# 当前分配的总内存
current_memory = torch.cuda.memory_allocated()
print(f"Current memory allocated: {current_memory} bytes")

# 峰值内存
peak_memory = torch.cuda.max_memory_allocated()
print(f"Peak memory allocated: {peak_memory} bytes")

优化和释放内存

删除不再需要的变量

x = torch.randn(1000, 1000, device="cuda")
y = torch.randn(1000, 1000, device="cuda")
z = x + y

del x, y
torch.cuda.empty_cache()  # 清空未使用的缓存

使用torch.no_grad()

在进行推理或任何不需要计算梯度的前向传播时，使用torch.no_grad()来减少内存使用：

with torch.no_grad():
    output = model(input)

原地操作

尽可能使用原地操作，减少内存分配：

x = torch.randn(1000, 1000, device="cuda")
x.add_(1)  # 原地加1

清空CUDA缓存

在确定某些内存不再被需要后，可以清空CUDA缓存：

torch.cuda.empty_cache()