一. 背景介紹
在字節(jié)跳動，基于深度學習的應用遍地開花，工程師關注模型效果的同時也需要關注線上服務一致性和性能，早期這通常需要算法專家和工程專家分工合作并緊密配合來完成，這種模式存在比較高的 diff 排查驗證等成本。
隨著
PyTorch/TensorFlow 框架的流行，深度學習模型訓練和在線推理完成了統(tǒng)一，開發(fā)者僅需要關注具體算法邏輯，調用框架的 Python
API 完成訓練驗證過程即可，之后模型可以很方便的序列化導出，并由統(tǒng)一的高性能 C++ 引擎完成推理工作。提升了開發(fā)者訓練到部署的體驗。
然而，完整的服務通常還存在大量的預處理/后處理等業(yè)務邏輯，這類邏輯通常是把各種輸入經(jīng)過加工處理轉變?yōu)?Tensor，再輸入到模型，之后模型的輸出 Tensor 再加工成目標格式，一些典型的場景如下：
我們的目標就是為以上端到端的過程，提供自動化且統(tǒng)一的訓練、推理方案，減輕人工開發(fā)推理過程、對齊 diff 等一系列問題，實現(xiàn)大規(guī)模的統(tǒng)一部署方案。
二. 核心問題
PyTorch/TensorFlow 等框架相對已經(jīng)解決了模型的訓練/推理統(tǒng)一的問題，因此模型計算本身不存在訓推一體的問題了（算子性能優(yōu)化不在本次討論范圍）。
核心要解決的問題就是：預處理和后處理需要提供高性能訓推一體的方案。
對于此類邏輯，TensorFlow 2.x 提供了 tf.function(還不完善)，PyTorch 提供了 TorchScript，其無一例外都是選擇了原生 Python 語法子集。 但即使強大如此，仍然存在不可忽略的問題：
- 性能：此方案大多基于虛擬機實現(xiàn)，虛擬機方案靈活并且非常可控，但深度學習框架中的虛擬機大多通常性能不夠優(yōu)良。補充說明一下，框架早期都是為
Tensor
計算設計，數(shù)組計算每個算子成本很高，虛擬機的派發(fā)和調度成本可以忽略。但是，移植到程序語言編程層面開銷難以忽略，代碼寫多了就會成為性能瓶頸。據(jù)測試，TorchScript
解釋器性能只有 Python 的 1/5 左右，tf.function 性能更差一些。
- 功能不全：事實上應用到真實場景中，我們仍然可以找出很多
tf.function/TorchScript 不支持的重要功能，比如：自定義的資源不能打包，只能序列化內置類型；字符串只能做 bytes
處理，中文等 unicode 會造成 diff；容器必須同構，不支持自定義類型等等...
再者，還有很多非深度學習任務，比如在自然語言處理中仍然有很多非深度學習的應用或者子任務，如序列標注，語言模型解碼，樹模型的人工特征構造等任務，這些通常具有更靈活的特征范式，但同時都沒有完整實現(xiàn)端到端的訓推一體方案，仍然有大量的開發(fā)以及正確性校驗工作。
為了解決上述問題，我們開發(fā)了一套基于編譯的預處理方案：MATXScript！
三. MATXScript
在深度學習算法開發(fā)中，開發(fā)者通常使用 Python 進行快速迭代和實驗，同時使用 C++ 開發(fā)高性能的線上服務，其中正確性校驗和服務開發(fā)都會成為較重負擔！
MatxScript（https://github.com/bytedance/matxscript）
是一個 Python 子語言的 AOT 編譯器，可以自動化將 Python 翻譯成 C++，并提供一鍵打包發(fā)布功能。使用 MATXScript
可以讓開發(fā)者快速進行模型迭代的同時以較低成本完成高性能服務的部署。
核心架構如下：
- 最底層是純 C++/CUDA 的基礎庫，由高性能算子專家開發(fā)。
- 在基礎庫之上，準守約定封裝出來 Python 的 庫，可以用在 training 過程中。
- 需要 inferencing 時，利用 MATXScript 可以把 Python 代碼，翻譯成對等的 C++ 代碼，編譯成動態(tài)鏈接庫，加上模型及其他依賴的資源，一起打包發(fā)布即可。
其中，編譯器作用非常關鍵，其核心流程如下：
通過以上流程，用戶所編寫的預處理代碼，可以被編譯成 Pipeline 中的一個 JitOp，為了把前后處理和模型聯(lián)動，我們還開發(fā)了 tracing 系統(tǒng)（接口設計上參考了 PyTorch），架構如下：
基于 MATXScript，我們可以訓練和推理使用同一套代碼，大大降低了模型部署的成本。同時，架構和算法得到了解耦，算法同學完全使用 Python 工作即可，架構同學專注于編譯器開發(fā)及 Runtime 優(yōu)化，在字節(jié)跳動，此方案得到了大規(guī)模部署驗證！
四. 小試牛刀
此處以最簡單的英文文本預處理為例，展示一下 MATXScript 如何使用。
目標：把一段英文文本轉成 indexes
- 編寫一個基本的查字典的邏輯
class Text2Ids:
def __init__(self) -> None:
self.table: Dict[str, int] = {
"hello": 0,
"world": 1,
"[UNK]": 2,
}
def lookup(self, word: str)
return self.table.get(word, 2)
def __call__ (self, words: List[str])
return [self.lookup(w) for w in words]
- 編寫 Pipeline
import matx
class WorkFlow:
def __init__(self):
# 此處會進行代碼編譯，Python 代碼自動編譯封裝為 Callable 對象
self.text2ids = matx.script(Text2Ids)()
def process(self, texts):
ids = self.text2ids(texts)
return ids
# test
handler = WorkFlow()
print(handler.process("hello world unknown"))
# output: [0, 1, 2]
- Trace 導出到 磁盤
# dump
mod = matx.trace(handler.process, "hello world")
print(mod.run({"texts": "hello world"}))
mod.save('./my_dir')
# load
mod = matx.load('./my_dir', -1)
print(mod.run({"texts": "hello world"}))
- C++ 加載
#include <string>
#include <vector>
#include <map>
#include <iostream>
#include <matxscript/pipeline/tx_session.h>
using namespace ::matxscript::runtime;
int main()
{
// test case
std::unordered_map<std::string, RTValue> feed_dict;
feed_dict.emplace("texts", Unicode(U"hello world"));
std::vector<std::pair<std::string, RTValue>> result;
const char* module_path = "./my_dir";
const char* module_name = "model.spec.json";
{
// -1 mean cpu
auto sess = TXSession::Load(module_path, module_name, -1);
auto result = sess->Run(feed_dict);
for (auto& r : result) {
std::cout << "key: " << r.first << ", value: " << r.second << std::endl;
}
}
return 0;
}
完整的代碼見：https://github.com/bytedance/matxscript/tree/main/examples/text2ids
小結：以上是一個非常簡單的純 Python 實現(xiàn)的預處理邏輯，且能被一段通用的 C++ 代碼加載運行，下面我們結合模型展示一個實際的多模態(tài)端到端案例！
五. 多模態(tài)案例
此處以圖文多模態(tài)(Bert+Resnet)為例，模型使用 PyTorch 編寫，展示訓練和部署中實際的工作。
- 配置環(huán)境
a. 配置 gcc/cuda 等基礎設施（通常是運維同學已經(jīng)搞定）
b. 安裝 MATXScript 及基于此開發(fā)的基礎庫(text、vision等) - 編寫模型代碼
a. 此處省略，大家可以參考論文或其他開源實現(xiàn)自行搞定 - 編寫預處理代碼
a. text
from typing import List, Dict, Tuple
import libcut
import matx
class Vocabulary:
...
def utf8_decoder(s: List[bytes]):
return [x.decode() for x in s]
class TextNDArrayBuilder:
...
class TextPipeline:
def __init__(self, mode: str = "eval"):
self.mode = mode
self.cut_engine = libcut.Cutter('/path/to/cut_models', ...)
self.vocab = matx.script(Vocabulary)('/path/to/vocab.txt')
self.decoder = matx.script(utf8_decoder)
self.input_builder = matx.script(TextNDArrayBuilder)(self.vocab)
def process(self, text: List[bytes]):
# List[bytes] 是對齊 C++ 的 vector<string>
text: List[str] = self.decoder(text)
words: List[List[str]] = self.cut_engine(text)
batch_ids: List[List[int]] = self.vocab(words)
input_ids, segment_ids, mask_ids = self.input_builder(batch_ids, 32)
if self.mode == "train":
return input_ids.torch(), segment_ids.torch(), mask_ids.torch()
return input_ids, segment_ids, mask_ids
b. vision
from typing import List, Dict, Tuple
import matx
from matx import vision
class VisionPipeline:
def __init__(self,
device_id: int = 0,
mode: str = "eval",
image_size: int = 224,):
self.is_training = mode == 'train'
self.mode = mode
...
def process(self, image,):
if self.is_training:
decode_nds = self.random_crop_decode(image)
flip_nds = self.random_flip(decode_nds)
resize_nds = self.resize(flip_nds)
transpose_nd = self.transpose_norm(resize_nds, vision.SYNC)
else:
decode_nds = self.decode(image)
resize_nds = self.resize(decode_nds)
crop_nds = self.center_crop(resize_nds)
transpose_nd = self.transpose_norm(crop_nds, vision.SYNC)
if self.mode == "trace":
return transpose_nd
return transpose_nd.torch()
- 接入 DataLoader
a. TextPipeline 可以當成一個正常的 Python Class 接入 Dataset 即可
b. VisionPipeline 涉及到 GPU 預處理，更適合按 batch 進行處理，需要自己單獨構造一個 DataLoader（這里埋個點，之后會開源字節(jié)跳動內部基于多線程的 DataLoader） - 加上模型代碼，開始訓練吧
- 導出端到端的 Inference Model
class MultimodalEvalPipeline:
def __init__(self):
self.text_pipe = TextPipeline(mode="eval", ...)
self.vision_pipe = VisionPipeline(mode="eval", ...)
self.torch_model = torch.jit.load('/path/to/multimodal.jit', map_locatinotallow='cuda:0')
self.tx_model_op = matx.script(self.torch_model, device=0)
def eval(self, texts: List[bytes], images: List[bytes])
input_ids, segment_ids, mask_ids = self.text_pipe.process(texts)
images = self.vision_pipe.process(images)
scores = self.tx_model_op(input_ids, segment_ids, mask_ids, images)
return scores
# examples
example_batch_size = 8
text_examples = ['hello, world'.encode()] * example_batch_size
with open('/path/image.jpg', 'rb') as f:
image_example = f.read()
image_examples = [image_example] * example_batch_size
# pipeline instance
pipe = MultimodalEvalPipeline(...)
mod = matx.trace(pipe.eval, text_examples, image_examples)
# test
print(mod.run({"texts": text_examples, "images": image_examples}))
# save
mod.save('/path/to/my_multimodal')
小結：經(jīng)過以上步驟，我們即可完成端到端的訓練&發(fā)布工作，且整個過程是純 Python 代碼完成的，可以完全由算法同學自己控制。當然，如果模型計算本身還有性能問題，也是可以在背后通過自動改圖優(yōu)化工作完成。
注：完整代碼示例見 https://github.com/bytedance/matxscript/tree/main/examples/e2e_multi_modal
六. 統(tǒng)一Server
在上個章節(jié)，我們得到了一個算法同學發(fā)布的模型包，本章節(jié)論述如果用統(tǒng)一的服務進行加載和運行。
完整的 Server 包括：IDL 協(xié)議、Batching 策略、進/線程調度和排布、模型推理...
這里，我們只討論模型推理這塊，其他的都是可以按約定開發(fā)即可。我們以一個 main 函數(shù)來示例模型加載和運行的過程：
#include <string>
#include <vector>
#include <map>
#include <iostream>
#include <matxscript/pipeline/tx_session.h>
using namespace ::matxscript::runtime;
int main()
{
// test case
std::unordered_map<std::string, RTValue> feed_dict;
feed_dict.emplace("texts", List({String("hello world")}));
feed_dict.emplace("images", List({String("......")}));
std::vector<std::pair<std::string, RTValue>> result;
const char* module_path = "/path/to/my_multimodal";
const char* module_name = "model.spec.json";
{
// cuda:0
auto sess = TXSession::Load(module_path, module_name, 0);
auto result = sess->Run(feed_dict);
for (auto& r : result) {
std::cout << "key: " << r.first << ", value: " << r.second << std::endl;
}
}
return 0;
}
以上代碼就是最簡單的一個 C++ 加載多模態(tài)模型的案例，對 Server 開發(fā)的同學來說，只需進行簡單的抽象和約定，即可把上述代碼改造成一個統(tǒng)一的 C++ 模型服務框架。
七. 更多信息
我們是字節(jié)跳動-AML-機器學習系統(tǒng)團隊，致力于為公司提供統(tǒng)一的高性能訓推一體化框架，同時也會通過火山引擎機器學習平臺服務于合作企業(yè)，火山引擎機器學習平臺預計
2023 年起提供 MATX
的相關支持，包括預置鏡像環(huán)境、常用場景的公開樣例、企業(yè)接入和使用過程中的技術保障等，可以達到訓練和推理場景低成本加速和一體化的效果。歡迎在 https://www.volcengine.com/product/ml-platform 詳細了解我們的產(chǎn)品。
