PyTorch Tensor to NumPy Array and Back

NumPy to PyTorch

PyTorch is designed to be pretty compatible with NumPy. Because of this, converting a NumPy array to a PyTorch tensor is simple:

import torch
import numpy as np

x = np.eye(3)

torch.from_numpy(x)

# Expected result
# tensor([[1., 0., 0.],
#         [0., 1., 0.],
#         [0., 0., 1.]], dtype=torch.float64)

All you have to do is use the torch.from_numpy() function.

Once the tensor is in PyTorch, you may want to change the data type:

x = np.eye(3)

torch.from_numpy(x).type(torch.float32)

# Expected result
# tensor([[1, 0, 0],
#         [0, 1, 0],
#         [0, 0, 1]])

All you have to do is call the .type() method. Easy enough.

Or, you may want to send the tensor to a different device, like your GPU:

x = np.eye(3)

torch.from_numpy(x).to("cuda")

# Expected result
# tensor([[1., 0., 0.],
#         [0., 1., 0.],
#         [0., 0., 1.]], device='cuda:0', dtype=torch.float64)

The .to() method sends a tensor to a different device. Note: the above only works if you’re running a version of PyTorch that was compiled with CUDA and have an Nvidia GPU on your machine. You can test whether that’s true with torch.cuda.is_available().

PyTorch to NumPy

Going the other direction is slightly more involved because you will sometimes have to deal with two differences between a PyTorch tensor and a NumPy array:

  1. PyTorch can target different devices (like GPUs).
  2. PyTorch supports automatic differentiation.

In the simplest case, when you have a PyTorch tensor without gradients on a CPU, you can simply call the .numpy() method:

x = torch.eye(3)

x.numpy()

# Expected result
# array([[1., 0., 0.],
#        [0., 1., 0.],
#        [0., 0., 1.]], dtype=float32)

But, if the tensor is part of a computation graph that requires a gradient (that is, if x.requires_grad is true), you will need to call the .detach() method:

x = torch.eye(3)
x.requires_grad = True

x.detach().numpy()

# Expected result
# array([[1., 0., 0.],
#        [0., 1., 0.],
#        [0., 0., 1.]], dtype=float32)

And if the tensor is on a device other than "cpu", you will need to bring it back to the CPU before you can call the .numpy() method. We saw this above when sending a tensor to the GPU with .to("cuda"). Now, we just go in reverse:

x = torch.eye(3)
x = x.to("cuda")

x.to("cpu").numpy()

# Expected result
# array([[1., 0., 0.],
#        [0., 1., 0.],
#        [0., 0., 1.]], dtype=float32)

Both the .detach() method and the .to("cpu") method are idempotent. So, if you want to, you can plan on calling them every time you want to convert a PyTorch tensor to a NumPy array, even when it’s not strictly necessary:

x = torch.eye(3)

x.detach().to("cpu").numpy()

# Expected result
# array([[1., 0., 0.],
#        [0., 1., 0.],
#        [0., 0., 1.]], dtype=float32)

By the way, if you want to perform image transforms on a NumPy array directly you can! All you need is to have a transform that accepts NumPy arrays as input. Check out my post on TorchVision transforms if you want to learn more.

Connect

Contact

ben [at] sparrow [dot] dev

Email List