服务器多GPU指定使用

最新推荐文章于 2024-07-09 15:30:57 发布

学习飞的菜鸟

最新推荐文章于 2024-07-09 15:30:57 发布

阅读量2.7k

点赞数 2

文章标签：服务器深度学习 pytorch

本文链接：https://blog.csdn.net/weixin_43429393/article/details/121337941

版权

os.environ["CUDA_VISIBLE_DEVICES"] = '2,3'
model = nn.DataParallel(model,device_ids=[0,1])
device = [0,1] if torch.cuda.device_count() > 1 else [0]

代码中 CUDA_VISIBLE_DEVICES 这里设定为’2,3’，代表程序可见的GPU编号是第三块和第四块，也就是说如果你想使用第三块和第四块GPU跑程序的话，你需要设置device id 为 [0,1]，因为此时程序只能看见第三块和第四块GPU，因此此时GPU3的device id 为0，GPU4的device id 为1。