谈起gpu id的设置,通常来说我们使用如下代码即可设置:
device = torch.device("cuda:id")
如果你有两张卡,那么上述代码中id写入你想使用的设备号即可,例如你有两张卡,gpu0和gpu1,那么如果你想使用gpu0,则写法为:
device = torch.device("cuda:0")
使用其余gpu同理。
但重点来了,假设你使用的是封装好的代码,例如transformer中的Trainer、TrainingArguments等,它默认的是多卡并行运行,单机多卡方式,而在你的设备上又只能使用一张卡时,这种封装好的代码修改起来就不那么的友好了。
案例:服务器上有两个设备gpu0、gpu1,分配给你的只有gpu0,而你又使用的是单机多卡封装方式的代码,这时运行无论你怎么设置设备id,都是两张卡一起跑,因此提出如下解决方案:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
通过如上方式设置后,你就可以只使用gpu0进行模型训练了,上述代码一定切记要添加在主函数main中才生效,设置在其他代码中不起作用!
如果有3张卡,而你想使用gpu0与gpu2,那么设置方式如下:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2"
其他同理,切记一定要添加在主函数中!