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

Question

我正在使用modified predict.py测试pruned SqueezeNet Model

[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)