Ошибка типа: невозможно преобразовать тензор CUDA в numpy. Используйте Tensor.cpu(), чтобы сначала скопировать тензор в память хоста.

Question

Я использую модифицированный файл gnast.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()

Это означает, что данные сначала перемещаются в процессор, а затем преобразуются в массив numpy.

Answer 2

Я узнал, что я могу просто использовать

output.argmax()

Answer 3

Ты можешь использовать torch.max работают следующим образом:

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