Torch 1.9.1 DDP 并行优化与多模块调用问题

最新推荐文章于 2024-07-05 17:49:04 发布

williamyi96

最新推荐文章于 2024-07-05 17:49:04 发布

阅读量1.1k

点赞数 1

文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/williamyi96/article/details/120889727

版权

DDP 基础实现

由于 DataParallel (DP) 采取的是多线程并行，出于其特性，会造成通信瓶颈（GIL 限制），因此更高效的方式是使用 DistributedDataParallel 实现更高效的 GPU 使用。DDP 相关基础实现参考此处文章，亲测可以使用。

注：目前使用 1.7+ pytorch nccl 初始化 DDP 会报错，亲测也有该问题，因此建议使用 ‘gloo’，虽然相对速度可能较慢。

多模块调用问题

假设考虑如下问题：

import torch

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

import torch.optim as optim
import torch.nn as nn

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1, type=int)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank

torch