pytorch-指定多gpu训练

最新推荐文章于 2025-07-13 21:22:13 发布

Iriving_shu

最新推荐文章于 2025-07-13 21:22:13 发布

阅读量9.4k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： pytorch踩坑之路

本文链接：https://blog.youkuaiyun.com/Iriving_shu/article/details/89256905

pytorch踩坑之路专栏收录该内容

1 篇文章

订阅专栏

本文详细介绍了在PyTorch中如何配置多GPU环境，包括两种主要方法：一是通过解析命令行参数并直接指定GPU ID；二是利用环境变量设置可见GPU。文章深入探讨了如何避免运行时错误，确保模型能在指定的GPU上正确运行。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

如何指定到 4,5,6,7号卡

方法一：

步骤：

    # gpu init
    multi_gpus = False
    if ',' in args.gpus:
        gpu_ids = [int(id) for id in args.gpus.split(',')]
        multi_gpus = True
    else:
        gpu_ids = [int(args.gpus)]

    device = torch.device('cuda:{}'.format(gpu_ids[0]) if torch.cuda.is_available() else 'cpu')
    # 此处如果使用 下面一行代码，则会报错，RuntimeError: all tensors must be on devices[0]
    # 默认情况下 device 为0，因此需要指定 device id.
   #  device = torch.device('cuda'  if torch.cuda.is_available() else 'cpu')
    ....
      if multi_gpus:
       net = DataParallel(net, device_ids=gpu_ids).to(device)
       margin = DataParallel(margin, device_ids=gpu_ids).to(device)
     else:
        net = net.to(device)
        margin = margin.to(device)

方法二

    if len(args.gpus.split(',')) > 1:
        multi_gpus = True
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus

    #device = torch.device('cuda:{}'.format(gpu_ids[0]) if torch.cuda.is_available() else 'cpu')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')