TypeError：无法将 CUDA 张量转换为 numpy。首先使用 Tensor.cpu() 将张量复制到主机内存

Question

我正在使用修改后的predict.py来测试修剪的SqueezeNet模型

[phung@archlinux SqueezeNet-Pruning]$ python predict.py --image 3_100.jpg --model model_prunned --num_class 2
prediction in progress
Traceback (most recent call last):
File “predict.py”, line 66, in
prediction = predict_image(imagepath)
File “predict.py”, line 52, in predict_image
index = output.data.numpy().argmax()
TypeError: can’t convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
[phung@archlinux SqueezeNet-Pruning]$

我了解 numpy 还不支持 GPU。

如何修改代码来避免此错误而不调用张量复制数据操作，

tensor.cpu()

？

Answer 1

改变

index = output.data.numpy().argmax()

到

index = output.cpu().data.numpy().argmax()

这意味着数据首先移动到CPU，然后转换为numpy数组。

Answer 2

我发现我可以使用

output.argmax()

Answer 3

您可以使用

torch.max

功能，如下所示：

value, index = torch.max(output,1)

Answer 4

我在下面遇到了同样的错误：

TypeError：无法将 cuda:0 设备类型张量转换为 numpy。首先使用 Tensor.cpu() 将张量复制到主机内存。

当我尝试使用 numpy() 将 CUDA(GPU) PyTorch 张量转换为 NumPy 数组时，如下所示：

import torch

my_tensor = torch.tensor([0., 1., 2.], device='cuda:0')

my_tensor.numpy() # Error

所以首先，我使用 cpu() 将 PyTorch 张量的设备转换为 CPU，如错误消息所示，或者使用 to() 或将

force=True

设置为

numpy()

，然后我可以将 PyTorch 张量转换为带有

numpy()

的 NumPy 数组，如下所示。 *

force=True

可以创建张量的副本：

import torch

my_tensor = torch.tensor([0., 1., 2.], device='cuda:0')

my_tensor.cpu().numpy()
my_tensor.to(device='cpu').numpy()
my_tensor.numpy(force=True)
# array([0., 1., 2.], dtype=float32)