Переход на ЦП с помощью Google Colab

Question

Я использую модель в Google Colab. Последний шаг, который я хотел бы сделать, это распечатать изображение и показать 5 главных классификационных прогнозов модели. Вот код:

image = process_image(imgpath)

index = 17
plot = imshow(image, ax = plt)
plot.axis('off')
plot.title(cat_to_name[str(index)])
plot.show()

axes = predict(image, model)

yaxis = [cat_to_name[str(i)] for i in np.array(axes[1][0])]
y_pos = np.arange(len(yaxis))
xaxis = np.array(axes[0][0])   

plt.barh(y_pos, xaxis)
plt.xlabel('probability')
plt.yticks(y_pos, yaxis)
plt.title('probability of flower classification')

plt.show()

Я получаю эту ошибку при запуске этой ячейки:

TypeError                                 Traceback (most recent call last)
<ipython-input-19-d0bb6f461eec> in <module>()
     11 axes = predict(image, model)
     12 
---> 13 yaxis = [cat_to_name[str(i)] for i in np.array(axes[1][0])]
     14 y_pos = np.arange(len(yaxis))
     15 xaxis = np.array(axes[0][0])

/usr/local/lib/python3.6/dist-packages/torch/tensor.py in __array__(self, dtype)
    447     def __array__(self, dtype=None):
    448         if dtype is None:
--> 449             return self.numpy()
    450         else:
    451             return self.numpy().astype(dtype, copy=False)

TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

Есть ли способ временно использовать процессор в Google Colab и на этом конкретном этапе ? Мне не нужно переключаться обратно на GPU, потому что это последний шаг в моем коде.

Answer 1

Попробуйте следующее:

yaxis = [cat_to_name[str(i)] for i in axes[1][0].cpu()]
xaxis = axes[0][0].cpu().numpy() 

Answer 2

Я думаю, вам просто нужно изменить

axes = predict(image, model)

на

axes = predict(image, model).cpu()