对于多张GPU卡,如何选择指定的卡进行训练
本文最后更新于104 天前,其中的信息可能已经过时,如有错误请发送邮件到1729915388@qq.com

1. 利用CUDA_VISIBLE_DEVICES设置可用显卡

放在网络net开始训练的代码之前

os.environ['CUDA_VISIBLE_DEVICES'] = gpu_ids 

需要放在import torch之前?否则没有效果。

该行代码的意思是只选择指定标号GPU进行使用,执行完该行代码之后python环境无法检测到指定GPU之外的其他GPU

2 .cuda()方法和torch.cuda.set_device()

还可以使用.cuda()[包括model.cuda()/loss.cuda()/tensor.cuda()]方法和torch.cuda.set_device()来把模型和数据加载到对应的gpu上。

(1) .cuda()

以model.cuda()为例,加载方法为:

model.cuda(gpu_id) # gpu_id为int类型变量,只能指定一张显卡
model.cuda('cuda:'+str(gpu_ids)) #输入参数为str类型,可指定多张显卡
model.cuda('cuda:1,2') #指定多张显卡的一个示例

(2) torch.cuda.set_device()

使用torch.cuda.set_device()可以更方便地将模型和数据加载到对应GPU上, 直接定义模型之前加入一行代码即可

torch.cuda.set_device(gpu_id) #单卡
torch.cuda.set_device('cuda:'+str(gpu_ids)) #可指定多卡

多卡数据并行torch.nn.DataParallel

多卡数据并行一般使用

torch.nn.DataParallel(model,device_ids)
# self.net = torch.nn.DataParallel(self.net, device_ids=list(range(torch.cuda.device_count())))

其中model是需要运行的模型,device_ids指定部署模型的显卡,数据类型是list

device_ids中的第一个GPU(即device_ids[0])和model.cuda()或torch.cuda.set_device()中的第一个GPU序号应保持一致,否则会报错。此外如果两者的第一个GPU序号都不是0,比如设置为:

model=torch.nn.DataParallel(model,device_ids=[2,3])
model.cuda(2)

那么程序可以在GPU2和GPU3上正常运行,但是还会占用GPU0的一部分显存(大约500M左右),这是由于pytorch本身的bug导致的(截止1.4.0,没有修复这个bug)。

device_ids的默认值是使用可见的GPU,不设置model.cuda()或torch.cuda.set_device()等效于设置了model.cuda(0)

文末附加内容

评论

  1. 博主
    Windows Chrome
    3 月前
    2024-4-16 15:24:27

    os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘1’
    os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘1,2’

    来自中国

发送评论 编辑评论


				
|´・ω・)ノ
ヾ(≧∇≦*)ゝ
(☆ω☆)
(╯‵□′)╯︵┴─┴
 ̄﹃ ̄
(/ω\)
∠( ᐛ 」∠)_
(๑•̀ㅁ•́ฅ)
→_→
୧(๑•̀⌄•́๑)૭
٩(ˊᗜˋ*)و
(ノ°ο°)ノ
(´இ皿இ`)
⌇●﹏●⌇
(ฅ´ω`ฅ)
(╯°A°)╯︵○○○
φ( ̄∇ ̄o)
ヾ(´・ ・`。)ノ"
( ง ᵒ̌皿ᵒ̌)ง⁼³₌₃
(ó﹏ò。)
Σ(っ °Д °;)っ
( ,,´・ω・)ノ"(´っω・`。)
╮(╯▽╰)╭
o(*////▽////*)q
>﹏<
( ๑´•ω•) "(ㆆᴗㆆ)
😂
😀
😅
😊
🙂
🙃
😌
😍
😘
😜
😝
😏
😒
🙄
😳
😡
😔
😫
😱
😭
💩
👻
🙌
🖕
👍
👫
👬
👭
🌚
🌝
🙈
💊
😶
🙏
🍦
🍉
😣
Source: github.com/k4yt3x/flowerhd
颜文字
Emoji
小恐龙
花!
上一篇
下一篇