从品牌网站建设到网络营销策划,从策略到执行的一站式服务
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)无效。
成都网站建设公司地址:成都市青羊区太升南路288号锦天国际A座10层 建设咨询028-86922220
成都快上网科技有限公司-四川网站建设设计公司 | 蜀ICP备19037934号 Copyright 2020,ALL Rights Reserved cdkjz.cn | 成都网站建设 | © Copyright 2020版权所有.
专家团队为您提供成都网站建设,成都网站设计,成都品牌网站设计,成都营销型网站制作等服务,成都建网站就找快上网! | 成都网站建设哪家好? | 网站建设地图