我该如何在Google Colab中启用pytorch GPU支持？

Question

如何启用pytorch在GPU上工作？

我已经成功在Google Colab笔记本中安装了pytorch： TensorFlow报告GPU已就绪：

但是，torch.device函数出现了错误：

我该如何解决这个问题？

- Stepan Yakovenko

---

当我尝试时，它可以工作。 - Omegastick

Answer 1

我遇到了同样的问题。

尝试按照以下方式安装Torch：

# http://pytorch.org/
from os import path
from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag
platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())

accelerator = 'cu80' #'cu80' if path.exists('/opt/bin/nvidia-smi') else 'cpu'
print('Platform:', platform, 'Accelerator:', accelerator)

!pip install --upgrade --force-reinstall -q http://download.pytorch.org/whl/{accelerator}/torch-0.4.0-{platform}-linux_x86_64.whl torchvision

import torch
print('Torch', torch.__version__, 'CUDA', torch.version.cuda)
print('Device:', torch.device('cuda:0'))

输出应该是:

平台: cp36-cp36m 加速器: cu80 Torch 0.4.0 CUDA 8.0.61
设备: cuda:0

一些流传的代码片段使用 torch-0.3.0.post4-{platform}-linux_x86_64.whl，这会导致相同的错误，因为 device 是 Torch 4 的特性。如果您已经安装了错误的版本，则可能需要执行 !pip uninstall torch。

还要确保在 编辑 > 笔记本设置 > 硬件加速器 下启用 GPU。

- Adam Bittlingmayer

Answer 2

您可以通过点击“运行时”菜单下的“更改运行时类型”来启用GPU。现在还提供了“TPU”的支持。

您可以使用 torch.device 来定义 device：

import torch

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

- abdullahselek

Answer 3

您可以使用此教程：https://medium.com/@nrezaeis/pytorch-in-google-colab-640e5d166f13。

例如，对于CUDA 9.2和Python 3.6:

!pip3 install http://download.pytorch.org/whl/cu92/torch-0.4.1-cp36-cp36m-linux_x86_64.whl
!pip3 install torchvision

现在使用PyTorch检查GPU设备：

torch.cuda.get_device_name(0)

我在Google Colab中的结果是Tesla K80。

- NRR

Answer 4

除了在“运行时” -> “更改运行时类型”菜单下启用GPU外，还可以通过以下方式启用GPU支持：

import torch

if torch.cuda.is_available():
  device = torch.device("cuda")
else:
  device = torch.device("cpu")

- Christian