PyTorch 旨在与 NumPy 非常兼容。因此，PyTorch 的 Tensor 与 Numpy 的 array 能够非常方便的进行互相转换。本篇以 Jupyterlab 为平台进行介绍。

numpy to pytorch

常规转换

1
2
import numpy as np
import torch
1
2
x = np.eye(3)
x
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])

1
torch.from_numpy(x)
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]], dtype=torch.float64)

转换时改变数据类型

1
2
x = np.eye(3)
x
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])

1
torch.from_numpy(x).type(torch.float32)
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])

转换时发送到不同的设备上，如 GPU

1
2
x = np.eye(3)
x
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])

1
2
3
if torch.cuda.is_available():
    y = torch.from_numpy(x).to("cuda")
y
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]], device='cuda:0', dtype=torch.float64)

注意，当使用锁页内存（pytorch 中数据加载器的锁页内存 pinned memory）的方式加载数据时，数据放入 GPU 的时候，应该把 non_blocking=True，这样能够缩减访问时间，加快训练。

1
2
3
4
# 方法
if torch.cuda.is_available():
    y = torch.from_numpy(x).to("cuda", non_blocking=True)
y

pytorch to numpy

转向另一个方向会稍微复杂一些，因为有时您必须处理 PyTorch 张量和 NumPy 数组之间的两个差异：

  1. PyTorch 可以针对不同的设备（如 GPU）；
  2. PyTorch 支持自动微分。

在最简单的情况下，当你在 CPU 上有一个没有梯度的 PyTorch 张量时，你可以简单地调用 .numpy() 方法：

1
2
x = torch.eye(3)
x
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])

1
x.numpy()
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]], dtype=float32)

如果张量是需要梯度的计算图的一部分（也就是说，如果 x.requires_grad=True），则需要调用 .detach() 方法：

1
2
x = torch.eye(3, requires_grad=True)
x
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]], requires_grad=True)

1
x = torch.eye(3)
1
x
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])

1
2
x.requires_grad = True
x
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]], requires_grad=True)

1
x.detach().numpy()
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]], dtype=float32)

如果 Tensor 位于 “cpu” 以外的设备上，则需要先将其带回 CPU，然后才能调用 .numpy() 方法。我们在上面使用 .to(“cuda”) 向 GPU 发送张量时看到了这一点。现在，我们只是反过来：

1
2
x = torch.eye(3)
x
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])

1
2
x = x.to("cuda")
x
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]], device='cuda:0')

1
x.to('cpu').numpy()
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]], dtype=float32)

1
x
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]], device='cuda:0')

1
x.cpu().numpy()
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]], dtype=float32)

最完全最常用的将 Tensor 转成 numpy array 的方法如下：

1
2
x = torch.eye(3)
x
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])

1
x.detach().cpu().numpy()
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]], dtype=float32)

1
x.detach().to('cpu').numpy()
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]], dtype=float32)

转换后共享内存

注意，转换后的 pytorch tensor 与 numpy array 指向同一地址，所以，对一方的值改变另一方也随之改变

除 chartensor 外所有 tensor 都可以转换为 numpy

1
2
import torch
import numpy as np
1
2
a = torch.ones(5)
a
tensor([1., 1., 1., 1., 1.])

1
b = a.numpy()
1
b
array([1., 1., 1., 1., 1.], dtype=float32)

1
2
a.add_(1)
a
tensor([2., 2., 2., 2., 2.])

1
b
array([2., 2., 2., 2., 2.], dtype=float32)

1
b += 1
1
b
array([3., 3., 3., 3., 3.], dtype=float32)

1
a
tensor([3., 3., 3., 3., 3.])

1
2
c = np.ones(5)
c
array([1., 1., 1., 1., 1.])

1
2
d = torch.from_numpy(c)
d
tensor([1., 1., 1., 1., 1.], dtype=torch.float64)

1
2
np.add(c, 1, out=c)
c
array([2., 2., 2., 2., 2.])

1
d
tensor([2., 2., 2., 2., 2.], dtype=torch.float64)

1
2
d.add_(1)
d
tensor([3., 3., 3., 3., 3.], dtype=torch.float64)

1
c
array([3., 3., 3., 3., 3.])

参考文献

  1. PyTorch Tensor to NumPy Array and Back
  2. pytorch tensor与numpy转换