模型的预处理操作应该部署在哪里？

前言

这两天在做一个事情：使用 NVIDIA Triton 将训练好的文本分类模型 BERT + TextCNN 部署到服务器上。部署模型的过程中，发现模型的输入预处理操作，可以部署到客户端上，也可以部署到服务端上。因此，有了本文的标题，模型的预处理操作应该部署在哪里？

对于文本分类模型，预处理操作需要进行分词，将一个句子变成一个一个单词。除了文本分类模型，其他模型可能还需要后处理操作，比如目标检测模型需要 NMS，去除掉重叠的框框。这些模型的预处理、后处理操作应该部署在哪里呢？

对于这个问题，需要仔细考虑 tradeoff，部署到服务端有什么好处、坏处？部署到客户端又有什么好处、坏处？再根据具体模型，具体的部署环境，做出选择。

tradeoff

部署到服务端

优点

1. 从调用者的角度来看，部署到服务端可以获得一个统一的 API 入口，调用者无需关注内部实现。
2. 接着第一点，模型的迭代更新，预处理后处理操作可能发生改变，统一的 API 将无需更新客户端。
3. 服务端性能高、速度快，软件环境齐全。

缺点

1. 增加服务器负载。
2. 原始数据可能很大，经过预处理操作之后，数据传输量变小。传原始数据的时间会更长。

部署到客户端

优点

1. 充分利用客户端性能，减少服务器负载。

缺点

1. 客户端需要知道预处理的逻辑，或者服务端需要为每一种语言提供 SDK。
2. 预处理的速度受到客户端设备的影响，可能预处理的速度太慢，导致整体推理时间变长。
3. 客户端的机器上，不一定有预处理操作需要的环境。

实验

本节通过实验，对比预处理操作放在客户端和服务端的性能差异。在部署前，将 torch 训练好的模型，导出成 ONNX，然后转换成 tensorrt/openvino 模型。另外，模型的预处理操作，依赖了 transformers/tokenizers，需要 pip 安装 tokenizers，然后分别开发客户端和服务端的逻辑。

实验设置 1.1：将预处理放在客户端，服务端部署一个 tensorrt 模型。

实验设置 1.2：将预处理操作部署到服务端，服务端部署 tensorrt 模型和预处理模型，并使用 triton ensemble 特性搭建数据处理流水线。

实验设置 2.1：将预处理放在客户端，服务端部署一个 openvino 模型。

实验设置 2.2：将预处理操作部署到服务端，服务端部署 openvino 模型和预处理模型，并使用 triton ensemble 特性搭建数据处理流水线。

实验结果

send ids 表示客户端预处理后发送 ids，send str 表示客户端直接发送 string。

1.1 和 1.2 的结果：
send ids time cost:  0.006104469299316406 [[-2.2179837  2.3859138]]
send str time cost:  0.007252931594848633 [[-2.2179837  2.3859138]]

2.1 和 2.2 的结果：
send ids time cost:  0.02684760093688965 [[-2.217986   2.3859158]]
send str time cost:  0.027203798294067383 [[-2.217986   2.3859158]]

实验结果分析

我们可以看到，不管是 GPU 还是 CPU，预处理操作放在服务端会变慢一点。初步怀疑是数据传输导致的速度变慢。

进一步实验

将发送到服务端的 string 复制几份，让文本变长。同时也实验了文本只发送一个 "great!" 这样的超短文本。

"great!" 的结果：
send ids time cost:  0.0056743621826171875 [[-1.4489281  1.5539637]]
send str time cost:  0.005866050720214844 [[-1.4489281  1.5539637]]

复制几份的结果：
send ids time cost:  0.009465217590332031 [[-2.2179837  2.3859138]]
send str time cost:  0.012185335159301758 [[-2.2179837  2.3859138]]

这一组实验结果表明：

1. triton ensemble 存在 overhead，预处理操作即使放在服务端，在文本变短、数据量发送变少的情况，也仍然会慢一点。
2. 预处理放在客户端可以减少数据发送量，减少数据传输的延时。

总结

预处理操作放在哪里，要看具体的部署场景。一般来说，部署到服务端的好处是更加统一，方便升级。部署到客户端的好处是，充分利用边缘设备的性能。

实验代码

完整的项目看这个地方，从训练到部署：https://github.com/zzk0/models/tree/main/nlp/text-classification

import time
import numpy as np
import tritonclient.http as httpclient
from tokenizers import Tokenizer


device_to_service = {
    'cpu': ['text_cnn_bert_openvino', 'text_cnn_bert_pipeline_cpu'],
    'gpu': ['text_cnn_bert_tensorrt', 'text_cnn_bert_pipeline'],
}

tokenizer = Tokenizer.from_file('./pretrained/bert-base-cased/tokenizer.json')
tokenizer.enable_truncation(max_length=128)
tokenizer.enable_padding(length=128)


def preprocess(text: str):
    inputs = tokenizer.encode_batch(text)
    inputs_ids = []
    for input in inputs:
        inputs_ids.append(input.ids)
    inputs_ids = np.array(inputs_ids).astype(np.int32)
    return inputs_ids


def send_ids(triton_client, text: str, device: str = 'gpu'):
    input_ids = preprocess(text)
    inputs = []
    inputs.append(httpclient.InferInput('input', input_ids.shape, "INT32"))
    inputs[0].set_data_from_numpy(input_ids, binary_data=False)
    outputs = []
    outputs.append(httpclient.InferRequestedOutput('output', binary_data=False))
    results = triton_client.infer(device_to_service[device][0], inputs=inputs, outputs=outputs)
    output_data0 = results.as_numpy('output')
    return output_data0


def send_text(triton_client, text: str, device: str = 'gpu'):
    text_bytes = []
    for sentence in text:
        text_bytes.append(str.encode(sentence, encoding='UTF-8'))
    text_np = np.array([text_bytes], dtype=np.object_)
    inputs = []
    inputs.append(httpclient.InferInput('input', text_np.shape, "BYTES"))
    inputs[0].set_data_from_numpy(text_np, binary_data=False)
    outputs = []
    outputs.append(httpclient.InferRequestedOutput('output', binary_data=False))
    results = triton_client.infer(device_to_service[device][1], inputs=inputs, outputs=outputs)
    output_data0 = results.as_numpy('output')
    return output_data0



if __name__ == '__main__':
    device = 'gpu'

    triton_client = httpclient.InferenceServerClient(url='127.0.0.1:8000')
    text = ["I went and saw this movie last night after being coaxed to by a few friends of mine. I'll admit that I was reluctant to see it because from what I knew of Ashton Kutcher he was only able to do comedy. I was wrong. Kutcher played the character of Jake Fischer very well, and Kevin Costner played Ben Randall with such professionalism. The sign of a good movie is that it can toy with our emotions. This one did exactly that. The entire theater (which was sold out) was overcome by laughter during the first half of the movie, and were moved to tears during the second half. While exiting the theater I not only saw many women in tears, but many full grown men as well, trying desperately not to let anyone see them crying. This movie was great, and I suggest that you go see it before you judge."]

    _ = send_ids(triton_client, text, device)
    t0 = time.time()
    res = send_ids(triton_client, text, device)
    t1 = time.time()
    print('send ids time cost: ', t1 - t0, res)

    _ = send_text(triton_client, text, device)
    t0 = time.time()
    res = send_text(triton_client, text, device)
    t1 = time.time()
    print('send str time cost: ', t1 - t0, res)