python - 使用 Google Colab 改回 CPU

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 上以及在此特定步骤中临时使用 CPU？我真的不需要切换回 GPU，因为这是我代码中的最后一步。

Answer 1

尝试以下操作:

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