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.cuda.set_device(local_rank)
dist.init_process_group(backend='gloo')
device = torch.device("cuda", local_rank)
class Net(torch.nn.Module):
def __init__(self):
super().__init__()
self.device = device
self.nn1 = nn.Linear(10, 20)
self.nn2 = nn.Linear(20, 10)
def forward(self, x):
mid = self.nn1(x)
model = self.nn2(mid)
return model
def middle(self, x):
return self.nn1(x)
model = Net().to(device)
model = DDP(model, device_ids=[local_rank], output_device=local_rank)
x = torch.randn(20, 10).to(local_rank)
outputs = model(x)
mid = model.middle(x)
labels = torch.randn(20, 10).to(local_rank)
loss_fn = nn.MSELoss()
loss_fn(outputs, labels).backward()
optimizer = optim.SGD(model.parameters(), lr=0.001)
optimizer.step()
如果没有 model_ddp = DDP(model, device_ids=[local_rank], output_device=local_rank)
的操作,上述程序是个简单例子。注意到此处 Net 模型有一个 intrinsic forwar() 函数 和 其他函数 middle(),在调用 mid = model.middle(x)
会报错,因为 model 转换为 DDP 形式之后只会执行 forward()。
考虑如下解决方式:
DDP.module.func()
查看 DDP 实现源码可以发现,DDP 可以通过模块引用从而调用 其他函数 middle()。将
mid = model.middle(x)
改为
mid = model.module.middle(x)
可以找到 middle() 函数,注意此处 mid 是非 DDP 形式的。
并行非并行分离
类似 StyleGAN2, 我们可以考虑创建并行 model_ddp() 模型和 model() 非并行模型,由于其他函数一般不参与模型学习,而是中间状态输出,因此这种情况下我们可以使用双模型。如下:
model = Net().to(device)
model_ddp = DDP(model, device_ids=[local_rank], output_device=local_rank)
forward() 条件引入 middle
由于问题是 DDP 只会执行 forward(), 因此我们可以考虑将 middle 进行条件引入:
def forward(self, x, run_middle=False):
if run_middle:
return self.middle(x)
mid = self.nn1(x)
model = self.nn2(mid)
return model
此为一种最为通用的策略。
上述三种策略根据需求可以结合使用。