Pytorch 使用torch.nn.parallel.DistributedDataParallel多GPU分布式训练model

最新推荐文章于 2023-05-11 22:50:33 发布

努力沉淀

最新推荐文章于 2023-05-11 22:50:33 发布

阅读量429

点赞数

文章标签： pytorch 深度学习人工智能

本文链接：https://blog.csdn.net/qq_39129717/article/details/124988426

版权

使用torch.nn.parallel.DistributedDataParallel必须要使用终端命令进行模型训练，命令模板为：

CUDA_VISIBLE_DEVICES=1,2,3  python -m torch.distributed.launch --nproc_per_node=3  train.py

这里 CUDA_VISIBLE_DEVICES=1,2,3表示你需要使用的gpu标号，--nproc_per_node=？直接填写gpu数量，train.py为要运行的文件

必须要添加一个参数 --locak_rank，可以不在终端命令中指示，命令中的torch.distributed.launch

import argparse

parser = argparse.ArgumentParser(description='model description')
parser.add_argument("--local_rank", type=int)

torch.distributed.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)

data = YourData_Class()

train_sampler = torch.utils.data.distributed.DistributedSampler(data)
dataloader = torch.utils.data.DataLoader(
        dataset=data,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=16,
        pin_memory=True,
        sampler=train_sampler)