是否需要清除 PyTorch 中的 GPU 张量？

Question

I am new to PyTorch, and I am exploring the functionality of .to() method. As per the documentation for the CUDA tensors, I see that it is possible to transfer the tensors between the CPU and GPU memory.

# let us run this cell only if CUDA is available
if torch.cuda.is_available():

    # creates a LongTensor and transfers it to GPU as torch.cuda.LongTensor
    a = torch.full((10,), 3, device=torch.device("cuda"))
    # transfers it to CPU, back to being a torch.LongTensor
    b = a.to(torch.device("cpu"))

In this context, I would like to know if it always necessary to transfer back the tensors from the GPU to CPU, perhaps to free the GPU memory? Doesn't, the runtime automatically clear the GPU memory?

Apart from its use of transferring data between CPU and GPU, I would like to know the recommended usage for .to() method (from the memory perspective). Thanks in advance.

Answer 1

In this context, I would like to know if it always necessary to transfer back the tensors from the GPU to CPU, perhaps to free the GPU memory?

No, it's not always necessary. Memory should be freed when there are no more references to GPU tensor. Tensor should be cleared automatically in this case:

def foo():
    my_tensor = torch.tensor([1.2]).cuda()
    return "whatever"

smth = foo()

but it won't in this case:

def bar():
    return torch.tensor([1.2]).cuda()

tensor = bar()

In the second case (tensor being passed around, possibly accumulated or added to list), you should cast it to CPU in order not to waste GPU memory.

Apart from its use of transferring data between CPU and GPU, I would like to know the recommended usage for .to() method (from the memory perspective)

Not sure what you mean here. What you should be after is the least to calls as they require copying the array (O(n) complexity), but shouldn't be too costly anyway (in comparison to pushing data through neural network for example) and probably not worth to be too hardcore about this micro-optimization.

Usually data loading is done on CPU (transformations, augmentations) and each batch copied to GPU (possibly with pinned memory) just before it is passed to neural network.

Also, as of 1.5.0 release, pytorch provides memory_format argument in .to method. This allows you to specify whether (N, C, H, W) (PyTorch default) or channel last (N, H, W, C) should be used for tensor and model (convolutional models with torch.nn.Conv2d to be precise). This could further speed up your models (for torchvision midels speedup of 16% was reported IIRC), see here for more info and usage.