PyTorch：掌握CUDA内存使用情况

PyTorch 是一个广泛使用的深度学习框架，它支持使用 CUDA（Compute Unified Device Architecture）进行GPU加速计算。在使用 PyTorch 进行深度学习时，我们经常需要知道 CUDA 已经占用了多少内存。本文将介绍如何输出 CUDA 已经占用的内存量。
首先，我们需要导入 torch 库，并确保已经安装了 torch 和 torchvision。

import torch

接下来，我们可以使用 torch.cuda.memory_allocated() 函数来获取当前 CUDA 已经分配的内存量。该函数返回一个整数，表示当前被 PyTorch 的 CUDA 后端使用的字节数。以下是一个示例代码片段：

memory = torch.cuda.memory_allocated()
print(f"CUDA已经占用的内存量为：{memory / (1024 * 1024)} MB")

在这个示例中，我们将 torch.cuda.memory_allocated() 返回的内存量除以 1024 * 1024，将其转换为 MB。如果您希望将内存量转换为其他单位（如 GB），可以相应地调整除数。
如果您希望输出所有可用的 CUDA 设备的内存量信息，可以使用 torch.cuda.device_properties() 函数。该函数返回一个包含有关所有可用的 CUDA 设备的详细信息的元组。以下是一个示例代码片段：

properties = torch.cuda.device_properties()
for i, device in enumerate(properties):
print(f"Device {i}:")
print(f"  Name: {device.name}")
print(f"  Total memory: {device.total_memory / (1024 * 1024)} MB")
print(f"  Memory allocated: {torch.cuda.memory_allocated(device.id) / (1024 * 1024)} MB")

在这个示例中，我们使用 torch.cuda.device_properties() 返回的元组，遍历所有可用的 CUDA 设备，并输出每个设备的名称、总内存量和当前已分配的内存量。注意，我们在 torch.cuda.memory_allocated() 中传递了设备的 ID，以便获取特定设备的内存分配情况。