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

來源丨h(huán)ttps://zhuanlan.zhihu.com/p/360697168編輯丨極市平臺

由于最近的程序?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()

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

PyTorch中，nn.DataParallel為我提供了一個(gè)簡單的接口，可以很簡單的實(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對模型進(jìn)行混合精度運(yùn)算需要進(jìn)行一些特殊的配置，不然模型是無法實(shí)現(xiàn)數(shù)據(jù)并行化的。autocast 設(shè)計(jì)為 “thread local” 的，所以只在 main thread 上設(shè) autocast 區(qū)域是不 work 的。借鑒自（https://zhuanlan.zhihu.com/p/348554267） 這里先給出錯(cuò)誤的操作：

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. 另一個(gè)正確姿勢是在 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) ```

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

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

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

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

啟動程序 由于博主目前也只是實(shí)踐了單機(jī)多卡操作，因此主要對單機(jī)多卡進(jìn)行介紹。區(qū)別于平時(shí)簡單的運(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ù)，告訴每個(gè)進(jìn)程自己的位置，要使用哪張GPU

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

# 為這個(gè)進(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í)，設(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 
# 將會讓程序在開始時(shí)花費(fèi)一點(diǎn)額外時(shí)間，為整個(gè)網(wǎng)絡(luò)的每個(gè)卷積層搜索最適合它的卷積實(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的更新

對于比較消耗顯存的訓(xùn)練任務(wù)時(shí)，往往單卡上的相對批量過小，影響模型的收斂效果。跨卡同步 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ī)種子的狀況。為了保險(xiǎn)起見，建議確保每個(gè)模型初始化的隨機(jī)種子相同，保證每個(gè)GPU進(jìn)程上的模型是同步的。

站在巨人的肩膀上，對前段時(shí)間自學(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 裝飾器的八種寫法，你隨便問~