cupy与numpy、pytorch Tensor转换

一、cupy与numpy互转

import cupy as cp
import numpy as np

#cupy->numpy
numpy_data = cp.asnumpy(cupy_data)

#numpy->cupy
cupy_data = cp.asarray(numpy_data)
           

注：cupy->numpy过程较慢

二、cupy与pytorch Tensor互转

需要借助中间库 dlpack，三者关系

Cupy.array<–>Dlpack.Tensor<–>torch.Tensor

from cupy.core.dlpack import toDlpack
from cupy.core.dlpack import fromDlpack
from torch.utils.dlpack import to_dlpack
from torch.utils.dlpack import from_dlpack
import torch

#tensor->cupy
cupy_data = fromDlpack(to_dlpack(tensor_data))

#cupy->tensor
tensor_data = from_dlpack(toDlpack(cupy_data))
           

三、numpy与pytorch Tensor互转

import numpy as np
import torch

#tensor->numpy
numpy_data = tensor_data.numpy()

#numpy->tensor
tensor_data = torch.from_numpy(numpy_data)