【深度學(xué)習(xí)】多卡訓(xùn)練__單機(jī)多GPU詳解（torch.nn.DataParallel、torch.distributed）

1. 介紹

多GPU訓(xùn)練能夠加快模型的訓(xùn)練速度，而且在單卡上不能訓(xùn)練的模型可以使用多個(gè)小卡達(dá)到訓(xùn)練的目的。
多GPU訓(xùn)練可以分為單機(jī)多卡和多機(jī)多卡這兩種，后面一種也就是分布式訓(xùn)練——訓(xùn)練方式比較麻煩，而且要關(guān)注的性能問(wèn)題也有很多，據(jù)網(wǎng)上的資料有人建議能單機(jī)訓(xùn)練最好單機(jī)訓(xùn)練，不要使用多機(jī)訓(xùn)練。本文主要對(duì)單機(jī)多卡訓(xùn)練的實(shí)現(xiàn)展開(kāi)說(shuō)明。

2. 單機(jī)多GPUの方法

2.1 方法1：torch.nn.DataParallel

這是最簡(jiǎn)單最直接的方法，代碼中只需要一句代碼就可以完成單卡多GPU訓(xùn)練了。其他的代碼和單卡單GPU訓(xùn)練是一樣的。

2.1.1 API

import torch
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

參數(shù)：

• module：即模型，此處注意，雖然輸入數(shù)據(jù)被均分到不同gpu上，但每個(gè)gpu上都要拷貝一份模型。
• device_ids：即參與訓(xùn)練的gpu列表，例如三塊卡， device_ids = [0，1，2]。
• output_device：指定輸出gpu，一般省略。在省略的情況下，默認(rèn)為第一塊卡，即索引為0的卡。此處有一個(gè)問(wèn)題，輸入計(jì)算是被幾塊卡均分的，但輸出loss的計(jì)算是由這一張卡獨(dú)自承擔(dān)的，這就造成這張卡所承受的計(jì)算量要大于其他參與訓(xùn)練的卡。
• dim：其表示tensors被分散的維度，默認(rèn)是0，nn.DataParallel將在dim0（批處理維度）中對(duì)數(shù)據(jù)進(jìn)行分塊，并將每個(gè)分塊發(fā)送到相應(yīng)的設(shè)備。

2.1.2 特點(diǎn)

• 優(yōu)點(diǎn)：特別簡(jiǎn)單，實(shí)現(xiàn)起來(lái)容易；
• 缺點(diǎn)：也很明顯，就是每個(gè)batch中，模型的權(quán)重都是在單一的線(xiàn)程上算出來(lái)的，然后分發(fā)到多個(gè)GPU上，這里就有一個(gè)GPU通信瓶頸，使得GPU的利用率不是很高，模型訓(xùn)練的速度也不快。

2.1.3 例子與解釋

import torch
net = torch.nn.Linear(100,1)
print(net)
print('---------------------')
net = torch.nn.DataParallel(net, device_ids=[0,3])
print(net)

'''輸出'''
Linear(in_features=10, out_features=1, bias=True)
---------------------
DataParallel(
  (module): Linear(in_features=10, out_features=1, bias=True)
)

nn.DataParallel是怎么做的：

• 首先在前向過(guò)程中，
  • 你的輸入數(shù)據(jù)會(huì)被劃分成多個(gè)子部分（以下稱(chēng)為副本）送到不同的device中進(jìn)行計(jì)算，
  • 而你的模型module是在每個(gè)device上進(jìn)行復(fù)制一份，也就是說(shuō)，輸入的batch是會(huì)被平均分到每個(gè)device中去，但是你的模型module是要拷貝到每個(gè)devide中去的，每個(gè)模型module只需要處理每個(gè)副本即可，當(dāng)然你要保證你的batch size大于你的gpu個(gè)數(shù)。
• 然后在反向傳播過(guò)程中，每個(gè)副本的梯度被累加到原始模塊中。

概括來(lái)說(shuō)就是：DataParallel會(huì)自動(dòng)幫我們將數(shù)據(jù)切分 load 到相應(yīng) GPU，將模型復(fù)制到相應(yīng) GPU，進(jìn)行正向傳播計(jì)算梯度并匯總。

還有一句話(huà)，官網(wǎng)中是這樣描述的：

The parallelized module must have its parameters and buffers on device_ids[0] before running this DataParallel module.

意思是：在運(yùn)行此DataParallel模塊之前，并行化模塊必須在device_ids [0]上具有其參數(shù)和緩沖區(qū)。在執(zhí)行DataParallel之前，會(huì)首先把其模型的參數(shù)放在device_ids[0]上，一看好像也沒(méi)有什么毛病，其實(shí)有個(gè)小坑。

• 舉個(gè)例子，服務(wù)器是八卡的服務(wù)器，剛好前面序號(hào)是0的卡被別人占用著，于是你只能用其他的卡來(lái)，比如你用2和3號(hào)卡，如果你直接指定device_ids=[2, 3]的話(huà)會(huì)出現(xiàn)模型初始化錯(cuò)誤，類(lèi)似于module沒(méi)有復(fù)制到在device_ids[0]上去。那么你需要在運(yùn)行train之前需要添加如下兩句話(huà)指定程序可見(jiàn)的devices，如下：

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

當(dāng)你添加這兩行代碼后，那么device_ids[0]默認(rèn)的就是第2號(hào)卡，你的模型也會(huì)初始化在第2號(hào)卡上了，而不會(huì)占用第0號(hào)卡了。

• 也就是設(shè)置上面兩行代碼后，那么對(duì)這個(gè)程序而言可見(jiàn)的只有2和3號(hào)卡，和其他的卡沒(méi)有關(guān)系，這是物理上的號(hào)卡，邏輯上來(lái)說(shuō)其實(shí)是對(duì)應(yīng)0和1號(hào)卡，即device_ids[0]對(duì)應(yīng)的就是第2號(hào)卡，device_ids[1]對(duì)應(yīng)的就是第3號(hào)卡。

當(dāng)然你要保證上面這兩行代碼需要定義在下面這兩行代碼之前，一般放在train.py中import一些package之后：

device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

優(yōu)化器同樣可以使用nn.DataParallel，如下兩行代碼：

optimizer = torch.optim.SGD(net.parameters(), lr=lr)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

2.1.4 說(shuō)明

1）多GPU計(jì)算減少了程序運(yùn)行的時(shí)間？
很多時(shí)候發(fā)現(xiàn)在進(jìn)行多GPU運(yùn)算時(shí)，程序花費(fèi)的時(shí)間反而更多了，這其實(shí)是因?yàn)槟愕腷atch_size太小了，因?yàn)閠orch.nn.DataParallel()這個(gè)函數(shù)是將每個(gè)batch的數(shù)據(jù)平均拆開(kāi)分配到多個(gè)GPU上進(jìn)行計(jì)算，計(jì)算完再返回來(lái)合并。這導(dǎo)致GPU之間的開(kāi)關(guān)和通訊過(guò)程占了大部分的時(shí)間開(kāi)銷(xiāo)。

我們可以使用 watch -n 1 nvidia-smi 這個(gè)命令來(lái)查看每1s各個(gè)GPU的運(yùn)行情況，如果發(fā)現(xiàn)每個(gè)GPU的占用率均低于50%，基本可以肯定你使用多GPU計(jì)算所花的時(shí)間要比單GPU計(jì)算花的時(shí)間更長(zhǎng)了。

2）如何保存和加載多GPU網(wǎng)絡(luò)？

• 如何來(lái)保存和加載多GPU網(wǎng)絡(luò)，它與普通網(wǎng)絡(luò)有一點(diǎn)細(xì)微的不同：

import torch
net = torch.nn.Linear(10,1)  # 先構(gòu)造一個(gè)網(wǎng)絡(luò)
net = torch.nn.DataParallel(net, device_ids=[0,3])  #包裹起來(lái)
torch.save(net.module.state_dict(), './networks/multiGPU.h5') #保存網(wǎng)絡(luò)

# 加載網(wǎng)絡(luò)
new_net = torch.nn.Linear(10,1)
new_net.load_state_dict(torch.load("./networks/multiGPU.h5"))

因?yàn)镈ataParallel實(shí)際上是一個(gè)nn.Module，所以我們?cè)诒４鏁r(shí)需要多調(diào)用了一個(gè)net.module，模型和優(yōu)化器都需要使用net.module來(lái)得到實(shí)際的模型和優(yōu)化器。

3）為什么第一塊卡的顯存會(huì)占用的更多一些？
最后一個(gè)參數(shù)output_device一般情況下是省略不寫(xiě)的，那么默認(rèn)就是在device_ids[0]，也就是第一塊卡上，也就解釋了為什么第一塊卡的顯存會(huì)占用的比其他卡要更多一些。

• 也就是當(dāng)你調(diào)用nn.DataParallel的時(shí)候，只是在你的input數(shù)據(jù)是并行的，但是你的output loss卻不是這樣的，每次都會(huì)在第一塊GPU相加計(jì)算，這就造成了第一塊GPU的負(fù)載遠(yuǎn)遠(yuǎn)大于剩余其他的顯卡。

4）直接使用nn.DataParallel的時(shí)候，訓(xùn)練采用多卡訓(xùn)練，會(huì)出現(xiàn)一個(gè)warning？

UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; 
will instead unsqueeze and return a vector.

說(shuō)明：

• 每張卡上的loss都是要匯總到第0張卡上求梯度，更新好以后把權(quán)重分發(fā)到其余卡。但是為什么會(huì)出現(xiàn)這個(gè)warning，這其實(shí)和nn.DataParallel中最后一個(gè)參數(shù)dim有關(guān)，
  • 其表示tensors被分散的維度，默認(rèn)是0，nn.DataParallel將在dim0（批處理維度）中對(duì)數(shù)據(jù)進(jìn)行分塊，并將每個(gè)分塊發(fā)送到相應(yīng)的設(shè)備。
• 單卡的沒(méi)有這個(gè)warning，多卡的時(shí)候采用nn.DataParallel訓(xùn)練會(huì)出現(xiàn)這個(gè)warning，由于計(jì)算loss的時(shí)候是分別在多卡計(jì)算的，那么返回的也就是多個(gè)loss，你使用了多少個(gè)gpu，就會(huì)返回多少個(gè)loss。（有人建議DataParallel類(lèi)應(yīng)該有reduce和size_average參數(shù)，比如用于聚合輸出的不同loss函數(shù)，最終返回一個(gè)向量，有多少個(gè)gpu，返回的向量就有幾維。）

關(guān)于這個(gè)問(wèn)題在pytorch官網(wǎng)的issues上有過(guò)討論，下面簡(jiǎn)單摘出一些：

• 有人提出求loss平均的方式會(huì)在不同數(shù)量的gpu上訓(xùn)練會(huì)以微妙的方式影響結(jié)果。模塊返回該batch中所有損失的平均值，如果在4個(gè)gpu上運(yùn)行，將返回4個(gè)平均值的向量。然后取這個(gè)向量的平均值。但是，如果在3個(gè)GPU或單個(gè)GPU上運(yùn)行，這將不是同一個(gè)數(shù)字，因?yàn)槊總€(gè)GPU處理的batch size不同！
• 舉個(gè)簡(jiǎn)單的例子（就直接摘原文出來(lái)）：

A batch of 3 would be calculated on a single GPU and results 
would be [0.3, 0.2, 0.8] and model that returns the loss would return 0.43.

If cast to DataParallel, and calculated on 2 GPUs, [GPU1 - batch 0,1], [GPU2 - batch 2] 
- return values would be [0.25, 0.8] (0.25 is average between 0.2 and 0.3)
- taking the average loss of [0.25, 0.8] is now 0.525!

Calculating on 3 GPUs, one gets [0.3, 0.2, 0.8] as results and average is back to 0.43!

這么求平均loss確實(shí)有不合理的地方。那么有什么好的解決辦法呢，可以使用size_average=False，reduce=True作為參數(shù)。每個(gè)GPU上的損失將相加，但不除以GPU上的批大小。然后將所有平行損耗相加，除以整批的大小，那么不管幾塊GPU最終得到的平均loss都是一樣的。

pytorch貢獻(xiàn)者也實(shí)現(xiàn)了這個(gè)loss求平均的功能，即通過(guò)gather的方式來(lái)求loss平均：
https://github.com/pytorch/pytorch/pull/7973/commits/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be

如果它們?cè)谝粋€(gè)有2個(gè)GPU的系統(tǒng)上運(yùn)行，DP將采用多GPU路徑，調(diào)用gather并返回一個(gè)向量。如果運(yùn)行時(shí)有1個(gè)GPU可見(jiàn)，DP將采用順序路徑，完全忽略gather，因?yàn)檫@是不必要的，并返回一個(gè)標(biāo)量。

2.2 方法2：torch.nn.parallel.DistributedDataParallel

這種方法旨在緩解nn.DataParallel方法GPU使用效率低的問(wèn)題。

• 這方法會(huì)使得GPU的顯存分配更加平衡一點(diǎn)，
• 同時(shí)這個(gè)方法是多線(xiàn)程的，顯卡的利用效率自然也就高一點(diǎn)。

2.2.1 API

1）首先第一步就是要進(jìn)行init_process_group的初始化，聲明GPU的NCCL通信方式。

import torch
torch.distributed.init_process_group(backend='nccl')

2）其次，由于是多線(xiàn)程的，因此數(shù)據(jù)加載和模型加載也要做對(duì)應(yīng)的修改如下：

train_data = ReadDataSet('train.tsv', args, sentences_count = None)
train_sample = torch.utils.data.distributed.DistributedSampler(train_data)
train_loader = DataLoader(dataset=train_data, batch_size=args.batch_size, shuffle=(train_sample is None), sampler=train_sample)
 
model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], find_unused_parameters=True) #多進(jìn)程多GPU并行

2.2.2 注意事項(xiàng)

1）首先就是代碼使用bash腳本啟動(dòng)的時(shí)候是不一樣的，要像下面這么定義：

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 main_gpus.py \
			# 后面加一些你要傳入的參數(shù)

在 pyton關(guān)鍵字之前把可用顯卡號(hào)用它CUDA_VISIBLE_DEVICES=0,1來(lái)指定；同時(shí)python關(guān)鍵字之后-m torch.distributed.launch --nproc_per_node=2 來(lái)指定 分布式啟動(dòng)和采用的節(jié)點(diǎn)數(shù)，也就是有幾個(gè)顯卡也就用幾個(gè)節(jié)點(diǎn)。

2）其次就是日志信息的打印，在代碼中直接打印的話(huà)就會(huì)打印nproc_per_node次指定的輸出信息，這個(gè)時(shí)候就需要指定進(jìn)程號(hào)。

if (step+1)%200 == 0 and args.local_rank==0:
        print('Train Epoch[{}/{}],step[{}/{}],tra_acc{:.6f} %,loss:{:.6f}'.format(epoch,epochs,step,len(train_iter),two_pro_train_acc*100,two_pro_loss))

這樣就只會(huì)打印進(jìn)程為0的對(duì)應(yīng)各種信息。

3）再次就是loss和準(zhǔn)確率的合并，這里有多個(gè)線(xiàn)程，肯定就需要對(duì)一個(gè)batch中多個(gè)線(xiàn)程對(duì)應(yīng)的不同loss和準(zhǔn)確率進(jìn)行合并。實(shí)現(xiàn)方式如下：

def reduce_tensor(tensor: torch.Tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    rt /= dist.get_world_size()  # 總進(jìn)程數(shù)
    return rt

把各自的loss或者accuracy做分布式操作的加法，然后在求平均值。

4）最后，關(guān)于batch_size和lr的設(shè)置，這里一般可以采用batch_size = n*batch_size_base的方式；而lr = (1,n)*lr_base的方式。

2.2.3 主要代碼（可以參照改成自己的）

import torch
from torch import nn
from unet.unet_transfer import UNet16, UNetResNet
from pathlib import Path
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset, random_split
import torch.nn.functional as F
from torch.autograd import Variable
import shutil
from data_loader import ImgDataSet
import os
import argparse
import tqdm
import numpy as np
import scipy.ndimage as ndimage
import torch.distributed as dist

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def create_model(type ='vgg16'):
    if type == 'vgg16':
        print('create vgg16 model')
        model = UNet16(pretrained=True)
    elif type == 'resnet101':
        encoder_depth = 101
        num_classes = 1
        print('create resnet101 model')
        model = UNetResNet(encoder_depth=encoder_depth, num_classes=num_classes, pretrained=True)
    elif type == 'resnet34':
        encoder_depth = 34
        num_classes = 1
        print('create resnet34 model')
        model = UNetResNet(encoder_depth=encoder_depth, num_classes=num_classes, pretrained=True)
    else:
        assert False
    model.eval()
    return model

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

def find_latest_model_path(dir):
    model_paths = []
    epochs = []
    for path in Path(dir).glob('*.pt'):
        if 'epoch' not in path.stem:
            continue
        model_paths.append(path)
        parts = path.stem.split('_')
        epoch = int(parts[-1])
        epochs.append(epoch)

    if len(epochs) > 0:
        epochs = np.array(epochs)
        max_idx = np.argmax(epochs)
        return model_paths[max_idx]
    else:
        return None

def train(train_loader, model, criterion, optimizer, validation, args):
    latest_model_path = find_latest_model_path(args.model_dir)

    best_model_path = os.path.join(*[args.model_dir, 'model_best.pt'])

    if latest_model_path is not None:
        state = torch.load(latest_model_path)
        epoch = state['epoch']
        model.load_state_dict(state['model'])
        epoch = epoch

        #if latest model path does exist, best_model_path should exists as well
        assert Path(best_model_path).exists() == True, f'best model path {best_model_path} does not exist'
        #load the min loss so far
        best_state = torch.load(latest_model_path)
        min_val_los = best_state['valid_loss']

        print(f'Restored model at epoch {epoch}. Min validation loss so far is : {min_val_los}')
        epoch += 1
        print(f'Started training model from epoch {epoch}')
    else:
        print('Started training model from epoch 0')
        epoch = 0
        min_val_los = 9999

    valid_losses = []

    for epoch in range(epoch, args.n_epoch + 1):

        adjust_learning_rate(optimizer, epoch, args.lr)

        tq = tqdm.tqdm(total=(len(train_loader) * args.batch_size))
        tq.set_description(f'Epoch {epoch}')

        losses = AverageMeter()

        model.train()
        for i, (input, target) in enumerate(train_loader):
            two_pro_loss = 0
            input_var  = Variable(input).cuda(args.local_rank, non_blocking=True)
            target_var = Variable(target).cuda(args.local_rank, non_blocking=True)

            masks_pred = model(input_var)

            masks_probs_flat = masks_pred.view(-1)
            true_masks_flat  = target_var.view(-1)

            loss = criterion(masks_probs_flat, true_masks_flat)
            two_pro_loss += reduce_tensor(loss).item() # 有多個(gè)進(jìn)程，把進(jìn)程0和1的loss加起來(lái)平均
            losses.update(two_pro_loss)
            tq.set_postfix(loss='{:.5f}'.format(losses.avg))
            tq.update(args.batch_size)

            # compute gradient and do SGD step
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        valid_metrics = validation(model, valid_loader, criterion)
        valid_loss = valid_metrics['valid_loss']
        valid_losses.append(valid_loss)
        print(f'\tvalid_loss = {valid_loss:.5f}')
        tq.close()

        #save the model of the current epoch
        epoch_model_path = os.path.join(*[args.model_dir, f'model_epoch_{epoch}.pt'])
        torch.save({
            'model': model.state_dict(),
            'epoch': epoch,
            'valid_loss': valid_loss,
            'train_loss': losses.avg
        }, epoch_model_path)

        if valid_loss < min_val_los:
            min_val_los = valid_loss

            torch.save({
                'model': model.state_dict(),
                'epoch': epoch,
                'valid_loss': valid_loss,
                'train_loss': losses.avg
            }, best_model_path)


def reduce_tensor(tensor: torch.Tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    rt /= dist.get_world_size()  # 總進(jìn)程數(shù)
    return rt

def validate(model, val_loader, criterion):
    losses = AverageMeter()
    model.eval()
    with torch.no_grad():
        for i, (input, target) in enumerate(val_loader):
            two_pro_loss = 0
            input_var = Variable(input).cuda(args.local_rank, non_blocking=True)
            target_var = Variable(target).cuda(args.local_rank, non_blocking=True)

            output = model(input_var)
            loss = criterion(output, target_var)
            two_pro_loss += reduce_tensor(loss)
            losses.update(loss.item(), input_var.size(0))

    return {'valid_loss': losses.avg}

def save_check_point(state, is_best, file_name = 'checkpoint.pth.tar'):
    torch.save(state, file_name)
    if is_best:
        shutil.copy(file_name, 'model_best.pth.tar')

def calc_crack_pixel_weight(mask_dir):
    avg_w = 0.0
    n_files = 0
    for path in Path(mask_dir).glob('*.*'):
        n_files += 1
        m = ndimage.imread(path)
        ncrack = np.sum((m > 0)[:])
        w = float(ncrack)/(m.shape[0]*m.shape[1])
        avg_w = avg_w + (1-w)

    avg_w /= float(n_files)

    return avg_w / (1.0 - avg_w)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
    parser.add_argument('--n_epoch', default=10, type=int, metavar='N', help='number of total epochs to run')
    parser.add_argument('--lr', default=0.001, type=float, metavar='LR', help='initial learning rate')
    parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum')
    parser.add_argument('--print_freq', default=20, type=int, metavar='N', help='print frequency (default: 10)')
    parser.add_argument('--weight_decay', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)')
    parser.add_argument('--batch_size',  default=4, type=int,  help='weight decay (default: 1e-4)')
    parser.add_argument('--num_workers', default=4, type=int, help='num_workers')
    parser.add_argument('--data_dir',type=str, default='dataset', help='input dataset directory')
    parser.add_argument('--model_dir', type=str, default='model', help='output model directory')
    parser.add_argument('--model_type', type=str, required=False, default='vgg16', choices=['vgg16', 'resnet101', 'resnet34'])
    parser.add_argument('--local_rank', type=int, default=-1)
    args = parser.parse_args()

    torch.cuda.set_device(args.local_rank)

    # os.environ['LOCAL_RANK'] = -1

    dist.init_process_group(backend='nccl')

    os.makedirs(args.model_dir, exist_ok=True)

    DIR_IMG  = os.path.join(args.data_dir, 'images')
    DIR_MASK = os.path.join(args.data_dir, 'masks')

    img_names  = [path.name for path in Path(DIR_IMG).glob('*.jpg')]
    mask_names = [path.name for path in Path(DIR_MASK).glob('*.jpg')]

    print(f'total images = {len(img_names)}')

    model = create_model(args.model_type)

    optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    criterion = nn.BCEWithLogitsLoss().cuda()

    # ori
    channel_means = [0.485, 0.456, 0.406]
    channel_stds  = [0.229, 0.224, 0.225]

    # dam
    # channel_means = [0.595, 0.608, 0.604]
    # channel_stds =  [0.047, 0.047, 0.047]

    train_tfms = transforms.Compose([transforms.ToTensor(),
                                     transforms.Normalize(channel_means, channel_stds)])

    val_tfms = transforms.Compose([transforms.ToTensor(),
                                   transforms.Normalize(channel_means, channel_stds)])

    mask_tfms = transforms.Compose([transforms.ToTensor()])

    dataset = ImgDataSet(img_dir=DIR_IMG, img_fnames=img_names, img_transform=train_tfms, mask_dir=DIR_MASK, mask_fnames=mask_names, mask_transform=mask_tfms)
    train_size = int(0.85*len(dataset))
    valid_size = len(dataset) - train_size
    train_dataset, valid_dataset = random_split(dataset, [train_size, valid_size])

    train_sample = torch.utils.data.distributed.DistributedSampler(train_dataset)
    valid_sample = torch.utils.data.distributed.DistributedSampler(valid_dataset)

    train_loader = DataLoader(train_dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, sampler=train_sample)
    valid_loader = DataLoader(valid_dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, sampler=valid_sample)

    model.cuda(args.local_rank)

    train(train_loader, model, criterion, optimizer, validate, args)

?。。。。。?！你可能會(huì)運(yùn)行出錯(cuò)，記得一定要按照下面的去運(yùn)行代碼！?。。。。。。。?/strong>

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 main_gpus.py \
			# 后面加一些你要傳入的參數(shù)

for i, (inputs, labels) in enumerate(training_set):
  loss = model(inputs, labels)              # 計(jì)算loss
  optimizer.zero_grad()								      # 清空梯度
  loss.backward()                           # 反向計(jì)算梯度
  optimizer.step()                          # 更新參數(shù)

for i, (inputs, labels) in enumerate(training_set):
  loss = model(inputs, labels)                    # 計(jì)算loss
  loss = loss / accumulation_steps                # Normalize our loss (if averaged)
  loss.backward()                                 # 反向計(jì)算梯度，累加到之前梯度上
  if (i+1) % accumulation_steps == 0:             
      optimizer.step()                            # 更新參數(shù)
      model.zero_grad()                           # 清空梯度

單機(jī)多卡
--accumulation_steps 2
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 686.4237
tra_acc73.325002 %,dev_acc76.400002 %,best_dev_acc76.400002 %
*******************************************************************************
 
--accumulation_steps 5
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 578.8834
tra_acc73.329997 %,dev_acc75.500000 %,best_dev_acc76.100006 %
*******************************************************************************
 
--accumulation_steps 10
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 579.5692
tra_acc71.015000 %,dev_acc75.400402 %,best_dev_acc77.300002 %
*******************************************************************************
 
--accumulation_steps 20
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 613.6300s
tra_acc64.775002 %,dev_acc78.199997 %,best_dev_acc78.199997 %
*******************************************************************************
 
--accumulation_steps 20
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 580.7058
tra_acc64.754999 %,dev_acc77.400002 %,best_dev_acc77.400002 %
*******************************************************************************
 
--accumulation_steps 50
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 621.0073
tra_acc53.034997 %,dev_acc71.900002 %,best_dev_acc71.900002 %
*******************************************************************************
 
--accumulation_steps 80
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 568.5933
tra_acc43.325001 %,dev_acc67.199997 %,best_dev_acc67.199997 %
*******************************************************************************
 
--accumulation_steps 80
roberta_large textclassification task train and dev 2 epochs with grad accmulation time is 575.0746
tra_acc44.005001 %,dev_acc67.500000 %,best_dev_acc67.500000 %
*******************************************************************************
 
--accumulation_steps 0
roberta_large textclassification task train and dev 2 epochs time is 718.4363s
tra_acc74.285001 %,dev_acc73.199997 %,best_dev_acc73.199997 %
*******************************************************************************
 
--accumulation_steps 0
roberta_large textclassification task train and dev 2 epochs time is 694.9744
tra_acc74.559999 %,dev_acc74.000000 %,best_dev_acc74.000000 %
 
 
 
單卡單GPU
*******************************************************************************
trian and eval model time is 1023.3577s
tra_acc64.715000 %,dev_acc71.400000 %,best_dev_acc71.400000 %
 
*******************************************************************************
trian and eval model time is 1034.7063
tra_acc72.760000 %,dev_acc74.300000 %,best_dev_acc74.300000 %
 
 
 
*******************************************************************************
結(jié)論：
單卡3090耗時(shí)：1029s
 
雙卡3090耗時(shí)：707s——提升：45.5%
 
雙卡3090+梯度累加耗時(shí)： 580s——提升77.4%，21.9%