[PyTorch] .detach().cpu().numpy()와 .cpu().data.numpy() ?

PyTorch를 사용하다보면, Module을 통해 나온 Tensor을 후처리에 사용하거나, 계산된 loss Tensor을 logging 하는 일이 많다. 

for i, data in enumerate(dataloader):
    # 생략
    embeddings = model(images)
    
    # total_embeddings에 embeddings append
    # total_embeddings.append(embeddings) ???

 

이때 embeddings는 GPU에 올라가있는 Tensor 이기 때문에 numpy 혹은 list로의 변환이 필요하다. 오픈소스를 보면 detach(), cpu(), data, numpy(), tolist() 등을 조합해서 변환을 한다. 하지만 stackoverflow나 pytorch discuss를 보면 이 methods/attributes 의 순서가 중요하다고 한다. 이번 포스팅에는 각 메서드에 대한 소개와 조합 순서의 차이에 대해 설명한다.

Why do we call .detach() before calling .numpy() on a Pytorch Tensor?

.detach()

Automatic differentiation package - torch.autograd — PyTorch 1.8.1 documentation

PyTorch docs에 따르면 "Returns a new Tensor, detached from the current graph. The result will never require gradient." 즉 graph에서 분리한 새로운 tensor를 리턴한다. 파이토치는 tensor에서 이루어진 모든 연산을 추적해서 기록해놓는다(graph). 이 연산 기록으로 부터 도함수가 계산되고 역전파가 이루어지게 된다. detach()는 이 연산 기록으로 부터 분리한 tensor을 반환하는 method이다.

.data

Release Trade-off memory for compute, Windows support, 24 distributions with cdf, variance etc., dtypes, zero-dimensional Tensor

What about .data?

.data was the primary way to get the underlying Tensor from a Variable. After this merge, calling y = x.data still has similar semantics. So y will be a Tensor that shares the same data with x, is unrelated with the computation history of x, and has requires_grad=False.

However, .data can be unsafe in some cases. Any changes on x.data wouldn't be tracked by autograd, and the computed gradients would be incorrect if x is needed in a backward pass. A safer alternative is to use x.detach(), which also returns a Tensor that shares data with requires_grad=False, but will have its in-place changes reported by autograd if x is needed in backward.

Variable에서 값을 얻는 유용한 attribute 였다. v0.4.0 부터 Variable과 Tensor가 합쳐지면서 .detach()를 권장한다고 한다.

.cpu()

Returns a copy of this object in CPU memory.

If this object is already in CPU memory and on the correct device, then no copy is performed and the original object is returned

GPU 메모리에 올려져 있는 tensor를 cpu 메모리로 복사하는 method

.numpy()

Returns self tensor as a NumPy ndarray. This tensor and the returned ndarray share the same underlying storage. Changes to self tensor will be reflected in the ndarray and vice versa.

tensor를 numpy로 변환하여 반환. 이때 저장공간을 공유하기 때문에 하나를 변경하면 다른 하나도 변경된다.

import torch
import numpy as np

embeddings_np = np.ones((2, 2))
embeddings_t = torch.Tensor(embeddings_np)

embeddings_numpy = embeddings_t.numpy()
# embeddings_numpy
# array([[1., 1.],
#        [1., 1.]], dtype=float32)

embeddings_t.add_(1)
# embeddings_numpy
# array([[2., 2.],
#        [2., 2.]], dtype=float32)

 

또한 cpu 메모리에 올려져 있는 tensor만 .numpy() method를 사용할 수 있다.

.detach() .cpu() .numpy() 순서

위 .numpy() 예제를 보면 GPU 메모리에 올려져 있는 tensor를 numpy로 변환하기 위해서는 우선 cpu 메모리로 옮겨야 한다. 따라서 .numpy() 이전에 .cpu()가 실행되야 한다. 그러면 .detach()와 .cpu() 중 어떤 것을 먼저 해야할까?

Should it really be necessary to do var.detach().cpu().numpy()?

위에서 언급한대로, gradient가 계산될 tensor의 경우 graph로 기록되어있다(node는 Tensor, edge는 입력 Tensor로부터 출력 Tensor를 만들어내는 함수). 만약 .cpu().detach()를 수행하게 되면, cpu 를 만드는 edge가 생성된다. 이 edge는 곧 없어지긴 한다. 반면 .detach().cpu()의 경우 이 작업을 수행하지 않기 때문에 주로 .detach().cpu()를 사용한다. (하지만 사실상 별 차이 없을듯...)

 

결론: .detach().cpu().numpy()를 사용하자!!

 

pytorch graph에 대해 더 알고 싶으면 아래의 자료들을 추천한다.

A Gentle Introduction to torch.autograd — PyTorch Tutorials 1.8.1+cu102 documentation

PyTorch Autograd