多卡训练模板

# main.py
import torch
import argparse
import torch.distributed as dist
 
from apex.parallel import DistributedDataParallel
 
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()
 
dist.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)
 
train_dataset = ...
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
 
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)
 
model = ...
model.cuda()
model, optimizer = amp.initialize(model, optimizer)
model = DistributedDataParallel(model)
 
optimizer = optim.SGD(model.parameters())
 
for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda()
      target = target.cuda()
      ...
      output = model(images)
      loss = criterion(output, target)
      optimizer.zero_grad()
      with amp.scale_loss(loss, optimizer) as scaled_loss:
         scaled_loss.backward()
      optimizer.step()

多卡训练模板，命令行启动方法

代码中设置os.environ["CUDA_VISIBLE_DEVICES"]='0,1'

命令行：torchrun --nproc_per_node=4 main.py

如果要调用model中的init方法，要把

model._init_hidden_state(num_sample)改成

model.module._init_hidden_state(num_sample)

这是用apex库调用多卡训练的方法