PyTorch 查看内存占用
在深度学习中，模型的训练和推理过程需要大量的内存。了解如何查看内存占用是优化模型性能和调试内存泄漏的重要一步。本文将介绍如何使用 PyTorch 检查内存占用，并提供代码示例。
查看 GPU 内存占用
PyTorch 提供了 torch.cuda.memory_allocated() 和 torch.cuda.max_memory_allocated() 函数，可以用于查看当前已分配的 GPU 内存和最大已分配的 GPU 内存。
import torch
# 创建一个张量并将其存储在 GPU 上
tensor = torch.zeros((1000, 1000)).to('cuda')
# 查看当前已分配的 GPU 内存
allocated_memory = torch.cuda.memory_allocated()
print(f"Allocated GPU memory: {allocated_memory / 1024**3} GB")
# 执行一些操作，使内存占用增加
result = tensor.matmul(tensor)
# 查看最大已分配的 GPU 内存
max_allocated_memory = torch.cuda.max_memory_allocated()
print(f"Max allocated GPU memory: {max_allocated_memory / 1024**3} GB")
在上面的代码中，我们首先创建一个大小为 1000x1000 的零张量，并将其存储在 GPU 上。然后，我们使用 torch.cuda.memory_allocated() 函数查看当前已分配的 GPU 内存，并将结果除以 1024**3 将其转换为 GB 单位进行显示。接下来，我们执行一些操作，使内存占用增加，然后再次使用 torch.cuda.max_memory_allocated() 函数查看最大已分配的 GPU 内存。
查看模型和张量占用的内存
除了查看 GPU 内存占用外，我们还可以查看模型和张量本身占用的内存。PyTorch 提供了 torch.Tensor.numel() 函数可以用于获取张量元素的总数，而模型的内存占用可以通过计算其参数的总数来估计。以下是一个示例：
import torch
import torchvision.models as models
# 创建一个预训练的模型
model = models.resnet50(pretrained=True)
# 统计模型参数的总数
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total_params / 1024**2} MB")
# 创建一个大小为 1000x1000 的零张量
tensor = torch.zeros((1000, 1000))
# 统计张量的元素总数
total_elements = tensor.numel()
print(f"Total elements: {total_elements / 1024**2} MB")
在上面的代码中，我们首先使用 torchvision 库中的 resnet50 函数创建一个预训练的 ResNet-50 模型。然后，我们使用一个简单的生成器表达式，计算了模型参数的总数，并将结果除以 1024**2 将其转换为 MB 单位进行显示。接下来，我们创建了一个大小为 1000x1000 的零张量，并使用 torch.Tensor.numel() 函数计算了张量元素的总数，并将结果除以 1024**2 将其转换为 MB 单位进行显示。
内存占用可视化
除了查看内存占用，有时候我们也希望将内存占用可视化，以便更好地理解和分析。下面是一个使用 mermaid 语法绘制的关系图，用于表示 GPU 内存、模型内存和张量内存的关系：
erDiagram
GPU --o| Allocated
GPU --o| Max Allocated
Allocated --o| Model
Allocated --o| Tensor
上述关系图中，GPU 表示 GPU 内存，Allocated 表示已分配的 GPU 内存，Max Allocated 表示最大已分配的 GPU 内存，Model 表示模型内存，Tensor 表示张量内存。箭头表示了它们之间的关系。
另外，我们还可以使用 mermaid 语法绘制一个旅行图，以展示模型和张量如何在内存中占用空间