來源丨https://zhuanlan.zhihu.com/p/360697168編輯丨極市平臺

由于最近的程序對速度要求比較高，想要快速出結果，因此特地學習了一下混合精度運算和并行化操作，由于已經有很多的文章介紹相關的原理，因此本篇只講述如何應用torch實現混合精度運算、數據并行和分布式運算，不具體介紹原理。

自動混合精度訓練（auto Mixed Precision，AMP）可以大幅度降低訓練的成本并提高訓練的速度。在此之前，自動混合精度運算是使用NVIDIA開發的Apex工具。從PyTorch1.6.0開始，PyTorch已經自帶了AMP模塊，因此接下來主要對PyTorch自帶的amp模塊進行簡單的使用介紹。

## 導入amp工具包 
from torch.cuda.amp import autocast, GradScaler

model.train()

## 對梯度進行scale來加快模型收斂，
## 因為float16梯度容易出現underflow（梯度過小）
scaler = GradScaler()

batch_size = train_loader.batch_size
num_batches = len(train_loader)
end = time.time()
for i, (images, target) in tqdm.tqdm(
    enumerate(train_loader), ascii=True, total=len(train_loader)
):
    # measure data loading time
    data_time.update(time.time() - end)
    optimizer.zero_grad()
    if args.gpu is not None:
        images = images.cuda(args.gpu, non_blocking=True)

    target = target.cuda(args.gpu, non_blocking=True)
    # 自動為GPU op選擇精度來提升訓練性能而不降低模型準確度
    with autocast():
    # compute output
        output = model(images)

        loss = criterion(output, target)

    scaler.scale(loss).backward()
    # optimizer.step()
    scaler.step(optimizer)
    scaler.update()

當服務器有單機有多卡的時候，為了實現模型的加速（可能由于一張GPU不夠），可以采用單機多卡對模型進行訓練。為了實現這個目的，我們必須想辦法讓一個模型可以分布在多個GPU上進行訓練。

PyTorch中，nn.DataParallel為我提供了一個簡單的接口，可以很簡單的實現對模型的并行化，我們只需要用nn.DataParallel對模型進行包裝，在設置一些參數，就可以很容易的實現模型的多卡并行。

# multigpu表示顯卡的號碼
multigpu = [0,1,2,3,4,5,6,7] 
# 設置主GPU,用來匯總模型的損失函數并且求導，對梯度進行更新
torch.cuda.set_device(args.multigpu[0])
# 模型的梯度全部匯總到gpu[0]上來
model = torch.nn.DataParallel(model, device_ids=args.multigpu).cuda(
        args.multigpu[0]
        )

nn.DataParallel對模型進行混合精度運算需要進行一些特殊的配置，不然模型是無法實現數據并行化的。autocast 設計為 “thread local” 的，所以只在 main thread 上設 autocast 區域是不 work 的。借鑒自（https://zhuanlan.zhihu.com/p/348554267） 這里先給出錯誤的操作：

model = MyModel() 
dp_model = nn.DataParallel(model)

with autocast():     # dp_model's internal threads won't autocast.
     #The main thread's autocast state has no effect.     
     output = dp_model(input)     # loss_fn still autocasts, but it's too late...
     loss = loss_fn(output)

解決的方法有兩種，下面分別介紹：1. 在模型模塊的forward函數中加入裝飾函數

MyModel(nn.Module):
    ...
    @autocast()
    def forward(self, input):
       ...

2. 另一個正確姿勢是在 forward 的里面設 autocast 區域： python MyModel(nn.Module): ... def forward(self, input): with autocast(): ... 在對forward函數進行操作后，再在main thread中使用autocast ```python model = MyModel() dp_model = nn.DataParallel(model)

with autocast(): output = dp_model(input) loss = loss_fn(output) ```

在每個訓練的batch中，nn.DataParallel模塊會把所有的loss全部反傳到gpu[0]上，幾個G的數據傳輸，loss的計算都需要在一張顯卡上完成，這樣子很容易造成顯卡的負載不均勻，經常可以看到gpu[0]的負載會明顯比其他的gpu高很多。此外，顯卡的數據傳輸速度會對模型的訓練速度造成很大的瓶頸，這顯然是不合理的。因此接下來我們將介紹，具體原理可以參考單機多卡操作(分布式DataParallel，混合精度，Horovod)（https://zhuanlan.zhihu.com/p/158375055）

nn.DistributedDataParallel：多進程控制多 GPU，一起訓練模型。

每個進程控制一塊GPU，可以保證模型的運算可以不受到顯卡之間通信的影響，并且可以使得每張顯卡的負載相對比較均勻。但是相對于單機單卡或者單機多卡(nn.DataParallel）來說，就有幾個問題

1. 同步不同GPU上的模型參數，特別是BatchNormalization 2. 告訴每個進程自己的位置，使用哪塊GPU，用args.local_rank參數指定 3. 每個進程在取數據的時候要確保拿到的是不同的數據（DistributedSampler）

啟動程序 由于博主目前也只是實踐了單機多卡操作，因此主要對單機多卡進行介紹。區別于平時簡單的運行python程序，我們需要使用PyTorch自帶的啟動器 torch.distributed.launch 來啟動程序。

# 其中CUDA_VISIBLE_DEVICES指定機器上顯卡的數量
# nproc_per_node程序進程的數量
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

配置主程序

parser.add_argument('--local_rank', type=int, default=0，help='node rank for distributed training')
# 配置local_rank參數，告訴每個進程自己的位置，要使用哪張GPU

初始化顯卡通信和參數獲取的方式

# 為這個進程指定GPU
torch.cuda.set_device(args.local_rank)
# 初始化GPU通信方式NCLL和參數的獲取方式，其中env表示環境變量
# PyTorch實現分布式運算是通過NCLL進行顯卡通信的
torch.distributed.init_process_group(
    backend='nccl',
    rank=args.local_rank
)

重新配置DataLoader

kwargs = {"num_workers": args.workers, "pin_memory": True} if use_cuda else {}

train_sampler = DistributedSampler(train_dataset)
self.train_loader = torch.utils.data.DataLoader(
            train_dataset, 
            batch_size=args.batch_size, 
            sampler=train_sampler,  
            **kwargs
        )

# 注意，由于使用了Sampler方法，dataloader中就不能加shuffle、drop_last等參數了
'''
PyTorch dataloader.py 192-197 代碼
        if batch_sampler is not None:
            # auto_collation with custom batch_sampler
            if batch_size != 1 or shuffle or sampler is not None or drop_last:
                raise ValueError('batch_sampler option is mutually exclusive '
                                 'with batch_size, shuffle, sampler, and '
                                 'drop_last')'''

pin_memory就是鎖頁內存，創建DataLoader時，設置pin_memory=True，則意味著生成的Tensor數據最開始是屬于內存中的鎖頁內存，這樣將內存的Tensor轉義到GPU的顯存就會更快一些。

模型的初始化

torch.cuda.set_device(args.local_rank)
device = torch.device('cuda', args.local_rank)
model.to(device)
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = torch.nn.parallel.DistributedDataParallel(
        model,
        device_ids=[args.local_rank],
        output_device=args.local_rank,
        find_unused_parameters=True,
        )
torch.backends.cudnn.benchmark=True 
# 將會讓程序在開始時花費一點額外時間，為整個網絡的每個卷積層搜索最適合它的卷積實現算法，進而實現網絡的加速
# DistributedDataParallel可以將不同GPU上求得的梯度進行匯總，實現對模型GPU的更新

DistributedDataParallel可以將不同GPU上求得的梯度進行匯總，實現對模型GPU的更新

對于比較消耗顯存的訓練任務時，往往單卡上的相對批量過小，影響模型的收斂效果。跨卡同步 Batch Normalization 可以使用全局的樣本進行歸一化，這樣相當于‘增大‘了批量大小，這樣訓練效果不再受到使用 GPU 數量的影響。參考自單機多卡操作(分布式DataParallel，混合精度，Horovod) 幸運的是，在近期的Pytorch版本中，PyTorch已經開始原生支持BatchNormalization層的同步。

• torch.nn.SyncBatchNorm
• torch.nn.SyncBatchNorm.convert_sync_batchnorm：將BatchNorm-alization層自動轉化為torch.nn.SyncBatchNorm實現不同GPU上的BatchNormalization層的同步

具體實現請參考模型的初始化部分代碼 python model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

同步模型初始化的隨機種子

目前還沒有嘗試過不同進程上使用不同隨機種子的狀況。為了保險起見，建議確保每個模型初始化的隨機種子相同，保證每個GPU進程上的模型是同步的。

站在巨人的肩膀上，對前段時間自學模型加速，踩了許多坑，最后游行都添上了，最后對一些具體的代碼進行了一些總結，其中也參考了許多其他的博客。希望能對大家有一些幫助。

引用（不分前后）：

1. PyTorch 21.單機多卡操作(分布式DataParallel，混合精度，Horovod)
2. PyTorch 源碼解讀之 torch.cuda.amp: 自動混合精度詳解
3. PyTorch的自動混合精度（AMP）
4. 訓練提速60%！只需5行代碼，PyTorch 1.6即將原生支持自動混合精度訓練
5. torch.backends.cudnn.benchmark ?!
6. 惡補了 Python 裝飾器的八種寫法，你隨便問~