Python-GPU运行指定

目录

GPU指定总结

1. os.environ[‘CUDA_VISIBLE_DEVICES’]=‘x’

2.torch.cuda.set_device(x)

3.命令行运行时指定

4. torch.nn.DataParallel

5. export 环境变量


转载 

在使用pytorch的时候利用下面的语句指定GPU为仅为"6",但是用nvidia-smi查看GPU使用时,仍默认为"0"号

import pytorch
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '6'


解决方案:
将上述语句放到当前这个python文件的最开头,即import torch 之前

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
import torch

GPU指定总结


1. os.environ[‘CUDA_VISIBLE_DEVICES’]=‘x’


如上述,可以在运行的 .py 开头加入

os.environ['CUDA_VISIBLE_DEVICES'] = 'x'

优点:可以指定多卡;强制设置程序可见某几块板子,对于其他不可见的板子就完全没有影响,
缺点:不方便,只能放在import torch之前,即程序最开头指定;由于很多时候我们会在.py文件里面写很多function,然后写一个main函数入口运行整个文件,这样一来import torch不得不在os.environ['CUDA_VISIBLE_DEVICES'] = 'x'之前,这种情况下,每次想要换卡的时候,还需要点开.py文件手动修改,不能通过外部传参修改

2.torch.cuda.set_device(x)

在构建网络之前,调用torch的set_device指定

torch.cuda.set_device(x)
优点:比较灵活,可以在程序不同地方指定,适用于需要在一个程序里,先后将多个模型分别放到多块板子上的情况;可以传参换卡
缺点:只能指定一块板子


3.命令行运行时指定


在命令行或者.sh文件里,可以在运行.py文件前直接指定:

CUDA_VISIBLE_DEVICES=x python  xxx.py
优点:其实就是类似于os.environ['CUDA_VISIBLE_DEVICES']='x',只不过它用起来更加方便
缺点:还是只能指定一块,而且不能再程序运行中途换卡,适用于一次run,全程只需要占单卡的.py文件

4. torch.nn.DataParallel


当某些情况下需要用到并行训练的时候,可以用下列语句指定多卡

torch.nn.DataParallel(module, device_ids=[x1,x2,x3,x4,x5,x6,x7])
x1,x2不仅可以是 int. . 还可以是 torch.device

5. export 环境变量


最推荐的是这种,可以在shell里面export CUDA的可见环境变量

gpu=$1
echo "export CUDA_VISIBLE_DEVICES=${gpu}"
export CUDA_VISIBLE_DEVICES=${gpu}

python main.py  # 你要运行的python进程

比方说上述这几条cmd是写在一个shell:run.sh里面的,那么只要

sh run.sh 1,2

那么这一整个sh里面的python进程都只能可视1,2块板子
相比于其他方法,这种指定方式是最便捷的,也可以指定多张GPU.

其他注意点
当有小伙伴想要在程序里混合使用上述指定GPU的语句时,需要注意:

os.environ['CUDA_VISIBLE_DEVICES']是设定程序对哪几张卡可视,一般设定成功之后,接下来程序中任何有关卡号的指定都是相对的。
例如:os.environ['CUDA_VISIBLE_DEVICES'] = '1,2,3',生效之后,再设置torch.cuda.set_device(0),此时pytorch将会使用1号cuda.
其他类似…

torch.cuda.set_device(x)设置GPU之后,torch.cuda.device_count()返回的还是你所有的板子数目,这样一来后面的程序需要注意一下,可能会有问题
例如:Hugging face 或者google-rearch这种开源的预训练模型使用Demo,会在程序里count device数目n_gpu,然后会根据n_gpu调整是否使用torch.nn.DataParallel并行训练BERT,还会影响batch_size。如果你自己设置 torch.cuda.set_device(x),记得最好手动修改n_gpu为1,
其他类似…

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

静静喜欢大白

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值