pytorch使用指定GPU训练

本文适合多GPU的机器,并且每个用户需要单独使用GPU训练。

虽然pytorch提供了指定gpu的几种方式,但是使用不当的话会遇到out of memory的问题,主要是因为pytorch会在第0块gpu上初始化,并且会占用一定空间的显存。这种情况下,经常会出现指定的gpu明明是空闲的,但是因为第0块gpu被占满而无法运行,一直报out of memory错误。解决方案如下:

  1. 指定环境变量,屏蔽第0块gpu
CUDA_VISIBLE_DEVICES = 1 main.py

这句话表示只有第1块gpu可见,其他gpu不可用,此时要注意第1块gpu已经变成第0块,因此代码里应该使用编号0来指定gpu。如果依然使用cuda:1会报invalid device ordinal
当然也可以在代码里使用os模块达到同样效果:

import os

#多块使用逗号隔开
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
  1. 使用torch.cuda.device
with torch.cuda.device(1):
	...
注意这种情况下gpu编号也会改变

关于这个问题的讨论,感兴趣的也可以去这里查看cuda out of memory error when GPU0 memory is fully utilized

PyTorch是一个深度学习框架,它支持使用GPU进行训练,以加速模型的计算。在PyTorch中,可以使用`torch.cuda`模块来管理GPU使用。如果你有多块GPU,可以通过设置环境变量`CUDA_VISIBLE_DEVICES`来指定使用哪些GPU。此外,在创建张量(tensor)或者模型时,可以指定`.to('cuda')`来将数据或模型转移到GPU上。 下面是一个简单的例子,展示了如何在PyTorch指定使用GPU进行训练: ```python import torch import torch.nn as nn import torch.optim as optim # 检查是否有可用的GPU,并设置 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 创建模型,并转移到GPU上(如果可用) model = MyModel().to(device) # 创建损失函数 criterion = nn.CrossEntropyLoss() # 创建优化器 optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) # 训练模型 for epoch in range(num_epochs): for inputs, labels in dataloader: inputs, labels = inputs.to(device), labels.to(device) # 前向传播 outputs = model(inputs) loss = criterion(outputs, labels) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}') ``` 在这个例子中,`MyModel`应该是一个自定义的神经网络模型类。当`torch.cuda.is_available()`返回`True`时,模型和数据会被送到GPU上进行训练。如果当前机器上没有安装CUDA或者没有可用的GPU,`device`将会是`cpu`。
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值