實(shí)踐教程｜PyTorch訓(xùn)練加速技巧

由于最近的程序?qū)λ俣纫蟊容^高，想要快速出結(jié)果，因此特地學(xué)習(xí)了一下混合精度運(yùn)算和并行化操作，由于已經(jīng)有很多的文章介紹相關(guān)的原理，因此本篇只講述如何應(yīng)用torch實(shí)現(xiàn)混合精度運(yùn)算、數(shù)據(jù)并行和分布式運(yùn)算，不具體介紹原理。

混合精度

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

##?導(dǎo)入amp工具包?
from?torch.cuda.amp?import?autocast,?GradScaler

model.train()

##?對梯度進(jìn)行scale來加快模型收斂，
##?因?yàn)閒loat16梯度容易出現(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選擇精度來提升訓(xùn)練性能而不降低模型準(zhǔn)確度
????with?autocast():
????#?compute?output
????????output?=?model(images)

????????loss?=?criterion(output,?target)

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

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

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

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

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

nn.DataParallel使用混合精度運(yùn)算

nn.DataParallel對模型進(jìn)行混合精度運(yùn)算需要進(jìn)行一些特殊的配置，不然模型是無法實(shí)現(xiàn)數(shù)據(jù)并行化的。autocast 設(shè)計為 “thread local” 的，所以只在 main thread 上設(shè) 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 的里面設(shè) autocast 區(qū)域： python MyModel(nn.Module): ... def forward(self, input): with autocast(): ... 在對forward函數(shù)進(jìn)行操作后，再在main thread中使用autocast ```python model = MyModel() dp_model = nn.DataParallel(model)

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

nn.DataParallel缺點(diǎn)

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

分布式運(yùn)算

nn.DistributedDataParallel：多進(jìn)程控制多 GPU，一起訓(xùn)練模型。

優(yōu)點(diǎn)

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

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

使用方式介紹

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

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

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

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

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

同步BatchNormalization層

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

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

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

同步模型初始化的隨機(jī)種子

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

總結(jié)

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

引用（不分前后）：

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