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:
- PyTorch can target different devices (like GPUs).
- 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.