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

【问题标题】：TypeError: can’t convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory firstTypeError：无法将 CUDA 张量转换为 numpy。首先使用 Tensor.cpu() 将张量复制到主机内存
【发布时间】：2019-05-22 21:04:09
【问题描述】：

我正在使用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()的情况下，我应该如何修改代码以摆脱此错误？

【问题讨论】：

标签： numpy neural-network pytorch pruning

【解决方案1】：

改变

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

到

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

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

【讨论】：

【解决方案2】：

我发现我可以使用

output.argmax()

【讨论】：

我不认为这会给你 numpy 数组。其输出将是火炬张量
@UmangGupta 火炬张量是我需要的。请参阅gitlab.com/promach/Pruning-CNN/tree/master/SqueezeNet-Pruning 的工作示例
我被你的问题的标题弄糊涂了。

【解决方案3】：

你可以使用torch.max函数如下：

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

【讨论】：