Pytorch调用GPU训练两种方法
方法一 .cuda()
我们可以通过对网络模型,数据,损失函数这三种变量调用 .cuda() 来在GPU上进行训练。
-
# 将网络模型在gpu上训练
-
model = Model()
-
if torch.cuda.is_available():
-
model = model.cuda()
-
-
# 损失函数在gpu上训练
-
loss_fn = nn.CrossEntropyLoss()
-
if torch.cuda.is_available():
-
loss_fn = loss_fn.cuda()
-
-
# 数据在gpu上训练
-
for data in dataloader:
-
imgs, targets = data
-
if torch.cuda.is_available():
-
imgs = imgs.cuda()
-
targets = targets.cuda()
方法二 .to(device)
指定 训练的设备
-
device = torch.device("cpu") # 使用cpu训练
-
device = torch.device("cuda") # 使用gpu训练
-
device = torch.device("cuda:0") # 当电脑中有多张显卡时,使用第一张显卡
-
device = torch.device("cuda:1") # 当电脑中有多张显卡时,使用第二张显卡
-
-
#单卡推荐
-
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
使用 GPU 训练
-
model = model.to(device)
-
-
loss_fn = loss_fn.to(device)
-
-
for data in train_dataloader:
-
imgs, targets = data
-
imgs = imgs.to(device)
-
targets = targets.to(device)
查看GPU使用情况
nvidia-smi.exe -l 5
这篇好文章是转载于:学新通技术网
- 版权申明: 本站部分内容来自互联网,仅供学习及演示用,请勿用于商业和其他非法用途。如果侵犯了您的权益请与我们联系,请提供相关证据及您的身份证明,我们将在收到邮件后48小时内删除。
- 本站站名: 学新通技术网
- 本文地址: /boutique/detail/tanhghakgk
系列文章
更多
同类精品
更多
-
photoshop保存的图片太大微信发不了怎么办
PHP中文网 06-15 -
《学习通》视频自动暂停处理方法
HelloWorld317 07-05 -
word里面弄一个表格后上面的标题会跑到下面怎么办
PHP中文网 06-20 -
Android 11 保存文件到外部存储,并分享文件
Luke 10-12 -
photoshop扩展功能面板显示灰色怎么办
PHP中文网 06-14 -
微信公众号没有声音提示怎么办
PHP中文网 03-31 -
excel下划线不显示怎么办
PHP中文网 06-23 -
excel打印预览压线压字怎么办
PHP中文网 06-22 -
TikTok加速器哪个好免费的TK加速器推荐
TK小达人 10-01 -
怎样阻止微信小程序自动打开
PHP中文网 06-13