實踐教程｜PyTorch訓練加速技巧

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

混合精度

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

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

model.train()

##?對梯度進行scale來加快模型收斂，
##?因為float16梯度容易出現(xiàn)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()

數(shù)據(jù)并行

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

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

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

nn.DataParallel使用混合精度運算

nn.DataParallel對模型進行混合精度運算需要進行一些特殊的配置，不然模型是無法實現(xiàn)數(shù)據(jù)并行化的。autocast 設計為 “thread local” 的，所以只在 main thread 上設 autocast 區(qū)域是不 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函數(shù)中加入裝飾函數(shù)

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

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

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

nn.DataParallel缺點

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

分布式運算

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

優(yōu)點

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

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

使用方式介紹

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

# 其中CUDA_VISIBLE_DEVICES指定機器上顯卡的數(shù)量
# nproc_per_node程序進程的數(shù)量
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參數(shù)，告訴每個進程自己的位置，要使用哪張GPU

初始化顯卡通信和參數(shù)獲取的方式

#?為這個進程指定GPU
torch.cuda.set_device(args.local_rank)
#?初始化GPU通信方式NCLL和參數(shù)的獲取方式，其中env表示環(huán)境變量
#?PyTorch實現(xiàn)分布式運算是通過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等參數(shù)了
'''
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就是鎖頁內存，創(chuàng)建DataLoader時，設置pin_memory=True，則意味著生成的Tensor數(shù)據(jù)最開始是屬于內存中的鎖頁內存，這樣將內存的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?
#?將會讓程序在開始時花費一點額外時間，為整個網絡的每個卷積層搜索最適合它的卷積實現(xiàn)算法，進而實現(xiàn)網絡的加速
#?DistributedDataParallel可以將不同GPU上求得的梯度進行匯總，實現(xiàn)對模型GPU的更新

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

同步BatchNormalization層

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

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

具體實現(xiàn)請參考模型的初始化部分代碼 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 裝飾器的八種寫法，你隨便問~

如果覺得有用，就請分享到朋友圈吧！