In PyTorch is GPU memory freed when a GPU tensor is assigned a new value?

Viewed 2100

When a Cuda variable in PyTorch is assigned a new value, it becomes a CPU variable again (As illustrated by the code below). In this case, is the memory held by the variable on the GPU previously is freed automatically?

import torch

t1 = torch.empty(4,5)

if torch.cuda.is_available():
  t1 = t1.cuda()

print(t1.is_cuda)

t1 = torch.empty(4,5)
print(t1.is_cuda)

The output of the above code is:

True
False
1 Answers

In python an object is freed as soon as there are no remaining references to it. Since you assign t1 to reference a brand new tensor there are no more references to the original GPU tensor so that tensor is freed. That said, when PyTorch is instructed to free a GPU tensor it tends to cache that GPU memory for a while since it's usually the case that if we used GPU memory once we will probably want to use some again, and GPU memory allocation is relatively slow. If you want to force this cache of GPU memory to be cleared you can use torch.cuda.empty_cache. Using this won't directly increase the GPU memory available within a single PyTorch instance since PyTorch will call it automatically in an attempt to save you from an out of GPU memory error.

To reiterate, the GPU tensor doesn't actually "become" a CPU tensor. In python, variable names are references to objects. What your code really does is assign t1 to refer to a new CPU tensor object. Internally, python counts the number of references for each object. When that count goes to zero that object is immediately freed.

Caveat (Reference Cycles): Reference counting fails in the case of unreachable reference cycles. Unreachable reference cycles occur when objects contain references to one-another but no reference to any object in the cycle is reachable. To deal with this python employs a garbage collection module which executes intermittently. This module uses more sophisticated algorithms to detect and free objects that are part of unreachable reference cycles. In these cases, the memory is not necessarily freed when a cycle becomes unreachable and will instead be freed once the internal garbage collector is activated. This occurs automatically and relatively unpredictably. If desired, the garbage collector may be queried, configured, or manually executed using python's built in gc garbage collection interface.

Based on the previous discussion if you really want to ensure your unreachable GPU memory is freed in PyTorch (even in the case of unreachable reference cycles) you could use

import gc
gc.collect()
torch.cuda.empty_cache()
Related