记录平常最常用的三个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'))