前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >pytorch基础知识-GPU加速

pytorch基础知识-GPU加速

作者头像
用户6719124
发布2019-11-17 22:07:05
1.1K0
发布2019-11-17 22:07:05
举报
文章被收录于专栏:python pytorch AI机器学习实践

本节比较简单,介绍一个显卡加速功能。

一般我们在使用笔记本电脑或者台式机进行神经网络结构计算时,默认使用cpu计算,但cpu运算速度十分有限,一个专门搞学术研究的人常配备一个英伟达显卡来加速计算。

GPU加速功能可以将运算切入到显卡中进行,从而提高运算速度。

该方法在pytorch 0.3版本以前较麻烦,当时是在代码后面加入.cpu()进行。

在新版本的pytorch中,变为统一设置运算位置的形式。

如上段代码中可以加入以下代码来提高运算速度。

首先定义device(设备),再调用.to函数

在使用该项功能前 首先确认自己电脑有GPU英伟达显卡,且支持CUDA模块,

随后确认自己电脑里安装了CUDA,

可以使用该代码来查看当前环境是否支持CUDA

代码语言:javascript
复制
import torch
print(torch.cuda.is_available())
# 返回True代表支持,False代表不支持

具体在神经网络中,该这样使用

代码语言:javascript
复制
import torch
import torch.nn as nn
import optim


device = torch.device('cuda:0')
# 'cuda:0'当中的0为想要使用显卡的编号
# 这里的0表示使用的是第一张显卡
net = MLP().to(device)
# 使用.to函数将神经网络模块搬到MLP上进行运算
optimizer = optim.SGD(net.parameters(), lr=1e-3)
criteon = nn.CrossEntropyLoss().to(device)
# 同样将loss部分的计算转移到GPU上去

同样的,数据部分也可以转移到GPU上去

代码语言:javascript
复制
data, target = data.to(device), target.to(device)

这里要注意同一个数据在CPU和在GPU上建立后是完全不一样的。

本文参与 腾讯云自媒体同步曝光计划,分享自微信公众号。
原始发表:2019-11-01,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 python pytorch AI机器学习实践 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档