動轉靜兩大升級!一鍵轉靜成功率領先,重點模型訓練提速18%+

語言: CN / TW / HK

目前主流深度學習框架支援的程式設計方式有兩種,分別為動態圖和靜態圖。動態圖的Pythonic程式設計體驗更佳、更易除錯,但效能方面與靜態圖有一定差距。靜態圖先組網再執行,預先擁有完整網路結構,更利於全域性優化,雖除錯難度大,但執行效能更佳。

百度飛槳採用動靜統一的技術架構設計,提供了動轉靜(@to_static)模組功能,支援使用者動態圖程式設計,並可一鍵切換靜態圖訓練和部署。2022年11月,飛槳框架 2.4 版本(以下簡稱飛槳v2.4)正式釋出,動轉靜“轉換成功率”和“訓練效能”迎來全面升級,帶來了全新的使用者使用體驗。

  • 動轉靜成功率明顯提升,一鍵轉換成功率達到92.1%。

  • 轉靜訓練加速效果明顯,重點模型訓練可提速18%+。

 

一鍵轉靜成功率明顯提升

動轉靜的轉換成功率是動轉靜功能的一個重要指標,與使用者的使用體驗息息相關,飛槳v2.4從“動轉靜語法完善”和“API動靜行為統一”兩個方面進行了重點優化和升級:

動轉靜語法完善

JIT 式動態執行

新增 Shape、Len、Attr、List、Unpack、Indexable等JIT 形式介面,提升語法轉寫的魯棒性。

控制流語法重構

重構了控制流IF/For/While語法轉寫邏輯,完備支援複雜巢狀場景下變數名解析等疑難問題。

關鍵字語法優化

優化了控制流中提前return、break、continue 等關鍵字語法轉寫機制,有效減少了靜態圖中間表示多餘運算元的引入,提升執行效率。

 

API動靜行為統一

屬性引數可變

完成了20多箇中高頻動態圖API引數的升級,如

Reduce列的paddle.mean/sum/max/min API的引數 axis ,新增支援為Tensor型別,動態可變。

介面動靜統一

補齊了Tensor類靜態圖下缺失的介面,升級了paddle.to_tensor、paddle.grad等高頻API 功能,支援靜態圖呼叫。

Einsum 升級

實現了動靜統一愛因斯坦求和運算元,並支援Python二元、多元輸入,訓推一體。

圖片

動轉靜成功率和語法支援度

飛槳v2.4下,動轉靜具備了更豐富的語法支援,Python語法支援比例達到了90%,在80多個外部使用者真實論文復現模型集合上,動轉靜一鍵轉寫成功率提升至92.1%,功能完備性和易用性都有明顯提升。

如下是一個動轉靜匯出預測模型的樣例程式碼:

import paddle

class SimpleNet(paddle.nn.Layer):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.linear = paddle.nn.Linear(10, 3)

    @paddle.jit.to_static   # step 1: 新增裝飾器
    def forward(self, x):
        out = self.linear(x)
        out = out + 1
        return out

net = SimpleNet()
train(net)  # 此處略去了訓練過程

# step 2: 切換到 eval() 模式
net.eval()
# step 3: 呼叫 jit.save 介面
paddle.jit.save(net, path='./simple_net')

執行上述程式碼樣例後,在當前目錄下會生成三個檔案,即代表成功匯出預測模型:

simple_net.pdiparams        // 存放模型中所有的權重資料
simple_net.pdmodel          // 存放模型的網路結構
simple_net.pdiparams.info   // 存放額外的其他資訊

動轉靜匯出模型一般包括三個步驟:

  • 新增裝飾器

將@to_static裝飾器裝飾在forward函式上。

  • 切換 eval 模式

如Dropout 、LayerNorm 介面在 train() 和 eval() 模式下行為存在較大的差異,在模型匯出前,請務必確認模型已切換到正確的模式。

  • 呼叫 save 介面

    呼叫 paddle.jit.save介面匯出其對應的模型檔案和引數檔案。

飛槳動轉靜@to_staitc更多功能用法,可參考【擴充套件閱讀動轉靜使用樣例】

 

動轉靜訓練加速效果明顯

飛槳框架中,通常情況下使用動態圖訓練即可滿足大部分場景需求。飛槳v2.4優化了動轉靜訓練的相關邏輯,面向重點模型,動態圖訓練的效能已經可以和靜態圖媲美,例如在ResNet50、Transformer、YOLOv3等模型上,動轉靜訓練相較於動態圖有18.6%~21.5%的顯著加速效果。

 

圖片

重點模型加速效果

在如下場景中,開發者可以考慮使用動轉靜方式進行模型訓練,將會獲得較明顯的效能提升效果。

 

場景一:重排程模型

即每個API背後的GPU Kernel 計算耗時較少,在CPU端拉起後很快就執行完了,此類任務的特點:

  • PU 利用率較低(可通過watch -n 1 nvidia-smi命令檢視)。

  • 常見於NLP 領域或AMP/FP16 任務。

  • 訓練效能瓶頸點主要是Host端排程開銷。

圖片

如上圖是重排程模型的動態圖和動轉靜 Timeline 示意圖。從圖中可以看出:

  • 一個Batch的訓練耗時取決於 Host 端總耗時。

  • 動態圖每個Python API在執行時,都會產生一次Python 和C++互動,會產生較大的排程開銷。

  • 動轉靜之後,整體上切分為執行前向和反向的兩個Python C API,故減少了很多個API間的排程開銷。

  • 動轉靜核心執行器也經過了極致的優化(如Instruction快取等),Kernel launch效率也會比純動態圖模式要高。

對於想使用動態圖訓練程式碼的使用者來說,只需要在組網入口的forward函式處新增裝飾器@to_static,其他程式碼無需改動就可以一鍵切換為動轉靜訓練。@to_static裝飾器會將此函式內的所有subLayers 轉化為一個靜態子圖並執行。如下是一個動轉靜訓練樣例程式碼:

import paddle

class SimpleNet(paddle.nn.Layer):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.linear = paddle.nn.Linear(103)

    @paddle.jit.to_static   # 僅需一行程式碼
    def forward(self, x):
        out = self.linear(x)
        out = out + 1
        return out

# create network
net = SimpleNet()
adam = opt.Adam(learning_rate=0.001, parameters=net.parameters())

for batch_id, x in enumerate(data_loader()):
    out = net(x)
    loss = paddle.mean(out)
    loss.backward()
    opt.step()
    opt.clear_grad()
 

 

場景二:排程+計算共存模型

即模型訓練時同時存在計算量小和計算量大的GPU Kernel,且一個Batch的起始位置常為小Kernel,此類任務的特點:

  • GPU 利用率波動比較大(可通過watch -n 1 nvidia-smi命令檢視)。

  • 訓練效能瓶頸點同時受區域性排程和區域性 Kernel 計算效率影響。

     

圖片

如上圖是排程+計算共存的動態圖和動轉靜Timeline示意圖。從圖中可以看出:

  • 一個Batch的訓練耗時取決於 Max(Host端,GPU端)。

  • 動轉靜降低了Python C API 排程開銷,收益點大多在Batch前半部分,後半部分可能會被overlap,排程方面的收益會打折扣。

  • 動轉靜可藉助全域性圖優化技術,通過運算元融合等技術提升模型訓練的吞吐。

在此種場景下,飛槳動轉靜@to_static API 提供build_strategy引數,在動轉靜的全圖視角下,使用者可以通過build_strategy引數開啟不同的全域性圖優化策略。通過裝飾器@to_static(build_strategy=get_build_strategy())或者API呼叫paddle.jit.to_static(net, build_strategy=get_build_strategy())兩種方式開啟全域性圖優化策略,如下是一個簡單的使用樣例:

import numpy as np
import paddle
import paddle.nn as nn

def get_build_strategy():
    build_strategy = paddle.static.BuildStrategy()
    # 運算元融合策略
    build_strategy.fuse_elewise_add_act_ops = True
    # 梯度 addto 策略
    build_strategy.enable_addto = True
    os.environ['FLAGS_max_inplace_grad_add'] = "8"
    return build_strategy

class ResNet(paddle.nn.Layer):
    # 此處省略了模型定義
    @to_static(build_strategy=get_build_strategy()) # 方式一
    def forward(self, image):
        # 省略前向程式碼

# create network
net = ResNet()
# 藉助 build_strategy 引數自定義開啟「全域性圖優化」策略
net = paddle.jit.to_static(net, build_strategy=get_build_strategy()) # 方式二

adam = opt.Adam(learning_rate=0.001, parameters=net.parameters())

for batch_id, image in enumerate(data_loader()):
    out = layer(image)
    loss = paddle.mean(out)
    loss.backward()
    opt.step()
    opt.clear_grad()
 

飛槳動轉靜@to_static開啟更多全域性圖優化用法,可參考【擴充套件閱讀動轉靜訓練圖優化策略】

拓展閱讀

[1] 動轉靜使用樣例

http://www.paddlepaddle.org.cn/documentation/docs/zh/guides/jit/basic_usage_cn.html

[2] 動轉靜訓練圖優化策略

http://www.paddlepaddle.org.cn/documentation/docs/zh/guides/jit/basic_usage_cn.html#sidongzhuanjinggengduoyongfa

[3] 動轉靜轉換原理

http://www.paddlepaddle.org.cn/documentation/docs/zh/guides/jit/principle_cn.html

[4] 動轉靜報錯除錯

http://www.paddlepaddle.org.cn/documentation/docs/zh/guides/jit/debugging_cn.html

[5] 動轉靜Limitations

http://www.paddlepaddle.org.cn/documentation/docs/zh/develop/guides/jit/limitations_cn.html