pytorch 使用GPU

# 查看gpu是否可用，在pytorch中cuda指的就是gpu的相关操作
>>> torch.cuda.is_available()
True
# 获取cuda设备
>>> device = torch.device('cuda')
>>> device
device(type='cuda')
# 随机创建一个两行三列的正态分布的tensor，这个tensor是在cpu上
>>> x = torch.randn(2,3)
>>> x
tensor([[-1.2191, -0.2216, -1.3021],
        [ 0.6544,  1.4655,  0.0568]])
>>> x.device  # device是x是一个属性，而不是方法
device(type='cpu')
# 直接在gpu上创建一个tensor，需要指定device参数
>>> y = torch.ones_like(x, device=device)
>>> y
tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')
>>> y.device # y的device是'cuda'而且是第0个
device(type='cuda', index=0)
>>> x = x.to(device) # 把x发送到gpu上
>>> x
tensor([[-1.2191, -0.2216, -1.3021],
        [ 0.6544,  1.4655,  0.0568]], device='cuda:0')
>>> z = x + y # 计算
>>> z
tensor([[-0.2191,  0.7784, -0.3021],
        [ 1.6544,  2.4655,  1.0568]], device='cuda:0')
>>> z_cpu = z.to('cpu', dtype=torch.double) # 把z从gpu移动到cpu
>>> z_cpu
tensor([[-0.2191,  0.7784, -0.3021],
        [ 1.6544,  2.4655,  1.0568]], dtype=torch.float64)
>>> z.device
device(type='cuda', index=0)
>>> z_cpu.device
device(type='cpu')

另一种方法：直接使用cuda()函数或cpu()函数实现tensor从gpu和cup之间的切换。

# -*- coding: utf-8 -*-
import torch
from torch.autograd import Variable
 
# 将变量或者数据移到GPU
gpu_info = Variable(torch.randn(3,3)).cuda()
# 将变量或者数据移到CPU
cpu_info = gpu_info.cpu()

如果是cpu上的tensor与gpu上的tensor计算，则会报RuntimeError的错误：

>>> x
tensor([[-0.3749, -3.1064,  0.8493],
        [ 0.0398, -0.1842, -0.0846]])
>>> x.device
device(type='cpu')
>>> y.device
device(type='cuda', index=0)
>>> i = x + y
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: expected device cpu but got device cuda:0