pytorch 多卡并行训练

互联网 12 月前 0 3

设置 参数：

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', type=int, default=64, help='Batch size for training')
    parser.add_argument('--local_rank', type=int, default=0, help='Local rank of the process')
    parser.add_argument('--device', type=str, default='0', help='Local rank of the process')
    args = parser.parse_args()

    torch.distributed.init_process_group(backend='nccl', init_method='env://')

    local_rank = args.local_rank
    device = torch.device('cuda', local_rank)
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    epoch_num = 500
    batch_size = args.batch_size


    if torch.cuda.is_available():
        net.to(device)
    net = torch.nn.parallel.DistributedDataParallel(net, device_ids= 
    [local_rank],find_unused_parameters=True)

训练时参数：

# 在训练循环中使用本地 GPU 设备
for batch in dataloader:
    inputs, labels = batch
    inputs = inputs.to(device)
    labels = labels.to(device)
    # 在这里进行训练
    ...

调用 命令：

CUDA_VISIBLE_DEVICES=3,4 /data3/lbg/envs//aimet_3.8/bin/python3.8 -m torch.distributed.launch --master_port 49998 --nproc_per_node 2 train.py --device '' --batch_size 256