PyTorch 清空列表中 GPU 上的张量|极客笔记

PyTorch 清空列表中 GPU 上的张量

在使用 PyTorch 进行深度学习模型训练的过程中，经常会创建并操作多个张量。有时候我们需要清空一个列表中的张量，特别是当这些张量在 GPU 上进行计算时，需要手动释放它们的显存以避免内存泄漏问题。本文将详细介绍如何清空一个列表中的 GPU 张量并释放显存。

准备工作

在开始本文之前，首先需要确保你已经安装了 PyTorch 并了解如何在 GPU 上创建和使用张量。

import torch

# 检测是否有可用的 GPU
if torch.cuda.is_available():
    device = torch.device("cuda")
    print("CUDA is available!")
else:
    device = torch.device("cpu")
    print("CUDA is not available!")

创建一个包含 GPU 张量的列表

首先，让我们创建一个包含 GPU 张量的列表。我们使用 torch.randn 函数来生成一些随机的张量，并将它们放置在 GPU 上。

tensor_list = [torch.randn(3, 3).to(device) for _ in range(5)]
print("原始张量列表：", tensor_list)

清空列表中的 GPU 张量

清空列表中的 GPU 张量可以通过 del 关键字进行操作。释放一个张量可以直接调用其 to(device) 方法将其放置在 CPU 上，然后再调用 del 删除该张量。

for tensor in tensor_list:
    tensor.data = tensor.to("cpu")
    del tensor

检验 GPU 显存是否被释放

为了验证 GPU 显存是否被正确释放，我们可以使用 torch.cuda.memory_allocated 函数查看当前已分配的显存大小。在释放 GPU 张量后，该值应该为 0。

if torch.cuda.is_available():
    torch.cuda.empty_cache()  # 清空 GPU 缓存
    print("GPU 已释放显存：", torch.cuda.memory_allocated())

完整代码

import torch

# 检测是否有可用的 GPU
if torch.cuda.is_available():
    device = torch.device("cuda")
    print("CUDA is available!")
else:
    device = torch.device("cpu")
    print("CUDA is not available!")

# 创建一个包含 GPU 张量的列表
tensor_list = [torch.randn(3, 3).to(device) for _ in range(5)]
print("原始张量列表：", tensor_list)

# 清空列表中的 GPU 张量
for tensor in tensor_list:
    tensor.data = tensor.to("cpu")
    del tensor

# 检验 GPU 显存是否被释放
if torch.cuda.is_available():
    torch.cuda.empty_cache()  # 清空 GPU 缓存
    print("GPU 已释放显存：", torch.cuda.memory_allocated())

运行结果

当运行上述代码后，如果你的环境中有可用的 GPU，你应该会看到如下输出：

CUDA is available!
原始张量列表： [tensor([[ 0.2284, -1.4138,  0.4676],
        [-0.7423, -1.8184, -0.9588],
        [-0.3180,  1.0530,  0.9185]], device='cuda:0'), 
        tensor([[ 1.2860,  1.2799,  0.5475],
        [-0.7374,  0.7734, -0.3058],
        [-0.5322,  0.2575, -0.6601]], device='cuda:0'),
        ...]
GPU 已释放显存： 0

这表明我们成功地清空了列表中的 GPU 张量并释放了显存。

通过本文的介绍，你现在应该能够清空一个包含 GPU 张量的列表并释放显存了。记得在需要释放 GPU 显存的时候及时清空不再需要的张量，以避免内存泄漏问题。