当前位置:   article > 正文

pytorch指定使用多个GPU_pytorch调用多个gpu

pytorch调用多个gpu

方式一:指定使用所有GPU

  1. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  2. model = CreateModel()
  3. model= nn.DataParallel(model)
  4. model.to(device)

方式二:结合系统环境变量使用特定的GPU

  1. CUDA_VISIBLE_DEVICES=1,3
  2. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  3. model = CreateModel()
  4. model= nn.DataParallel(model)
  5. model.to(device)

方式三:直接指定特定的GPU

  1. device = torch.device("cuda:1,3" if torch.cuda.is_available() else "cpu") ## specify the GPU id's, GPU id's start from 0.
  2. model = CreateModel()
  3. model= nn.DataParallel(model,device_ids = [1, 3])
  4. model.to(device)

以上所有情况中,后面使用的数据都要映射到对应的GPU设备。

如果x,y是数据:

  1. X.to(device)
  2. y.to(device)

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/Monodyee/article/detail/708110
推荐阅读
相关标签
  

闽ICP备14008679号