pytorch 指定卡1_在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)无效。

同样滴，Variable变量和Module类型的模型也可以指定放在哪块GPU上

v = Variable(a)

v

Out[18]:

Variable containing:

.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]

v.cuda(0)

Out[