MLPerf世界紀(jì)錄技術(shù)分享：優(yōu)化卷積合并算法提升Resnet50推理性能

MLPerf是一套衡量機(jī)器學(xué)習(xí)系統(tǒng)性能的權(quán)威標(biāo)準(zhǔn)，將在標(biāo)準(zhǔn)目標(biāo)下訓(xùn)練或推理機(jī)器學(xué)習(xí)模型的時(shí)間，作為一套系統(tǒng)性能的測(cè)量標(biāo)準(zhǔn)。MLPerf推理任務(wù)包括圖像識(shí)別（ResNet50）、醫(yī)學(xué)影像分割（3D-UNet）、目標(biāo)物體檢測(cè)（SSD-ResNet34）、語音識(shí)別（RNN-T）、自然語言理解（BERT）以及智能推薦（DLRM）。在MLPerf V2.0推理競賽中，浪潮AI服務(wù)器基于ImageNet數(shù)據(jù)集在離線場景中運(yùn)行Resnet50，達(dá)到了449,856 samples/s的計(jì)算性能，位居世界第一。

本文將介紹浪潮在MLPerf推理競賽中使用的卷積合并計(jì)算算法。

Resnet是殘差網(wǎng)絡(luò)(Residual Network)的縮寫，該系列網(wǎng)絡(luò)廣泛用于目標(biāo)分類等領(lǐng)域以及作為計(jì)算機(jī)視覺任務(wù)主干經(jīng)典神經(jīng)網(wǎng)絡(luò)的一部分，典型的網(wǎng)絡(luò)有Resnet50、Resnet101等。在Resnet神經(jīng)網(wǎng)絡(luò)中，主要計(jì)算算子是卷積計(jì)算層。Resnet50神經(jīng)網(wǎng)絡(luò)具有4組殘差結(jié)構(gòu)，這4組殘差結(jié)構(gòu)包含48個(gè)卷積算子，通過設(shè)計(jì)卷積算子的計(jì)算算法，提高卷積算子的計(jì)算性能，可以減少Resnet50推理過程中的延遲?；?a href="http://www.ttokpm.com/article/zt/" target="_blank">最新GPU單卡的性能測(cè)試顯示，在BatchSize=2048的情況下，優(yōu)化后的卷積合并優(yōu)化算法相比原算法可帶來14.6%的性能提升。

MLPerf Resnet50推理流程

在MLPerf V2.0推理測(cè)試中，Resnet50模型需要在ImageNet2012測(cè)試集上達(dá)到FP32精度（76.46%）的99%以上。數(shù)據(jù)中心賽道設(shè)置了離線(Offline)與在線(Server)兩種模式，其中離線模式會(huì)產(chǎn)生一次推理時(shí)間大于10分鐘的samples請(qǐng)求，可直接反映機(jī)器和算法的推理性能。

Resnet50推理流程如下。首先在ImageNet2012測(cè)試集中讀取數(shù)據(jù)，并進(jìn)行數(shù)據(jù)預(yù)處理，隨后數(shù)據(jù)會(huì)加載到TensorRT中進(jìn)行實(shí)際的推理測(cè)試。測(cè)試分為兩方面，一是測(cè)試模型的精度；二是產(chǎn)生一次推理請(qǐng)求，TensorRT會(huì)將請(qǐng)求中的圖片全部推理完成得到總時(shí)間，根據(jù)計(jì)算時(shí)間得到每秒推理的樣本數(shù)量，即為最終的成績。

圖1 MLPerf Resnet50推理流程▲

卷積合并計(jì)算算法

▏2.1 算法優(yōu)化思路

在GPU上運(yùn)行Resnet50圖像推理模型時(shí)，需要將每一個(gè)算子（卷積、池化、全連接等）放在GPU的Kernel中進(jìn)行算子計(jì)算，由于GPU上運(yùn)行Kernel時(shí)共享內(nèi)存以及寄存器的資源有限，不可能將所有的計(jì)算過程數(shù)據(jù)放到Kernel中，而GPU的全局內(nèi)存一般都很大，所以會(huì)將比較大的過程數(shù)據(jù)放在全局內(nèi)存中。在進(jìn)行推理時(shí)，根據(jù)Kernel的計(jì)算將數(shù)據(jù)按需從全局內(nèi)存讀取到Kernel中進(jìn)行計(jì)算，每個(gè)算子在計(jì)算時(shí)會(huì)不可避免地產(chǎn)生Kernel與全局內(nèi)存的數(shù)據(jù)交換，由于全局內(nèi)存的讀寫訪問延遲較大，會(huì)使算子計(jì)算性能下降。

對(duì)于每個(gè)算子的Kernel計(jì)算，會(huì)產(chǎn)生兩部分的全局內(nèi)存訪問，一部分是最開始的全局內(nèi)存讀取，另一部分是Kernel計(jì)算完成后的全局內(nèi)存寫回。為了降低全局內(nèi)存訪問帶來的性能影響，有如下兩種辦法：

一是采用算子合并的方式。默認(rèn)的程序會(huì)將每個(gè)算子都放在單獨(dú)的Kernel中進(jìn)行計(jì)算，每個(gè)算子都會(huì)產(chǎn)生全局內(nèi)存讀和寫兩次訪問。如果將兩個(gè)算子放在一個(gè)Kernel中進(jìn)行計(jì)算，對(duì)于連續(xù)的兩個(gè)卷積計(jì)算，可以減少第一個(gè)卷積算子的寫回以及第二個(gè)算子的讀?。粚?duì)于卷積與Shortcut的合并，可以減少一次全局內(nèi)存的寫回操作，通過減少全局內(nèi)存的訪問可以提高程序的計(jì)算性能；

二是根據(jù)GPU不同架構(gòu)的計(jì)算特性對(duì)Kernel的內(nèi)部計(jì)算進(jìn)行合理的優(yōu)化設(shè)計(jì)。當(dāng)不可避免地需要對(duì)全局內(nèi)存進(jìn)行訪問時(shí)，做到全局內(nèi)存進(jìn)行連續(xù)線程的融合讀取，充分利用向量化讀取等加速對(duì)全局內(nèi)存的訪問，同時(shí)優(yōu)化計(jì)算流程，通過Double buffer用計(jì)算來隱藏內(nèi)存的訪問延遲，對(duì)于需求較晚的全局內(nèi)存數(shù)據(jù)，也可以通過GPU的新特性-全局內(nèi)存的異步復(fù)制來隱藏?cái)?shù)據(jù)讀取過程。

本文主要針對(duì)MLPerf推理中Resnet50卷積神經(jīng)網(wǎng)絡(luò)的第二組殘差結(jié)構(gòu)中的部分算子進(jìn)行計(jì)算合并，在充分考慮GPU計(jì)算特性前提下，進(jìn)行合理的算法設(shè)計(jì)，提高Resnet50卷積神經(jīng)網(wǎng)絡(luò)的性能。

▏2.2 Resnet50合并計(jì)算算法

在Resnet50神經(jīng)網(wǎng)絡(luò)中，第二組殘差結(jié)構(gòu)有Res3.1、Res3.2、Res3.3、Res3.4，共四部分的卷積計(jì)算，其中Res3.2、Res3.3、Res3.4三部分計(jì)算結(jié)構(gòu)一樣，如下圖所示：

圖2 Resnet50中第二組殘差結(jié)構(gòu)Res3.2示意圖▲

可以看到，Res3.1的輸出（input）作為Res3.2部分的輸入，輸入后會(huì)有兩部分分支，在右部分的分支中，會(huì)先后計(jì)算Conv1，Conv2，Conv3三個(gè)卷積，其中Conv1，Conv2兩個(gè)卷積后面都包含Bn和Relu過程，Conv3后面會(huì)有Bn的計(jì)算過程；在右邊分支計(jì)算完成后，會(huì)與input進(jìn)行Shortcut操作，主要進(jìn)行的是與輸入數(shù)據(jù)Sum和Relu操作，兩部分結(jié)果經(jīng)過Shortcut操作后會(huì)得到Res3.2的輸出完成這部分的計(jì)算。

本文介紹的合并算法對(duì)圖2虛線框中的計(jì)算進(jìn)行合并，主要是對(duì)Conv3以及Shortcut的過程進(jìn)行合并，包含Conv3+Bn+Sum+Relu過程。

卷積合并算法在GPU加速卡上的實(shí)現(xiàn)

Res3.2的計(jì)算參數(shù)主要如下：

通過上表可以看到，Conv3輸入Data的H*W為28*28，輸入通道Ic為128，輸入的權(quán)重Weight的h*w為1*1，輸入通道Ic為128，輸出通道Oc為512；Shortcut的輸入同Conv1，其中H*W為28*28，輸入通道Ic為512；兩部分計(jì)算合并之后的輸出Output的H*W為28*28，通道Oc為512。

▏ 3.1 關(guān)于data、weight、output的layout變換

本文采用的計(jì)算數(shù)據(jù)類型為int8，因此下文介紹的所有內(nèi)容都是基于int8開展的優(yōu)化。

算法對(duì)data以及weight進(jìn)行了提前處理以適應(yīng)GPU的計(jì)算特性，主要處理如下：

對(duì)于data，原始layout為[B, H, W, Ic]=[B, 28, 28, 128]，算法將Ic=128以32為單位進(jìn)行拆分為4組，形成[B, 4, H, W, Ic/32]=[B, 4(I1), 28, 28, 32(I2)]的layout，這樣做的目的是32個(gè)int8可以組成16B共128位數(shù)據(jù)的聯(lián)合向量化讀寫，提高GPU中全局內(nèi)存的通信速度。

對(duì)于weight，由于h*w=1*1，因此本文后續(xù)不再表示h*w，默認(rèn)的weight的layout為[Ic, Oc]=[128, 512]，算法將Ic以32為單位進(jìn)行拆分為4組，將4放在左數(shù)第二維，將32放在左數(shù)第四維，這樣做的目的也是為了程序在訪問全局內(nèi)存時(shí)做到16B共128位數(shù)據(jù)的聯(lián)合向量化讀寫；算法將Oc以128為單位進(jìn)行拆分為4組，將4放在左數(shù)第一維，將128放在左數(shù)第三維，這樣做的目的是將Oc拆成了4組放在了不同的block中進(jìn)行計(jì)算，這樣在每個(gè)block進(jìn)行計(jì)算的時(shí)候可以順序的由全局內(nèi)存加載weight，不會(huì)產(chǎn)生數(shù)據(jù)內(nèi)存位置的跳躍，這部分會(huì)在后面block的劃分中進(jìn)行介紹，這樣就形成[O1, I1, O2, I2] =[4, 4, 128, 32]的weight的layout。

對(duì)于output，原始layout為[B, H, W, Oc]=[B, 28, 28, 512]，這部分?jǐn)?shù)據(jù)類似于輸入data，將Oc以32為單位進(jìn)行拆分為16組，形成layout為[B, 16, H, W, Oc/32]=[B, 16(O1), 28, 28, 32(O2)]。

▏ 3.2 關(guān)于Grid以及Block的并行劃分

對(duì)于Grid的劃分，首先是x維度，由上文可知，對(duì)于Conv3的Oc為512，本文將Oc劃分為4組放到Grid.x維度，每組計(jì)算的Oc為128；對(duì)于y維度，將H*W=28*28=784分為49組放到Grid.y維度，每組計(jì)算的HW為16；對(duì)于z維度，將B分為B/4組放到Grid.z維度，每組計(jì)算B的數(shù)量是4。這樣經(jīng)過劃分，Grid的數(shù)量為[Grid.x, Grid.y, Grid.z]=[4, 49, B/4]，即共有4*49*B/4組計(jì)算同時(shí)并行進(jìn)行。GPU上SM數(shù)量有108個(gè)，當(dāng)B≥4時(shí)，一個(gè)kernel共需要啟動(dòng)大于4*49*4/4=196個(gè)Block，完全滿足Grid維度并行度的要求。

對(duì)于Block中的劃分，GPU中一個(gè)SM的warp schedule為4個(gè)，因此一個(gè)block中線程數(shù)量至少大于或等于128，為了實(shí)現(xiàn)更好的并行度，算法選擇一個(gè)Block中設(shè)置8個(gè)warp共256個(gè)線程。

▏ 3.3 關(guān)于Block內(nèi)部計(jì)算層次劃分

圖3 Block內(nèi)部計(jì)算層析劃分▲

由上文可知，一個(gè)Block中劃分的Output的計(jì)算shape為[B, H*W, Oc]=[4, 16, 128]，由于在1*1卷積計(jì)算中B維度與HW維度具有同等地位，因此將BHW合并為一個(gè)維度，此時(shí)本文用M表示BHW維度，即M=BHW=4*16=64，用N表示Oc維度，即N=Oc=128，此時(shí)一個(gè)Block中的計(jì)算維度變?yōu)閇M,N]=[64, 128]。

由前述data的layout的變換可知形成[B, 4, H, W, Ic/32]=[B, 4, 28, 28, 32]的layout，由于Grid.y以及Grid.x對(duì)B維度以及HW維度進(jìn)行了劃分，此時(shí)一個(gè)Block中data的輸入數(shù)據(jù)為[B, I1, HW, I2]=[4, 4, 16, 32]，用上段所述BHW合并為1維用M表示，即M=B*HW=4*16=64，K表示I1*I2維度，即K=I1*I2=128，則此時(shí)一個(gè)Block中data的計(jì)算維度變?yōu)閇M, K]=[64, 128]。

由前述weight的layout的變換可知，weight的layout為[O1, I1, O2, I2]=[4, 4, 128, 32]，由于對(duì)Oc按照32為單位劃分4組在Grid.x維度，因此每個(gè)Block中計(jì)算的Oc為128，此時(shí)一個(gè)Block中的weight的計(jì)算數(shù)據(jù)為[I1, O2, I2] =[4, 128, 32]，用N表示O2維度，即N=O2=128，用K表示I1*I2維度，即K=I1*I2=128，則此時(shí)一個(gè)Block中weight的計(jì)算維度變?yōu)閇N, K]=[128, 128]。

一個(gè)Block中實(shí)際要進(jìn)行的計(jì)算就變?yōu)橐粋€(gè)矩陣乘data[M, K]點(diǎn)乘weight[N, K]等于output[M, N]，即[64, 128]﹒[128,128]=[64, 128]，共4*49*B/4個(gè)Block并行完成所有整個(gè)卷積合并的計(jì)算，其中data的實(shí)際維度為[B, I1, HW, I2]=[4, 4, 16, 32]，weight的實(shí)際維度為[I1, O, I2]=[4, 128, 32]。

經(jīng)過前面的劃分，一個(gè)Thread Block層次實(shí)際計(jì)算量為[64, 128]﹒[128,128]=[64, 128]。

為了加速int8矩陣乘的計(jì)算，程序采用了CUDA中mma進(jìn)行加速計(jì)算，其中mma的計(jì)算形狀為[m ,n ,k]=[16, 8, 32]，為了配合共享內(nèi)存，寄存器以及mma形狀的匹配，程序?qū)?nèi)積方向的K維度128拆分為2組64進(jìn)行計(jì)算，每組64進(jìn)一步拆分為2組32(k)進(jìn)行計(jì)算，這樣最基礎(chǔ)的Thread Block層次進(jìn)行的計(jì)算就變?yōu)閳D3中左上角虛線框中所示的[M, k]﹒[N, k]=[M ,N]即[64, 32]﹒[128, 32]=[64 ,128]，由于一個(gè)Block中設(shè)置warp的數(shù)量為8，8個(gè)warp會(huì)對(duì)Thread Block中的計(jì)算任務(wù)進(jìn)行劃分，每個(gè)warp計(jì)算任務(wù)為[32, 32]﹒[32, 32]=[32, 32]的矩陣乘，經(jīng)過內(nèi)積方向的4次32的循環(huán)，在warp level便可以將內(nèi)積方向K=128完全計(jì)算得[M, N]=[32, 32]的計(jì)算結(jié)果，則8個(gè)warp合并可得[M, N]=[64, 128]的計(jì)算結(jié)果。

如上文所述，程序?qū)?nèi)積方向的K維度128拆分為2組64進(jìn)行計(jì)算，每組64進(jìn)一步拆分為2組32(k)進(jìn)行計(jì)算。這么做的目的是將data以及weight的全局內(nèi)存加載變成了Double buffer模式，即首先將第一組的數(shù)據(jù)由全局內(nèi)存加載到共享內(nèi)存，然后在利用第一組的數(shù)據(jù)進(jìn)行計(jì)算前，便提交第二組數(shù)據(jù)由全局內(nèi)存加載到共享內(nèi)存的過程，這樣可以利用第一組數(shù)據(jù)的計(jì)算過程時(shí)間來隱藏第二組數(shù)據(jù)的全局內(nèi)存加載到共享內(nèi)存的過程的時(shí)間，整體流程示意圖如下：

圖4 程序整體流程圖▲

如前所述，每個(gè)warp計(jì)算任務(wù)為[32, 32]﹒[32, 32]=[32, 32]的矩陣乘，因此在warp的計(jì)算層次配合[m ,n ,k]=[16, 8, 32]的形狀，需要進(jìn)行row=2，col=4，共row*col=8次mma的計(jì)算才可以得到warp 層次的計(jì)算結(jié)果，在計(jì)算時(shí)配合ldmatrix的使用可以進(jìn)一步提高程序的計(jì)算性能。

對(duì)于Mma層次的計(jì)算，根據(jù)mma的形狀，單次計(jì)算為[m , k]﹒[n, k]=[16, 32]﹒[8, 32]=[16, 8]。

▏ 3.4 Shoutcut的合并計(jì)算

經(jīng)過以上計(jì)算，每個(gè)Block程序會(huì)得到Conv3的[M, N]=[64, 128]的計(jì)算結(jié)果，由于程序?qū)n+Sum+Relu進(jìn)行了合并，因此需要對(duì)Res3.1輸出的原始數(shù)據(jù)進(jìn)行加載。根據(jù)Grid[x , y, z]的劃分，可以相應(yīng)的得到Shortcut的數(shù)據(jù)偏移，為了隱藏這部分?jǐn)?shù)據(jù)在全局內(nèi)存加載到共享內(nèi)存時(shí)通信延遲，程序利用了GPU異步復(fù)制(pipeline_memcpy_async)的新特性，在程序的最開始便提交了這部分?jǐn)?shù)據(jù)的加載，這樣可以最大程度上利用計(jì)算的時(shí)間同時(shí)進(jìn)行數(shù)據(jù)的加載以隱藏Shortcut的通信延遲，如圖4所示。完成數(shù)據(jù)的加載后，會(huì)以warp為單位對(duì)每一個(gè)計(jì)算結(jié)果進(jìn)行Bn+Sum+Relu的操作，最后將數(shù)據(jù)由寄存器寫回共享內(nèi)存，再寫回全局內(nèi)存完成整個(gè)卷積合并算法的計(jì)算。

性能提升效果

根據(jù)上文介紹的卷積合并優(yōu)化算法，在TensorRT中增加了關(guān)于卷積合并算法的plugin以替代原始算法，在最新GPU單卡進(jìn)行Conv3+Bn+Sum+Relu性能測(cè)試，在BatchSize=2048的情況下，原算法的性能為123TOPS，經(jīng)過優(yōu)化后的卷積合并優(yōu)化算法性能為141TOPS，算子相比較原算法可以帶來14.6%的性能提升。通過合并Res3.2、Res3.3、Res3.4三部分Conv3+Bn+Sum+Relu算子合并，可將Resnet50推理性能提升1%-2%。同樣該算法合并思路可以用到其他殘差結(jié)構(gòu)中，通過合理的算法設(shè)計(jì)帶來整體的程序性能提升。

在MLPerf V2.0推理競賽中，浪潮通過軟件與硬件優(yōu)化，基于ImageNet數(shù)據(jù)集Resnet50模型，在Offline場景中達(dá)到了449,856 samples/s的計(jì)算性能，位居世界第一。

審核編輯：湯梓紅