Ben Chuanlong Du's Blog

It is never too late to learn.

Convert a Tensor to a Numpy Array or List in PyTorch

Tips

There are multiple ways to convert a Tensor to a numpy array in PyTorch. First, you can call the method Tensor.numpy.

my_tensor.numpy()

Second, you can use the function numpy.array.

import numpy as np
np.array(my_tensor)

It is suggested that you use the function numpy.array to convert a Tensor to a numpy array. The reason is that numpy.array is more generic. You can also use it to convert other objects (e.g., PIL.Image) to numpy arrays while those objects might not have a method named numpy.

Notice that a Tensor on CUDA cannot be converted to a numpy array directly. You have to move it to CPU first and then convert to a numpy array.

import numpy as np
np.array(my_tensor.to("cpu"))

As a matter of fact, this is the suggested way to convert a Tensor to a numpy array as it works for Tensors on different devices.

Tensor to Numpy Array

In [1]:
import torch
import numpy as np
In [2]:
x = torch.tensor(
    [
        [1, 2, 3, 4, 5],
        [6, 7, 8, 9, 10],
    ]
)
x
Out[2]:
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])
In [3]:
x.numpy()
Out[3]:
array([[ 1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10]])
In [4]:
np.array(x)
Out[4]:
array([[ 1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10]])

Move a Tensor to CPU and then convert it to a numpy array. This is the suggested way to convert a Tensor to a numpy array.

In [7]:
np.array(x.to("cpu"))
Out[7]:
array([[ 1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10]])

Tensor to list

  1. You can also convert a Tensor to a list using the method Tensor.tolist. If you work across different devices (CPU, GPU, etc.), it is suggested that you use Tensor.cpu().tolist.

  2. list(tensor) put all first-level element of tensor into a list. It is the reverse operation of torch.stack. This is not what you want, geneerally speaking.

Convert the tensor x to a list.

In [3]:
x.tolist()
Out[3]:
[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]

Move the tensor to CPU first and then convert it to a list. This is the recommended way to convert a Tensor to a list.

In [6]:
x.cpu().tolist()
Out[6]:
[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]
In [4]:
list(x)
Out[4]:
[tensor([1, 2, 3, 4, 5]), tensor([ 6,  7,  8,  9, 10])]
In [5]:
list(torch.tensor([1, 2, 3]))
Out[5]:
[tensor(1), tensor(2), tensor(3)]

References

https://stackoverflow.com/questions/54268029/how-to-convert-a-pytorch-tensor-into-a-numpy-array

In [ ]:

Comments