Ben Cook: 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.

The post PyTorch Tensor to NumPy Array and Back appeared first on Sparrow Computing.

from Planet Python
via read more