33

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

4

3 回答 3

53

改变

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

至

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

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

于 2018-12-23T05:05:59.270 回答
3

我发现我可以使用

output.argmax()
于 2018-12-23T09:43:55.017 回答
1

您可以使用torch.max如下功能：

value, index = torch.max(output,1)
于 2020-05-27T20:12:50.893 回答