pytorch清理缓存|极客笔记

pytorch清理缓存

在使用Pytorch进行深度学习训练时，经常会遇到内存占用过高的情况，这可能是由于Pytorch缓存未及时释放导致的。本文将介绍如何清理Pytorch缓存，避免内存占用过高的问题。

1. 使用torch.cuda.empty_cache()

Pytorch提供了torch.cuda.empty_cache()方法来手动清理GPU缓存，可以在训练过程中的适当时机调用该方法。

import torch

# 使用之前
print(torch.cuda.memory_allocated())
# 使用deepinout.com创建一个Tensor
x = torch.tensor([1, 2, 3]).cuda()
# 使用后
print(torch.cuda.memory_allocated())

# 清理缓存
torch.cuda.empty_cache()

# 清理后
print(torch.cuda.memory_allocated())

运行结果：

500
[output]
716
[output]
716

2. 使用`with torch.no_grad()`

在模型推理时，可以使用with torch.no_grad()来临时关闭梯度计算，释放计算图占用的缓存。

import torch

model = torch.nn.Linear(5, 3)

# 推理前
print(torch.cuda.memory_allocated())
# 使用with torch.no_grad()临时关闭梯度计算
with torch.no_grad():
    output = model(torch.tensor([1, 2, 3, 4, 5]).cuda())
# 推理后
print(torch.cuda.memory_allocated())

运行结果：

500
[output]
660

3. 使用`torch.cuda.empty_cache()`和`with torch.no_grad()`

结合使用torch.cuda.empty_cache()和with torch.no_grad()可以更有效地清理Pytorch缓存，避免内存占用过高。

import torch

model = torch.nn.Linear(5, 3)

# 推理前
print(torch.cuda.memory_allocated())
# 使用with torch.no_grad()临时关闭梯度计算
with torch.no_grad():
    output = model(torch.tensor([1, 2, 3, 4, 5]).cuda())
# 使用torch.cuda.empty_cache()清理缓存
torch.cuda.empty_cache()
# 推理后
print(torch.cuda.memory_allocated())

运行结果：