Torch 中显存回收节省显存的方法

在进行模型推理时，需要考虑如何有效地利用和管理GPU显存。以下总结了常用的节省显存的方法。

将变量从显存释放

1. 使用.detach()方法：在PyTorch中，可以使用.detach()方法来将计算结果从计算图中分离出来，这样就可以防止PyTorch在后续计算中保存这些结果的梯度信息，从而节省显存。

    output = model(input).detach()

2. 使用torch.no_grad()上下文管理器：在进行推理时，通常不需要计算梯度，因此可以使用torch.no_grad()来禁止梯度计算，这也可以节省显存。

   with torch.no_grad():
        output = model(input)

3. 删除不再使用的变量：

    del variable
    torch.cuda.empty_cache()

4. 使用更小的数据类型：可以考虑使用更小的数据类型（例如，使用float16代替float32）来节省显存。但是请注意，这可能会影响计算的精度。

  model = model.half()  # convert model to float16
  input = input.half()  # convert input to float16

或使用模型的 bf16 等版本。

5. 批量处理数据：一次处理多个数据，而不是逐个处理。这样可以更有效地利用显存，因为一些显存是用来存储模型的参数和中间结果的，而这些显存的使用量通常不会随着批量大小的增加而增加。

    outputs = []
    for batch in dataloader:
        with torch.no_grad():
            output = model(batch)
        outputs.append(output.cpu())  # move output to CPU memory

6. 将模型和数据移动到CPU内存：如果你的模型和数据都在GPU显存中，那么你可以考虑在完成推理后将它们移动到CPU内存，以释放显存。

model = model.cpu()
data = data.cpu()

请注意，以上的方法可能会影响推理的速度和精度，因此在使用这些方法时，你需要根据你的具体需求来权衡显存的使用量、推理的速度和精度。

释放缓存

在PyTorch中，即使你已经将Tensor从GPU转移到CPU并删除了相应的变量，有时GPU的显存仍然可能不会立即释放。这是因为PyTorch使用了一种称为"缓存分配器"的机制来管理显存，这种机制可以减少显存的分配和释放操作，从而提高效率。

当你删除一个Tensor并释放了它占用的显存后，这部分显存并不会立即返回给操作系统，而是被缓存分配器保留下来，以便在后续的操作中重复使用。因此，即使你已经删除了所有的Tensor，你可能仍然会看到GPU的显存被占用。

如果你想要立即释放这部分显存，你可以使用torch.cuda.empty_cache()函数来清空缓存分配器。但是请注意，这个操作可能会降低你的代码的效率，因为在后续的操作中，PyTorch可能需要重新分配显存。

# your code here
output = model(input)
output = output.cpu()  # move output to CPU

del output  # delete the variable
torch.cuda.empty_cache()  # clear the cache

如果你发现即使使用了torch.cuda.empty_cache()，GPU的显存仍然持续增长，那么可能是因为你的代码中存在其他的问题，例如你可能在某些地方忘记了删除一些不再需要的Tensor，或者你可能在循环中创建了一些不必要的Tensor。你需要仔细检查你的代码，确保所有的Tensor都被正确地管理。