1、KubeAI介紹

KubeAI是得物AI平臺，是我們在容器化過程中，逐步收集和挖掘公司各業(yè)務(wù)域在AI模型研究和生產(chǎn)迭代過程中的需求，逐步建設(shè)而成的一個云原生AI平臺。KubeAI以模型為主線提供了從模型開發(fā)，到模型訓(xùn)練，再到推理(模型)服務(wù)管理，以及模型版本持續(xù)迭代的整個生命周期內(nèi)的解決方案。

在數(shù)據(jù)方面，KubeAI提供基于cvat的標(biāo)注工具，與數(shù)據(jù)處理及模型訓(xùn)練流程打通，助力線上模型快速迭代；提供任務(wù)/Pipeline編排功能，對接ODPS/NAS/CPFS/OSS數(shù)據(jù)源，為用戶提供一站式AI工作站。平臺自研推理引擎助力業(yè)務(wù)在提高模型服務(wù)性能的同時還能控制成本；自研訓(xùn)練引擎提高了模型訓(xùn)練任務(wù)吞吐量，縮短了模型的訓(xùn)練時長，幫助模型開發(fā)者加速模型迭代。

此外，隨著AIGC的火熱發(fā)展，我們經(jīng)過調(diào)研公司內(nèi)部AI輔助生產(chǎn)相關(guān)需求，上線了AI制圖功能，為得物海報、營銷活動、設(shè)計師團隊等業(yè)務(wù)場景提供了基礎(chǔ)能力和通用AI制圖能力。

此前，我們通過一文讀懂得物云原生AI平臺-KubeAI的落地實踐過程一文，向大家介紹了KubeAI的建設(shè)和在業(yè)務(wù)中的落地過程。本文，我們將重點介紹下KubeAI平臺在推理、訓(xùn)練和模型迭代過程中的核心引擎能力實踐經(jīng)驗。

2、AI推理引擎設(shè)計實現(xiàn)

2.1 推理服務(wù)現(xiàn)狀及性能瓶頸分析

Python語言以其靈活輕盈的特點，以及其在神經(jīng)網(wǎng)絡(luò)訓(xùn)練與推理領(lǐng)域提供了豐富的庫支持，在模型研究和開發(fā)領(lǐng)域被廣泛使用，所以模型推理服務(wù)也主要以Python GPU推理為主。模型推理過程一般涉及預(yù)處理、模型推理、后處理過程，單體進程的方式下CPU前/后處理過程，與GPU推理過程需要串行，或者假并行的方式進行工作，大致流程如下圖所示：

上述架構(gòu)的優(yōu)勢是代碼寫起來比較通俗易懂，但在性能上有很大的弊端，所能承載的QPS比較低。通過在CV域的模型上進行壓測，我們發(fā)現(xiàn)推理QPS很難達到5，深入分析發(fā)現(xiàn)造成這一問題的原因如下：

（1）單線程模式下，CPU邏輯與GPU邏輯相互等待，GPU Kernel函數(shù)調(diào)度不足，導(dǎo)致GPU使用率不高，無法充分提升服務(wù)QPS。這種情況下只能開啟更多進程來提升QPS，但是更多進程會帶來更大的GPU顯存開銷。

（2）多線程模式下，由于Python的GIL鎖的原因，Python的多線程實際上是偽的多線程，并不是真正的并發(fā)執(zhí)行，而是多個線程通過爭搶GIL鎖來執(zhí)行，這種情況下GPU Kernel Launch線程不能得到充分的調(diào)度。此外，在Python推理服務(wù)中開啟多線程反而會導(dǎo)致GPU Kernel Launch線程頻繁被CPU的線程打斷，所以GPU算力也會一直“萎靡不振”，持續(xù)低下。

以上問題使得 如果推理服務(wù)想要支撐更多的流量，只能做橫向的增加服務(wù)實例數(shù)，伴隨著成本的上漲。

2.2 自研推理服務(wù)統(tǒng)一框架kubeai-inference-framework

針對以上問題，KubeAI的解決方案是把CPU邏輯與GPU邏輯分離在兩個不同的進程中：CPU進程主要負(fù)責(zé)圖片的前處理與后處理，GPU進程則主要負(fù)責(zé)執(zhí)行CUDA Kernel 函數(shù)，即模型推理。

為了方便模型開發(fā)者更快速地接入我們的優(yōu)化方案，我們基于Python開發(fā)了一個CPU與GPU進程分離的統(tǒng)一框架kubeai-inference-framework，舊有Flask或Kserve的服務(wù)，稍作修改即可接入推理引擎統(tǒng)一框架，新增服務(wù)按照框架實現(xiàn)指定function即可。推理服務(wù)統(tǒng)一框架構(gòu)如下圖所示：

如前所述，推理服務(wù)統(tǒng)一框架的主要思路是把GPU邏輯與CPU邏輯分離到兩個進程，除此之外，還會拉起一個Proxy進程做路由轉(zhuǎn)發(fā)。

CPU進程

CPU進程主要負(fù)責(zé)推理服務(wù)中的CPU相關(guān)邏輯，包括前處理與后處理。前處理一般為圖片解碼，圖片轉(zhuǎn)換，后處理一般為推理結(jié)果判定等邏輯。CPU進程在前處理結(jié)束后，會調(diào)用GPU進程進行推理，然后繼續(xù)進行后處理相關(guān)邏輯。CPU進程與GPU進程通過共享內(nèi)存或網(wǎng)絡(luò)進行通信，共享內(nèi)存可以減少圖片的網(wǎng)絡(luò)傳輸。

GPU進程

GPU進程主要負(fù)責(zé)運行GPU推理相關(guān)的邏輯，它啟動的時候會加載很多模型到顯存，然后在收到CPU進程的推理請求后，直接觸發(fā)Kernel Lanuch調(diào)用模型進行推理。

kubeai-inference-framework框架中對模型開發(fā)者提供了一個Model類接口，他們不需要關(guān)心后面的調(diào)用邏輯，只需要填充其中的前處理，后處理的業(yè)務(wù)邏輯，就可以快速上線模型服務(wù)，自動拉起這些進程。

Proxy進程

Proxy進程是推理服務(wù)入口，對外提供調(diào)用接口，負(fù)責(zé)路由分發(fā)與健康檢查。當(dāng)Proxy進程收到請求后，會輪詢調(diào)用CPU進程，分發(fā)請求給CPU進程進行處理。

自研的推理服務(wù)統(tǒng)一框架，把CPU邏輯（圖片解碼，圖片后處理等）與GPU邏輯（模型推理）分離到兩個不同的進程中后，有效解決了Python GIL鎖帶來的GPU Kernel Launch調(diào)度問題，提升了GPU利用率，提高了推理服務(wù)性能。針對線上的某個推理服務(wù)，使用我們的框架進行了CPU與GPU進程分離，壓測得出的數(shù)據(jù)如下表所示，可以看到QPS提升了近7倍。

2.3 做的更好 — 引入TensorRT優(yōu)化加速

在支持推理服務(wù)接入kubeai-inference-framework統(tǒng)一框架的過程中，我們繼續(xù)嘗試在模型本身做優(yōu)化提升。經(jīng)過調(diào)研和驗證，我們將現(xiàn)有pth格式模型通過轉(zhuǎn)成TensorRT格式，并開啟FP16，在推理階段取得了更好的QPS提升，最高可到10倍提升。

TensorRT是由英偉達公司推出的一款用于高性能深度學(xué)習(xí)模型推理的軟件開發(fā)工具包，可以把經(jīng)過優(yōu)化后的深度學(xué)習(xí)模型構(gòu)建成推理服務(wù)部署在實際的生產(chǎn)環(huán)境中，并提供基于硬件級別的推理引擎性能優(yōu)化。業(yè)內(nèi)最常用的TensorRT優(yōu)化流程，是把pytorch / tensorflow等模型先轉(zhuǎn)成onnx格式，然后再將onnx格式轉(zhuǎn)成TensorRT（trt）格式進行優(yōu)化，如下圖所示：

TensorRT所做的工作主要在兩個時期，一個是網(wǎng)絡(luò)構(gòu)建期，另外一個是模型運行期。

• 網(wǎng)絡(luò)構(gòu)建期

1. 模型解析與建立，加載onnx網(wǎng)絡(luò)模型。
2. 計算圖優(yōu)化，包括橫向算子融合，或縱向算子融合等。
3. 節(jié)點消除，去除無用的節(jié)點。
4. 多精度支持，支持FP32/FP16/int8等精度。
5. 基于特定硬件的相關(guān)優(yōu)化。

• 模型運行期

6. 序列化，加載RensorRT模型文件。
7. 提供運行時的環(huán)境，包括對象生命周期管理，內(nèi)存顯存管理等

為了更好地幫助模型開發(fā)者使用TensorRT優(yōu)化，KubeAI平臺提供了kubeai-trt-helper工具，用戶可以使用該工具把模型轉(zhuǎn)成TensorRT格式，如果在模型轉(zhuǎn)換的過程中出現(xiàn)精度丟失等問題，也可以使用該工具進行問題定位與解決。kubeai-trt-helper主要在兩個階段為用戶提供幫助：一個是問題定位，另一個階段是模型轉(zhuǎn)換。

問題定位

問題定位階段主要是為了解決模型轉(zhuǎn)TensorRT開啟FP16模式時出現(xiàn)的精度丟失問題。一般分類模型，對精度的要求不是極致的情況下，盡量開啟FP16，F(xiàn)P16模式下，NVIDIA對于FP16有專門的Tensor Cores可以進行矩陣運算，相比FP32來說吞吐量提升一倍以上。比如在轉(zhuǎn)TensorRT時，開啟FP16出現(xiàn)了精度丟失問題，kubeai-trt-helper工具在問題定位階段的大致工作流程如下：

第1步：設(shè)定模型轉(zhuǎn)換精度要求后，標(biāo)記所有算子為輸出，然后對比所有算子的輸出精度。

第2步：找到最早的不符合精度要求的算子，對該算子進行如下幾種方式干預(yù)。

• 標(biāo)記該算子為FP32。
• 標(biāo)記其父類算子為FP32。
• 更改該算子的優(yōu)化策略。

循環(huán)通過以上2個步驟，最終找到符合目標(biāo)精度要求的模型參數(shù)。這些參數(shù)比如：需要額外開啟FP32的那些算子等。相關(guān)參數(shù)會輸出到配置文件中，如下：

模型轉(zhuǎn)換

模型轉(zhuǎn)換階段則直接使用上面問題定位階段得到的參數(shù)，調(diào)用TensorRT相關(guān)接口與工具進行轉(zhuǎn)換。此外，我們在模型轉(zhuǎn)換階段，針對TensorRT原有參數(shù)與API過于復(fù)雜的問題也做了一些封裝，提供了更為簡潔的接口，比如工具可以自動解析onnx，判斷模型的輸入與輸出shape，不需要用戶再提供相關(guān)shape信息等。

2.4 落地實踐成果

在實際應(yīng)用中，我們幫助算法域的模型開發(fā)同學(xué)，能夠?qū)σ粋€推理基于自研推理服務(wù)統(tǒng)一框架進行實現(xiàn)的同時，也開啟TensorRT優(yōu)化，這樣往往可以得到QPS兩次優(yōu)化的疊加效果。

2.4.1 分類模型，CPU與GPU分離，TensorRT優(yōu)化，并開啟FP16，得到10倍QPS提升

線上某個基于Resnet的分類模型，對精度損失可以接受誤差在0.001（誤差定義：median,atol,rtol）范圍內(nèi)。因此我們對該推理服務(wù)進行了3項性能優(yōu)化：

1. 使用kubeai-inference-framework統(tǒng)一框架，對CPU進程和GPU進程進行分離改造。
2. 對模型轉(zhuǎn)ONNX后，轉(zhuǎn)TensorRT。
3. 開啟FP16模式，并使用自研工具定位到中間出現(xiàn)精度損失的算子，把這些算子標(biāo)記為FP32。

經(jīng)過以上優(yōu)化，最終得到了10倍QPS的提升（與原來Pytorch直接推理比較），服務(wù)成本大幅削減。

2.4.2 檢測模型，CPU與GPU分離，TensorRT模型優(yōu)化，QPS提升4-5倍左右。

線上某個基于Yolo的檢查模型，由于對精度要求比較高，所以不能開啟FP16，我們直接在FP32的模式下進行了TensorRT優(yōu)化，并使用kubeai-inference-framework統(tǒng)一框架對GPU進程與CPU進程分離，最終得到QPS 4-5倍的提升。

2.4.3 模型推理進程多實例化，充分利用GPU算力資源

在實際的場景中，往往GPU的算力是充足的，而GPU顯存是不夠的。經(jīng)過TensorRT優(yōu)化后，模型運行時需要的顯存大小一般會降低到原來的1/3到1/2。所以為了充分利用GPU算力，kubeai-inference-framework統(tǒng)一框架進一步優(yōu)化，支持可以把GPU進程在一個容器內(nèi)復(fù)制多份，這種架構(gòu)即保證了CPU可以提供充足的請求給GPU，也保證了GPU算力充分利用。

線上某個模型，經(jīng)過TensorRT優(yōu)化后，顯存由原來的2.4G降低到只需要1.2G。在保持推理服務(wù)配置5G顯存不變的情況下，我們將GPU進程為復(fù)制4份，充分利用了5G顯存，使得服務(wù)吞吐達到了原來的4倍。

3、AI訓(xùn)練引擎優(yōu)化實踐

3.1 PyTorch框架概況

PyTorch是近年來較為火爆的深度學(xué)習(xí)框架，幾乎占據(jù)了CV（Computer Vision，計算機視覺）、NLP（Natural Language Processing，自然語言處理）領(lǐng)域各業(yè)務(wù)方向，算法同學(xué)基本都在使用PyTorch框架來進行模型訓(xùn)練。下圖是基于PyTorch框架進行模型訓(xùn)練時的代碼基本流程：

第1步：從pytorch dataloader中將本step訓(xùn)練過程中需要的數(shù)據(jù)拉出來。

第2步：將獲取到的數(shù)據(jù)，例如：樣本圖片、樣本標(biāo)簽的tensor等數(shù)據(jù)，復(fù)制到GPU顯存里。

第3步：開始正式的模型訓(xùn)練：前向計算、計算損失、計算梯度、 更新參數(shù)。

整個訓(xùn)練過程的耗時，也主要分布在上面3個步驟。通常第2步不會是瓶頸，因為大部分訓(xùn)練樣本圖片都是被resize變小之后才從內(nèi)存拷貝到到GPU顯存上的。但由于模型的差異性、訓(xùn)練數(shù)據(jù)的差異性，經(jīng)常是第1、2步會在訓(xùn)練過程中出現(xiàn)性能瓶頸，導(dǎo)致訓(xùn)練耗時長，GPU利用率低下，影響模型迭代效率。

3.2 Dataloader瓶頸分析及優(yōu)化

3.2.1 PyTorch Dataset/Dataloader分析

PyTorch訓(xùn)練讀取數(shù)據(jù)部分主要是通過Dataset、Dataloader的方式完成的，其中Dataset為用戶自定義讀取數(shù)據(jù)的類（繼承自 torch.utils.data.Dataset），而Dataloader是PyTorch實現(xiàn)的在訓(xùn)練過程中對Dataset的調(diào)度器。

torch.utils.data.dataloader
torch.utils.data.Dataset


train_loader = torch.utils.data.DataLoader( MyDataset, 
                      batch_size=16, 
                      num_workers=4, 
                      shuffle=True, 
                      drop_last=True,
                      pin_memory=False)


val_loader = torch.utils.data.DataLoader(MyDataset, 
                     batch_size=batch_size, 
                     num_workers=4, 
                     shuffle=False)

參數(shù)解釋如下：

• dataset(Dataset)：傳入的自定義Dataset（數(shù)據(jù)讀取的具體步驟）。
• batch_size(int, optional)：每個batch有多少個樣本，每個iter可以從dataloader中取出多少數(shù)據(jù)。
• shuffle(bool, optional)：在每個epoch開始的時候，對數(shù)據(jù)進行重新排序，可以使每個epoch讀取數(shù)據(jù)的組合和順序不同。
• num_workers (int, optional)：這個參數(shù)決定dataloader啟動幾個后臺進程來做數(shù)據(jù)拉取。0意味著所有的數(shù)據(jù)都會被load進主進程，默認(rèn)為0。
• collate_fn (callable, optional)：將一個list的sample組成一個mini-batch的函數(shù)，一般CV場景是concat函數(shù)。
• pin_memory (bool, optional)：如果設(shè)置為True，那么data loader將會在返回batch之前，將tensors拷貝到CUDA中的固定內(nèi)存（CUDA pinned memory）中， 這個參數(shù)某些場景下有妙用。
• drop_last (bool, optional)：該參數(shù)是對最后的未完成的batch來說的，比如batch_size設(shè)置為64，而一個epoch只有100個樣本，如果設(shè)置為True，那么訓(xùn)練的時候后面的36個就被扔掉了，否則會繼續(xù)正常執(zhí)行，只是最后的batch_size會小一點。默認(rèn)設(shè)置為False。

上述參數(shù)中，比較重要的是num_workers，Dataloader在構(gòu)造的時候，會啟動num_workers個worker進程，然后主進程會向worker進程分發(fā)讀取任務(wù)，worker進程讀到數(shù)據(jù)之后，再把數(shù)據(jù)放到隊列中供主進程取用。多進程模式使用的是torch.multiprocessing接口，可以實現(xiàn)worker進程與主進程之間共享內(nèi)存，而且共享內(nèi)存中可以存放tensor，這樣進程中如果返回tensor，可以通過共享內(nèi)存的方式直接將結(jié)果返回給主進程，減少多進程間的通訊開銷。

當(dāng)num_workers 為0 的時候： get_data()流程與train_model()過程是串行，效率非常低下，如下圖所示：

當(dāng)num_workers 大于0開啟多進程讀取數(shù)據(jù)， 并且讀取一個batch數(shù)據(jù)的時間小于一個step訓(xùn)練的時間時效率最高，GPU算力被充分利用，如下圖所示：

當(dāng)num_workers 大于0開啟多進程度數(shù)據(jù)， 但是讀取一個batch數(shù)據(jù)的時間大于一個step訓(xùn)練的時間時，會出現(xiàn)GPU訓(xùn)練過程等待數(shù)據(jù)拉取，就會出現(xiàn)GPU算力空閑，訓(xùn)練耗時增加，如下圖所示：

由此可見Dateset中的__getitem__函數(shù)非常重要，詳細(xì)分析它的源碼實現(xiàn)后我們發(fā)現(xiàn)，該函數(shù)的耗時主要包含2段時間：

• load_image_time：從磁盤或者遠(yuǎn)程盤上讀取數(shù)據(jù)的耗時。
• transform_image_time：將圖片或文本數(shù)據(jù)進行預(yù)處理的耗時。

3.2.2 解問題 — 設(shè)置合理的參數(shù)很重要

通過上一小節(jié)的分析，訓(xùn)練時相關(guān)參數(shù)的選擇至關(guān)重要?？偨Y(jié)如下：

• batch_size：根據(jù)數(shù)據(jù)量，以及期望訓(xùn)練時長，用戶合理自定義設(shè)置
• 訓(xùn)練環(huán)境（KubeAI Notebook/任務(wù)/流水線節(jié)點）的CPU配置：建議CPU配置為 GPU卡數(shù)*(單GPU卡配置的CPU核數(shù))。
• num_workers：參數(shù)最小設(shè)置為 訓(xùn)練環(huán)境的CPU配置-1，比如：任務(wù)配置為12C時，建議該參數(shù)設(shè)置為11 。另外，num_workers數(shù)值可以適當(dāng)調(diào)大，因為dataset iter中有部分時間是在網(wǎng)絡(luò)或者磁盤IO， 這部分不消耗CPU；但是也不能設(shè)置太大，因為數(shù)據(jù)預(yù)處理部分是CPU密集型任務(wù)，并行進程過多，會造成CPU爭搶從而降低預(yù)處理效率。

優(yōu)化案例一

線上一個基于MMDetection框架（其底層也是調(diào)用PyTorch框架）的CV模型訓(xùn)練任務(wù)，在做參數(shù)調(diào)整之前，單個step耗時不穩(wěn)定，平均在1.12s左右，其中拉取數(shù)據(jù)時長在0.3s左右：

mmengine - INFO - Epoch(train) [2][3050/6005]  time: 1.1128  data_time: 0.1032  
mmengine - INFO - Epoch(train) [2][3100/6005]  time: 1.0193  data_time: 0.0055  
mmengine - INFO - Epoch(train) [2][3150/6005]  time: 1.0928  data_time: 0.3230  
mmengine - INFO - Epoch(train) [2][3200/6005]  time: 0.9927  data_time: 0.2304  
mmengine - INFO - Epoch(train) [2][3250/6005]  time: 1.3224  data_time: 0.5135  
mmengine - INFO - Epoch(train) [2][3300/6005]  time: 1.1044  data_time: 0.3427  
mmengine - INFO - Epoch(train) [2][3350/6005]  time: 1.0574  data_time: 0.2842

調(diào)整參數(shù)之后，單個step耗時穩(wěn)定，平均在0.78 s左右，其中拉取數(shù)據(jù)耗時0.004s，基本可以忽略。

mmengine - INFO - Epoch(train) [1][100/5592]  time: 0.8508  data_time: 0.0049  
mmengine - INFO - Epoch(train) [1][150/5592]  time: 0.7743  data_time: 0.0043  
mmengine - INFO - Epoch(train) [1][200/5592]  time: 0.7736  data_time: 0.0044  
mmengine - INFO - Epoch(train) [1][250/5592]  time: 0.7880  data_time: 0.0044  
mmengine - INFO - Epoch(train) [1][300/5592]  time: 0.7761  data_time: 0.0045

該模型訓(xùn)練任務(wù)，通過上述優(yōu)化調(diào)整，數(shù)據(jù)拉取時間縮短為0，單個step的耗時從原來的1.12s降到0.78s，整體訓(xùn)練時間減少30%（從2天縮短到33小時），效果顯著。

優(yōu)化案例二

線上某個多模態(tài)模型（輸入包含圖片和文字）訓(xùn)練任務(wù)，使用2卡V100訓(xùn)練，參數(shù)調(diào)整如下：

batch_size=32
CPU = 12  ---> 調(diào)整為 24
num-workers = 4  ---> 調(diào)整為 11

調(diào)整后訓(xùn)練300 step總消耗時405s，整體訓(xùn)練時間減少45%左右（從10天縮短到5天左右）。

優(yōu)化案例三

線上某YoloX模型訓(xùn)練任務(wù)，使用單卡A100訓(xùn)練，參數(shù)調(diào)整如下：

batch size ：48  
num-workers = 4 ---> 調(diào)整為 16

調(diào)整后整體訓(xùn)練時長減少80%左右（從10天19小時，縮短至1天16小時）。

3.2.3 數(shù)據(jù)拉取IO瓶頸分析

當(dāng)前，KubeAI平臺為訓(xùn)練場景提供3種存儲介質(zhì)：

• 本地盤：空間小，讀寫性能最好，單盤500G~3T空間可用。
• NAS網(wǎng)絡(luò)存儲：空間大，讀寫性能較差，成本適中。
• CPFS并行文件系統(tǒng)存儲：空間大，讀寫性能好，成本高。

對于小數(shù)據(jù)集，可以先將數(shù)據(jù)一次性拉取到本地盤，然后每個epoch從本地盤來讀數(shù)據(jù)，這樣避免了每一個epoch重復(fù)的從遠(yuǎn)程NAS來拉取數(shù)據(jù)，相當(dāng)于整個訓(xùn)練只需要從遠(yuǎn)程NAS拉取一次數(shù)據(jù)。對于大數(shù)據(jù)集，有2種解決方案：

• 將大數(shù)據(jù)集提前進行resize，存儲比較小的圖片來進行訓(xùn)練，這樣避免了每個epoch都需要resize，而且resize之后，圖片變小，讀取更快。
• 將數(shù)據(jù)集放入并行文件系統(tǒng)CPFS存儲上，提高訓(xùn)練吞吐。實驗表明CPFS 在圖片場景下是NAS盤讀性能的3～6倍。 

3.3 TrainingModel優(yōu)化

數(shù)據(jù)部分優(yōu)化后，訓(xùn)練過程中的主要時間開銷就在GPU訓(xùn)練部分了。目前業(yè)內(nèi)有一些比較成熟的方法可以參考，我們總結(jié)如下。

3.3.1 混合精度訓(xùn)練（AMP）

PyTorch混合精度訓(xùn)練在PyTorch官網(wǎng)有詳細(xì)介紹，以及開啟混合精度訓(xùn)練的方法，可以閱讀這里獲取實現(xiàn)方法。當(dāng)前許多CV訓(xùn)練框架已經(jīng)支持AMP訓(xùn)練，比如：

• MMCV框架中AMP參數(shù)就是開啟混合精度訓(xùn)練的選項。
• Pytorch Vision中也有相關(guān)參數(shù)來開啟AMP訓(xùn)練。

需要說明的是，混合精度訓(xùn)練過程中并不是將所有模型參數(shù)都轉(zhuǎn)為FP16來計算，只有部分做轉(zhuǎn)換。混合精度之所以能加速訓(xùn)練過程，是因為大部分英偉達GPU機型在FP16這種數(shù)據(jù)格式的浮點算力比FP32要快一倍；此外，混合精度訓(xùn)練顯存占用會更小。

scaler = GradScaler()


for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()


        # Unscales the gradients of optimizer's assigned params in-place
        scaler.unscale_(optimizer)


        # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)


        # optimizer's gradients are already unscaled, so scaler.step does not unscale them,
        # although it still skips optimizer.step() if the gradients contain infs or NaNs.
        scaler.step(optimizer)


        # Updates the scale for next iteration.
        scaler.update()

3.3.2 單機多卡數(shù)據(jù)并行訓(xùn)練

Pytorch原生支持多卡數(shù)據(jù)并行訓(xùn)練，詳細(xì)開啟多卡訓(xùn)練的方式參考官方文檔。多卡訓(xùn)練過程中每一張卡的backword計算會多加一次多卡之間集合通訊all-reduce操作，用來計算多張卡上的梯度的平均值。

3.4 自研訓(xùn)練引擎框架kubeai-training-framework

通過前面的分析我們可以看到，雖然PyTorch框架本身已經(jīng)做的很好了，訓(xùn)練方式、參數(shù)支持豐富，但在實際的模型研究和生產(chǎn)過程中，由于模型的差異性、訓(xùn)練數(shù)據(jù)的差異性，以及模型開發(fā)者的經(jīng)驗差異性，PyTorch框架本身的優(yōu)勢不一定能夠發(fā)揮出來。

基于前述分析和實踐，KubeAI平臺開發(fā)了訓(xùn)練引擎框架kubeai-training-framework，幫助模型開發(fā)者更好地匹配訓(xùn)練腳本參數(shù)，快速接入使用合適的訓(xùn)練方式。kubeai-training-framework中包含PyTorch Dataloader優(yōu)化、GPU TrainModel（AMP）提速以及各種功能函數(shù)等。以Dataloader為例，用戶可通過以下方式使用：

import torch
from kubeai_training_framework.dataloader import Dataloader


def train(train_loader, model, criterion, optimizer, epoch):
    train_dataset = .......
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)


    model.train()
    my_train_loader = Dataloader(train_loader)
    input, target = my_train_loader.next()


    while input is not None:
        ## model train 代碼 input, target
        ..........
        input, target = my_train_loader.next()

4、AI Pipeline引擎助力AI業(yè)務(wù)快速迭代

通常模型的開發(fā)可以歸納為如下圖所示的過程：

可以看到，在需求場景確定、第一個模型版本上線之后，模型是需要反復(fù)迭代的，以期望取得更好的業(yè)務(wù)效果。KubeAI平臺在迭代建設(shè)的過程中，逐步上線了Notebook、模型管理、訓(xùn)練任務(wù)管理、推理服務(wù)管理等一個個相對獨立的功能模塊。隨著業(yè)務(wù)需求的不斷變化，模型迭代效率直接影響了業(yè)務(wù)的上線效率，KubeAI平臺建設(shè)了AI Pipeline能力，重點解決AI場景的周期性迭代類需求，提高生產(chǎn)效率。

AI Pipeline是在ArgoWorkflow基礎(chǔ)上做了二次開發(fā)，以滿足模型迭代、推理任務(wù)管理、數(shù)據(jù)處理等對定時需求、任務(wù)啟動觸發(fā)方式、通用模板任務(wù)、指定節(jié)點啟動等需求。AI Pipeline上線之前，一個迭代任務(wù)可能會被配置為多個分散的任務(wù)，維護工作量大，調(diào)試周期長。如下圖是做一個類似任務(wù)需要單獨配置的任務(wù)情況：

AI Pipeline可以將整個工作流設(shè)計成如下圖所示：

Pipeline編排的方式，減少了模型開發(fā)者浪費在重復(fù)工作上的時間，可以將更多的時間投入到模型研究上。同時，通過合理編排任務(wù)，可以對有限的資源進行充分地利用。

5、展望

KubeAI平臺從得物AI業(yè)務(wù)場景的實際需求出發(fā)，以三大核心引擎為建設(shè)目標(biāo)，著力解決AI模型研發(fā)過程中的訓(xùn)練、推理性能問題，以及模型版本迭代過程中的效率問題。

在推理服務(wù)性能上，我們會以kubeai-inference-framework為起點，繼續(xù)在模型量化、算子優(yōu)化、圖優(yōu)化等方面進行深入探索。在模型訓(xùn)練方面，我們會繼續(xù)在圖像數(shù)據(jù)預(yù)處理、Tensorflow GPU訓(xùn)練框架支持、NLP模型訓(xùn)練支持上發(fā)力，以kubeai-training-framework訓(xùn)練引擎框架為接口，為模型開發(fā)者提供更高效、性能更高的訓(xùn)練框架。此外，AI Pipeline引擎上，我們會支持更豐富的預(yù)置模型，以滿足通用數(shù)據(jù)處理任務(wù)、推理任務(wù)等需求。