Pytorch机器学习GPU使用(从CPU转换为GPU)

本文介绍了如何在PyTorch中检查GPU的可用性,转移数据和模型到GPU,以及如何进行多GPU训练。通过torch.cuda.is_available()和torch.cuda.device_count()检查GPU状态,使用.to(device)或.cuda()迁移张量和模型到GPU,利用torch.nn.DataParallel实现单主机多GPU训练。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

torch.cuda.is_available()的返回值判断GPU是否可用,返回True则具有能够使用的GPU。
torch.cuda.device_count()可以获得能够使用的GPU数量。
nvidia-smi 查看GPU配置信息。

把数据从内存转移到GPU,一般针对张量(我们需要的数据)和模型。 对张量(类型为FloatTensor或者是LongTensor等),一律直接使用方法.to(device)或.cuda()即可。
device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)
#或device = torch.device(“cuda:0”)
device1 = torch.device(“cuda:1”)
for batch_idx, (img, label) in enumerate(train_loader):
img=img.to(device)
label=label.to(device)

对于模型来说,也是同样的方式,使用.to(device)或.cuda来将网络放到GPU显存。
#实例化网络
model = Net()
model.to(device) #使用序号为0的GPU
#或model.to(device1) #使用序号为1的GPU

多GPU加速
这里我们介绍单主机多GPUs的情况,单机多GPUs主要采用的DataParallel函数,而不是DistributedParallel,后者一般用于多主机多GPUs,当然也可用于单机多GPU。
使用多卡训练的方式有很多,当然前提是我们的设备中存在两个及以上的GPU。
使用时直接用model传入torch.nn.DataParallel函数即可,如下代码:

#对模型 net = torch.nn.DataParallel(model)

这时,默认所有存在的显卡都会被使用。
如果你的电脑有很多显卡,但只想利用其中一部分,如只使用编号为0、1、3、4的四个GPU,那

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值