numpy、cupy、pytorch数组对象的相互转换

记录平常最常用的三个python对象之间的相互转换：numpy，cupy，pytorch三者的ndarray转换

1. numpy与cupy互换

import numpy as np
import cupy as cp
A = np.zeros((4,4))
B = cp.asarray(A) # numpy -> cupy
C = cp.asnumpy(B) # cupy -> numpy
print(type(A), type(B), type(C))

输出：

<class 'numpy.ndarray'> <class 'cupy._core.core.ndarray'> <class 'numpy.ndarray'>

2. numpy与pytorch互换

import torch
import numpy as np
A = np.zeros((4,4))
B = torch.tensor(A) # numpy -> pytorch 方法1
C = torch.from_numpy(A) # numpy -> pytorch 方法2
D = B.numpy() # pytorch -> numpy
print(type(A), type(B), type(C), type(D))

输出：

<class 'numpy.ndarray'> <class 'torch.Tensor'> <class 'torch.Tensor'> <class 'numpy.ndarray'>

将大小为1的张量转换为Python标量, 我们可以调用item函数或Python的内置函数。

a = torch.tensor([3.5])
a, a.item(), float(a), int(a)

输出:

(tensor([3.5000]), 3.5, 3.5, 3)

3. cupy与pytorch的互换

import cupy as cp
from torch.utils.dlpack import to_dlpack,from_dlpack
from cupy import fromDlpack
A = cp.zeros((4,4))
B = from_dlpack(A.toDlpack()) # cupy -> pytorch
C = fromDlpack(to_dlpack(B)) # pytorch -> cupy
print(type(A), type(B), type(C))

输出：

<class 'cupy._core.core.ndarray'> <class 'torch.Tensor'> <class 'cupy._core.core.ndarray'>

其他

numpy和cupy的默认数据类型是float64, pytorch默认是float32

import torch
import numpy as np
import cupy as cp
A = np.zeros((4,4))
B = torch.zeros((4,4))
C = cp.zeros((4,4))
A.dtype, B.dtype, C.dtype

输出：

(dtype('float64'), torch.float32, dtype('float64'))