在pytorch中为Module和Tensor指定GPU的例子-创新互联

pytorch指定GPU

在用pytorch写CNN的时候，发现一运行程序就卡住，然后cpu占用率100%，nvidia-smi 查看显卡发现并没有使用GPU。所以考虑将模型和输入数据及标签指定到gpu上。

pytorch中的Tensor和Module可以指定gpu运行，并且可以指定在哪一块gpu上运行，方法非常简单，就是直接调用Tensor类和Module类中的 .cuda() 方法。

import torch
from PIL import Image
import torch.nn as nn
import numpy as np
from torch.autograd import Variable

# 先看看有没有显卡
torch.cuda.is_available()
Out[16]: True
# 嗯，有显卡，可以指定，先生成一个Tensor
a = torch.Tensor(3,5)
a
Out[13]: 
.00000e-05 *
 0.0000 0.0000 2.0419 0.0000 2.0420
 0.0000 0.0000 0.0000 0.0000 0.0000
 0.0132 0.0000 0.0131 0.0000 0.0000
[torch.FloatTensor of size 3x5]
a.cuda()
Out[14]: 
.00000e-05 *
 0.0000 0.0000 2.0419 0.0000 2.0420
 0.0000 0.0000 0.0000 0.0000 0.0000
 0.0132 0.0000 0.0131 0.0000 0.0000
[torch.cuda.FloatTensor of size 3x5 (GPU 0)]
# 可以看到上面显示了（GPU 0）,也就是说这个Tensor是在第一个GPU上的
a.cuda(1)
Traceback (most recent call last):

 File "", line 1, in 
  a.cuda(1)

 File "/home/chia/anaconda2/lib/python2.7/site-packages/torch/_utils.py", line 57, in _cuda
  with torch.cuda.device(device):

 File "/home/chia/anaconda2/lib/python2.7/site-packages/torch/cuda/__init__.py", line 127, in __enter__
  torch._C._cuda_setDevice(self.idx)

RuntimeError: cuda runtime error (10) : invalid device ordinal at torch/csrc/cuda/Module.cpp:84
# 这个报错了，因为只有一块GPU，所以指定cuda(1)无效。