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

Question

如何启用pytorch在GPU上运行？

我已在google colab笔记本中成功安装pytorch： Tensorflow报告GPU到位：

但是torch.device函数以某种方式失败：

我怎样才能解决这个问题？

Answer 1

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

device您可以使用以下方式定义torch.device：

import torch

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

Answer 2

我遇到了同样的问题。

尝试像这样安装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火炬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 。