在pytorch中为Module和Tensor指定GPU的例子资源-CSDN文库

197 浏览量 2021-01-20 06:34:31 上传评论收藏 57KB PDF 举报

资源推荐

资源详情

资源评论

在在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 "<ipython-input-15-ef42531f63ca>", line 1, in <module>

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[19]:

Variable containing:

.00000e-05 *

0.0000 0.0000 2.0419 0.0000 2.0420

0.0000 0.0000 0.0000 0.0000 0.0000

本内容试读结束，登录后可阅读更多

下载后可阅读完整内容，剩余5页未读，立即下载

内容反馈

weixin_38746166

粉丝: 8
资源: 959

最新资源

资源上传下载、课程学习等过程中有任何疑问或建议，欢迎提出宝贵意见哦~我们会及时处理！点击此处反馈

feedback-tip