Pytorch下使用多GPU训练模型

2022-10-27 12:48:59

注:本文针对单个服务器上多块GPU的使用,不是多服务器多GPU的使用。
当你的电脑中有多块GPU时,可以增加batch_size来加快训练速度等。

涉及的代码

import torch

  1. 设置可以使用的GPU编号
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = '0,3' # 使用第0和第3编号的GPU
  2. 判断你的电脑或者pytorch是否能使用GPU来加速模型的训练
    torch.cuda.is_available() -->model = model.cuda()
    或者device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -->model = model.to(device)
  3. 判断你的电脑有几块GPU
    torch.cuda.device_count()
  4. torch.nn.DataParallel启用多GPU并行计算
    torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
    说明:
    该函数实现了在module级别上的数据并行使用,注意batch size要大于GPU的数量。
    参数 :
    module:需要多GPU训练的模型
    device_ids&
  • 作者:Micheal超
  • 原文链接:https://blog.csdn.net/qq_42887760/article/details/109471235
    更新时间:2022-10-27 12:48:59