PyTorch 中的 CUDA 内存

PyTorch 是一个深度学习框架，提供了对 CUDA 的支持，可以利用 GPU 进行加速计算。由于 GPU 的显存较小，因此在使用 PyTorch 进行深度学习时需要特别关注 GPU 内存的使用情况。本文将介绍 PyTorch 中的 CUDA 内存管理，帮助程序员更好地利用 GPU 资源。

1. PyTorch 中的 CUDA 内存分配

在 PyTorch 中，使用 torch.Tensor 创建的 Tensor 变量默认在 CPU 上分配内存。如果我们要将其移动到 GPU 上，则可以使用 to() 方法：

import torch

x = torch.randn(3, 3)
x_gpu = x.to('cuda')

我们也可以在创建 Tensor 变量时直接指定其存储设备：

x_gpu = torch.randn(3, 3, device='cuda')

在使用 PyTorch 进行深度学习时，常用的方式是先将数据加载到 CPU 上，然后在训练过程中逐步地将 Tensor 变量移动到 GPU 上。例如：

import torch
from torchvision import datasets, transforms

train_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
train_data = datasets.MNIST('data', train=True, download=True, transform=train_transform)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)

# 定义模型、损失函数和优化器

for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to('cuda'), target.to('cuda')
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()

在上述代码中，我们将 MNIST 数据集加载到 CPU 上，然后在每个 Batch 中将数据和标签移动到 GPU 上，以便在 GPU 上进行训练。这种方式可以最大限度地减少 GPU 内存的浪费，并且可以在需要时灵活地移动 Tensor 变量。

2. PyTorch 中的 CUDA 内存释放

在使用 PyTorch 进行深度学习时，我们创建了许多 Tensor 变量，这些变量占用了 GPU 内存。为了节省 GPU 资源，我们需要及时释放无用的 Tensor 变量。

2.1 del 关键字

Python 的 del 关键字可以删除 Tensor 变量以释放其占用的内存。例如：

import torch

x = torch.randn(3, 3)
x_gpu = x.to('cuda')
del x

在上述代码中，我们将 x 转换为 GPU 上的 Tensor 变量 x_gpu 后，使用 del 关键字删除了 x。这样，x 和 x_gpu 占用的内存就可以被释放了。

然而，del 关键字只会删除 Python 变量的引用，而不会显式地释放 GPU 内存。如果我们要显式地释放 GPU 内存，需要使用 PyTorch 的 torch.cuda.empty_cache() 方法。

2.2 torch.cuda.empty_cache()

torch.cuda.empty_cache() 方法可以释放 PyTorch 占用的所有 GPU 内存，包括未被 Python 引用的内存。例如：

import torch

x = torch.randn(3, 3)
x_gpu = x.to('cuda')
del x
torch.cuda.empty_cache()

在上述代码中，我们先将 x 转换为 GPU 上的 Tensor 变量 x_gpu，然后使用 del 关键字删除了 x，最后使用 torch.cuda.empty_cache() 方法释放了 GPU 内存。

3. PyTorch 中的 CUDA 内存检查

在开发深度学习模型时，GPU 内存不足是一个常见的问题。为了保证程序的稳定性和可靠性，我们需要经常检查 GPU 内存的使用情况。

3.1 torch.cuda.memory_allocated()

torch.cuda.memory_allocated() 方法可以返回当前已分配的 GPU 内存大小（单位为字节）。例如：

import torch

x = torch.randn(3, 3)
x_gpu = x.to('cuda')
mem_alloc = torch.cuda.memory_allocated()

在上述代码中，我们先将 x 转换为 GPU 上的 Tensor 变量 x_gpu，然后使用 torch.cuda.memory_allocated() 方法获取已分配的 GPU 内存大小，保存在 mem_alloc 变量中。

3.2 torch.cuda.max_memory_allocated()

torch.cuda.max_memory_allocated() 方法可以返回已分配的 GPU 内存的历史峰值（单位为字节）。例如：

import torch

max_mem_alloc = 0
for i in range(100):
    x = torch.randn(1024, 1024)
    x_gpu = x.to('cuda')
    mem_alloc = torch.cuda.memory_allocated()
    max_mem_alloc = max(max_mem_alloc, mem_alloc)
print(max_mem_alloc)

在上述代码中，我们随机生成 100 个大小为 [1024, 1024] 的 Tensor 变量，并将它们转换为 GPU 上的 Tensor 变量。在每次分配完内存后，我们都调用 torch.cuda.memory_allocated() 方法获取已分配的 GPU 内存大小，并更新 max_mem_alloc 变量。最终，我们可以通过 max_mem_alloc 变量获取已分配的 GPU 内存的历史峰值。

4. 总结

本文介绍了 PyTorch 中的 CUDA 内存分配、CUDA 内存释放和 CUDA 内存检查方法。在开发深度学习模型时，合理地利用 GPU 资源可以大幅提升训练速度和模型性能，并且避免 GPU 内存不足等问题的出现。