pytorch 分布式多卡

记录一下使用多卡训练时用的方法还有碰到的问题
使用dataparallel类相对比较简单，distributeddataparallel可以稍微提升效率，在单节点上面也可以跑。这里就只按照单节点写了。
更正式的代码参考pytorch官方给的imagenet分布式训练代码

代码：

import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 
import torch
from xxx.modeling.model import xxxNet
import cv2, os
import numpy as np
import argparse
from xxxnet.config import tcfg
import datasets.TEST.dataset as datareader
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

parser = argparse.ArgumentParser()
parser.add_argument('--epoch', type=int)
parser.add_argument('--batchsize', type=int)
parser.add_argument('--cuda', action='store_true')
parser.add_argument('--multi_gpu', action='store_true')
parser.add_argument('--data_dir', type=str)
parser.add_argument('--learningrate', type=float)
parser.add_argument('--local_rank', default=-1, type=int,
                            help='node rank for distributed training')
args = parser.parse_args()

if args.multi_gpu:
    torch.distributed.init_process_group(backend='nccl')
    local_rank = torch.distributed.get_rank()
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)


def adjust_learning_rate(optimizer, lr):
    """Sets the learning rate to the initial LR decayed by 10 every 2 epochs"""
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def train(args, model):
    
    print(device)

    dataset = datareader.Dataset(args.data_dir)
    dataloader = DataLoader(dataset, batch_size=args.batchsize, shuffle=True)
    if args.multi_gpu:
        dataloader = DataLoader(dataset, batch_size=args.batchsize,  sampler=DistributedSampler(dataset))
    else:
        dataloader = DataLoader(dataset, batch_size=args.batchsize, shuffle=True)

    if args.multi_gpu:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)



    optimizer = torch.optim.Adam(model.parameters(), lr=args.learningrate)
    criterion = torch.nn.CrossEntropyLoss()

    loss_pre = 1000000000000

    for i in range(args.epoch):
        model.train()
        iter = 0
        for img,  label in dataloader:
            if args.cuda:
                img, label = img.cuda(), label.cuda()
            pred_label = model(img)
           

            loss = criterion(pred_label, label)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if iter % 10 == 9:
                print(f'loss at iter {iter} is {loss}')
                if loss < loss_pre:
                    print('saving...')
                    loss_pre = loss
                    torch.save(model.state_dict(), './model_save/model.pkl')
            if loss < 3.23:
                adjust_learning_rate(optimizer, 1e-4)
                print(f"adjusted lr: {optimizer.param_groups[0]['lr']}")
            iter += 1



if __name__ == '__main__':

    model = xxxNet()
    if args.cuda:
        model = model.cuda()

    train(args, model)

执行命令：env CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 python -m torch.distributed.launch --nproc_per_node=6 tools/train_net.py --epoch 20 --batchsize 4 --cuda --data_dir ./out/ --learningrate 1e-3 --multi_gpu

其中，nproc_per_node是指一个节点上面进程数量，设置成和显卡数量一致。
代码中和单卡训练不一样的内容都在if args.multi_gpu:里面。
dataloader需要用distributeddatasampler包起来，并且在dataloader里面的shuffle需要置false否则会报错。
碰到的一些问题：

1. 报错：RuntimeError: Expected tensor for argument #1 ‘input’ to have the same device as tensor for argument #2 ‘weight’; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)
   原因
   if args.multi_gpu:
   torch.distributed.init_process_group(backend=‘nccl’)
   local_rank = torch.distributed.get_rank()
   torch.cuda.set_device(local_rank)
   device = torch.device(“cuda”, local_rank)
   这几行应该放在程序最开头！！！

2. 报错：error: unrecognized arguments: --local_rank=0
   添加参数 --local_rank
   每个进程分配一个 local_rank 参数，表示当前进程在当前主机上的编号。例如：rank=2, local_rank=0 表示第 3 个节点上的第 1 个进程。
   这个参数是torch.distributed.launch传递过来的，我们设置位置参数来接受，local_rank代表当前程序进程使用的GPU标号
   parser = argparse.ArgumentParser()
   parser.add_argument(’–local_rank’, default=-1, type=int,
   help=‘node rank for distributed training’)
   args = parser.parse_args()
   参考：https://www.cnblogs.com/yh-blog/p/12877922.html

3. 在这个代码中，参数指定batchsize指的是每张卡上面的数据量。如果想要在命令行中指定batchsize的总大小，需要计算一下，具体见pytorch官方的imagenet分布式代码

4. 在命令行中使用了torch.distributed.launch之后，可以不在代码中使用torch.multiprocessing.spawn

5. 保存checkpoint应该是要加上判断只保存一张卡的模型的，这里没写。

6. 如果有什么问题的话请在评论区告诉我，我也是第一次写分布式训练……