pytorch指定使用多个GPU
短信预约 -IT技能 免费直播动态提醒
方式一:指定使用所有GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = CreateModel()model= nn.DataParallel(model)model.to(device)
方式二:结合系统环境变量使用特定的GPU
CUDA_VISIBLE_DEVICES=1,3device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = CreateModel()model= nn.DataParallel(model)model.to(device)
方式三:直接指定特定的GPU
device = torch.device("cuda:1,3" if torch.cuda.is_available() else "cpu") ## specify the GPU id's, GPU id's start from 0.model = CreateModel()model= nn.DataParallel(model,device_ids = [1, 3])model.to(device)
以上所有情况中,后面使用的数据都要映射到对应的GPU设备。
如果x,y是数据:
X.to(device)y.to(device)
来源地址:https://blog.csdn.net/RuanChengfeng/article/details/128172761
免责声明:
① 本站未注明“稿件来源”的信息均来自网络整理。其文字、图片和音视频稿件的所属权归原作者所有。本站收集整理出于非商业性的教育和科研之目的,并不意味着本站赞同其观点或证实其内容的真实性。仅作为临时的测试数据,供内部测试之用。本站并未授权任何人以任何方式主动获取本站任何信息。
② 本站未注明“稿件来源”的临时测试数据将在测试完成后最终做删除处理。有问题或投稿请发送至: 邮箱/279061341@qq.com QQ/279061341