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

Question

我正在使用修改后的Forecast.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()