1. 背景介紹

在位元組跳動，基於深度學習的應用遍地開花，工程師關注模型效果的同時也需要關注線上服務一致性和效能，早期這通常需要演算法專家和工程專家分工合作並緊密配合來完成，這種模式存在比較高的 diff 排查驗證等成本。

隨著 PyTorch/TensorFlow 框架的流行，深度學習模型訓練和線上推理完成了統一，開發者僅需要關注具體演算法邏輯，呼叫框架的 Python API 完成訓練驗證過程即可，之後模型可以很方便的序列化匯出，並由統一的高效能 C++ 引擎完成推理工作。提升了開發者訓練到部署的體驗。

然而，完整的服務通常還存在大量的預處理/後處理等業務邏輯，這類邏輯通常是把各種輸入經過加工處理轉變為 Tensor，再輸入到模型，之後模型的輸出 Tensor 再加工成目標格式，一些典型的場景如下：

Bert

Resnet

我們的目標就是為以上端到端的過程，提供自動化且統一的訓練、推理方案，減輕人工開發推理過程、對齊 diff 等一系列問題，實現大規模的統一部署方案。

2. 核心問題

PyTorch/TensorFlow 等框架相對已經解決了模型的訓練/推理統一的問題，因此模型計算本身不存在訓推一體的問題了（運算元效能優化不在本次討論範圍）。

核心要解決的問題就是：預處理和後處理需要提供高效能訓推一體的方案。

對於此類邏輯，TensorFlow 2.x 提供了 tf.function(還不完善)，PyTorch 提供了 TorchScript，其無一例外都是選擇了原生 Python 語法子集。但即使強大如此，仍然存在不可忽略的問題：

效能：此方案大多基於虛擬機器實現，虛擬機器方案靈活並且非常可控，但深度學習框架中的虛擬機器大多通常效能不夠優良。補充說明一下，框架早期都是為 Tensor 計算設計，陣列計算每個運算元成本很高，虛擬機器的派發和排程成本可以忽略。但是，移植到程式語言程式設計層面開銷難以忽略，程式碼寫多了就會成為效能瓶頸。據測試，TorchScript 直譯器效能只有 Python 的 1/5 左右，tf.function 效能更差一些。
功能不全：事實上應用到真實場景中，我們仍然可以找出很多 tf.function/TorchScript 不支援的重要功能，比如：自定義的資源不能打包，只能序列化內建型別；字串只能做 bytes 處理，中文等 unicode 會造成 diff；容器必須同構，不支援自定義型別等等...

再者，還有很多非深度學習任務，比如在自然語言處理中仍然有很多非深度學習的應用或者子任務，如序列標註，語言模型解碼，樹模型的人工特徵構造等任務，這些通常具有更靈活的特徵正規化，但同時都沒有完整實現端到端的訓推一體方案，仍然有大量的開發以及正確性校驗工作。

為了解決上述問題，我們開發了一套基於編譯的預處理方案：MATXScript！

3. MATXScript

在深度學習演算法開發中，開發者通常使用 Python 進行快速迭代和實驗，同時使用 C++ 開發高效能的線上服務，其中正確性校驗和服務開發都會成為較重負擔！

MatxScript（http://github.com/bytedance/matxscript）是一個 Python 子語言的 AOT 編譯器，可以自動化將 Python 翻譯成 C++，並提供一鍵打包釋出功能。使用 MATXScript 可以讓開發者快速進行模型迭代的同時以較低成本完成高效能服務的部署。

核心架構如下：

最底層是純 C++/CUDA 的基礎庫，由高效能運算元專家開發。
在基礎庫之上，準守約定封裝出來 Python 的庫，可以用在 training 過程中。
需要 inferencing 時，利用 MATXScript 可以把 Python 程式碼，翻譯成對等的 C++ 程式碼，編譯成動態連結庫，加上模型及其他依賴的資源，一起打包釋出即可。

其中，編譯器作用非常關鍵，其核心流程如下：

通過以上流程，使用者所編寫的預處理程式碼，可以被編譯成 Pipeline 中的一個 JitOp，為了把前後處理和模型聯動，我們還開發了 tracing 系統（介面設計上參考了 PyTorch），架構如下：

基於 MATXScript，我們可以訓練和推理使用同一套程式碼，大大降低了模型部署的成本。同時，架構和演算法得到了解耦，演算法同學完全使用 Python 工作即可，架構同學專注於編譯器開發及 Runtime 優化，在位元組跳動，此方案得到了大規模部署驗證！

小試牛刀

此處以最簡單的英文文字預處理為例，展示一下 MATXScript 如何使用。

目標：把一段英文文字轉成 indexes

編寫一個基本的查字典的邏輯

``` Python class Text2Ids: def init(self) -> None: self.table: Dict[str, int] = { "hello": 0, "world": 1, "[UNK]": 2, }

def lookup(self, word: str) -> int: return self.table.get(word, 2) def call (self, words: List[str]) -> List[int]: return [self.lookup(w) for w in words] ```

編寫 Pipeline

``` Python 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 匯出到磁碟

``` Python

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++ 載入

``` C++

include

using namespace ::matxscript::runtime; int main() { // test case std::unordered_map feed_dict; feed_dict.emplace("texts", Unicode(U"hello world")); std::vector> 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; } ```

完整的程式碼見：http://github.com/bytedance/matxscript/tree/main/examples/text2ids

小結：以上是一個非常簡單的純 Python 實現的預處理邏輯，且能被一段通用的 C++ 程式碼載入執行，下面我們結合模型展示一個實際的多模態端到端案例！

5. 多模態案例

此處以圖文多模態(Bert+Resnet)為例，模型使用 PyTorch 編寫，展示訓練和部署中實際的工作。

配置環境

a. 配置 gcc/cuda 等基礎設施（通常是運維同學已經搞定）

b. 安裝 MATXScript 及基於此開發的基礎庫(text、vision等)
編寫模型程式碼

a. 此處省略，大家可以參考論文或其他開源實現自行搞定
編寫預處理程式碼

a. text

```Python 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 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

```Python 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（這裡埋個點，之後會開源位元組跳動內部基於多執行緒的 DataLoader）
加上模型程式碼，開始訓練吧
匯出端到端的 Inference Model

```Python 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_location='cuda:0') self.tx_model_op = matx.script(self.torch_model, device=0)

def eval(self, texts: List[bytes], images: List[bytes]) -> List[float]: 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') ```

小結：經過以上步驟，我們即可完成端到端的訓練&釋出工作，且整個過程是純 Python 程式碼完成的，可以完全由演算法同學自己控制。當然，如果模型計算本身還有效能問題，也是可以在背後通過自動改圖優化工作完成。

注：完整程式碼示例見 http://github.com/bytedance/matxscript/tree/main/examples/e2e_multi_modal

6. 統一Server

在上個章節，我們得到了一個演算法同學釋出的模型包，本章節論述如果用統一的服務進行載入和執行。

完整的 Server 包括：IDL 協議、Batching 策略、進/執行緒排程和排布、模型推理...

這裡，我們只討論模型推理這塊，其他的都是可以按約定開發即可。我們以一個 main 函式來示例模型載入和執行的過程：

``` C++

include

using namespace ::matxscript::runtime; int main() { // test case std::unordered_map feed_dict; feed_dict.emplace("texts", List({String("hello world")})); feed_dict.emplace("images", List({String("......")})); std::vector> 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++ 載入多模態模型的案例，對 Server 開發的同學來說，只需進行簡單的抽象和約定，即可把上述程式碼改造成一個統一的 C++ 模型服務框架。

7. 更多資訊

我們是位元組跳動-AML-機器學習系統團隊，致力於為公司提供統一的高效能訓推一體化框架，同時也會通過火山引擎機器學習平臺服務於合作企業，火山引擎機器學習平臺預計 2023 年起提供 MATX 的相關支援，包括預置映象環境、常用場景的公開樣例、企業接入和使用過程中的技術保障等，可以達到訓練和推理場景低成本加速和一體化的效果。歡迎在 http://www.volcengine.com/product/ml-platform 詳細瞭解我們的產品。

更多資訊，可訪問 Github：http://github.com/bytedance/matxscript

位元組跳動模型大規模部署實戰

1. 背景介紹

2. 核心問題

3. MATXScript

test

output: [0, 1, 2]

dump

load

include

include

include

include

include

5. 多模態案例

examples

pipeline instance

test

save

6. 統一Server

include

include

include

include

include

7. 更多資訊