我们都希望从代码中榨取更多的性能，对吧？

在现代，充斥着需要大量计算资源的复杂机器学习算法，因此，榨取每一点性能至关重要。

传统上，机器学习算法是在具有支持大量并行计算能力的 GPU 上进行训练的。但是，当部署我们训练过的模型进行推理时，将所有东西都部署在能够访问这种高端 GPU 的机器上是不切实际的

NSDT工具推荐： Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割 

1、实际示例

让我们举一个例子来更好地理解我们正在解决的问题。

假设我们正在为神经机器翻译做一个机器学习组件。虽然你可以用任何机器学习模型来做到这一点，但我们以谷歌的 T5模型为例。

这是一个文本到文本的转换模型，在大型语料库上进行训练以完成多项任务。

Google T5 多任务图

执行此任务的示例代码如下所示：

from transformers import T5Tokenizer, T5ForConditionalGeneration, pipeline

tokenizer = T5Tokenizer.from_pretrained("t5-small")
model = T5ForConditionalGeneration.from_pretrained("t5-small")
pipe = pipeline(
	task='text2text-generation', 
    model=model, 
    tokenizer = tokenizer)
print(pipe('Translate English to French: Hi, How are you ?')) 

# Output: [{'generated_text': 'Bonjour, Comment êtes-vous ?'}]

• 首先，我们从 transformers 包中导入所有必要的类和函数
• 然后我们实例化 tokenizer 和模型
• 我们将 tokenizer 和模型传递给管道辅助函数，
• 从高层次上讲，transformer 架构可以分为两半：编码器和解码器。
• 我们将文本输入句子转换为各自的 token 表示，并将这些 token 提供给 T5 模型的编码器部分，然后这些输出被提供给解码器部分，最后我们使用相同的 tokenizer 将输出 token 转换回其文本表示。
• 所有上述复杂性都很好地包装在我们上面使用的管道函数中

现在，让我们来测量此代码的性能：

from transformers import T5Tokenizer, T5ForConditionalGeneration, pipeline
from timeit import timeit

tokenizer = T5Tokenizer.from_pretrained("t5-small")
model = T5ForConditionalGeneration.from_pretrained("t5-small")
pipe = pipeline(task='text2text-generation', model=model, tokenizer = tokenizer)

globals = dict(map(
    lambda x: (x, eval(x)),
    dir()
))

time_taken = timeit(
    globals = globals,
    number = 5,
    stmt="pipe('Translate English to French: Hi, How are you ?')",
)

print(f"It took about {time_taken} seconds")

我们使用 Python 内置的 time-it 模块来测量代码翻译所需的时间，并重复 5 次以获得平均估计值。

因为我们没有使用任何 GPU 进行推理，所以这些数字肯定取决于机器的原始 CPU 能力。但是我们仍然可以提取相对测量值。

(venv-meta) tarun@Taruns-MacBook-Pro ML % /Users/tarun/ML/venv-meta/bin/python /Users/tarun/ML/
t5_torch.py
It took about 0.4970941249999996 seconds

大约需要 0.5 秒

现在让我们检查一下内存使用情况

from memory_profiler import profile

from transformers import T5Tokenizer, T5ForConditionalGeneration, pipeline
from timeit import timeit

tokenizer = T5Tokenizer.from_pretrained("t5-small")
model = T5ForConditionalGeneration.from_pretrained("t5-small")
pipe = pipeline(task='text2text-generation', model=model, tokenizer = tokenizer)

@profile
def perform_inference():
    return pipe('Translate English to French: Hi, How are you ?')

if __name__ == '__main__':
    perform_inference()

运行脚本我们会发现以下结果：

(venv-meta) tarun@Taruns-MacBook-Pro ML % python -m memory_profiler t5_torch_memory_profile.py
Filename: t5_torch_memory_profile.py
Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    10    847.9 MiB    847.9 MiB           1   @profile
    11                                         def perform_inference():
    12    859.2 MiB     11.3 MiB           1       return pipe('Translate English to French: Hi, How are you ?')

我们可以看到，推理任务需要大约 11MB 的内存

虽然这些数字在功能强大的计算机上还不错，但在低功耗移动设备上可能会更加昂贵，如果有一种方法可以让这些模型的部署成本降低很多计算要求，并且执行速度更快，同时基本保持其准确性，那会怎样呢？

2、了解 ONNX

ONNX 代表开放神经网络交换，它是一种由微软研究院开发的开源技术，用于帮助加速跨框架和语言的机器学习推理。

ONNX 的好处在于它是跨平台和跨语言的。

这意味着你可以将机器学习导出为 ONNX 格式，并且可以用你喜欢的任何语言使用它们。

除此之外，ONNX 还具有许多内置优化功能，可以利用现代硬件级功能来加速推理。

让我们从将模型导出到 ONNX 表示开始。

将 transformer 模型导出为 ONNX 运行时等效项的官方方法是使用optimum。

运行以下命令执行导出：

(venv-meta) tarun@Taruns-MacBook-Pro ML % optimum-cli export onnx \
--model t5-small \
--optimize O3 \
t5_small_onnx

在上面的命令中，我们使用 optimum-cli ，它是 optimal 库的命令行包装器。

我们首先指定要导出的模型，然后指定 ONNX 应该进行的优化级别。

但我们还没有完成，我们可以做得更好！

3、量化ONNX模型

如果我说我对某个结果的确定性为 85.90123456789%，而对同一结果的确定性为 85.88%，这里的小数精度会有所不同吗？

可能不会，这就是量化对机器学习模型的作用。

当我们训练这些模型时，我们通常使用更高的精度来训练它们，例如使用高达 64 位的浮点精度。

但在实践中，我们可以根据我们的偏好将这个精度降低到 32 位、16 位甚至 8 位。

量化是降低权重、偏差和激活的精度的过程，这样它们消耗的内存更少，运行速度更快！

为了量化我们新导出的模型，让我们再次运行 optimal-cli：

(venv-meta) tarun@Taruns-MacBook-Pro ML % optimum-cli onnxruntime \
quantize \
--onnx_model t5_onnx_small \
--arm64 \
--output t5_onnx_small_quantized

...

Saving quantized model at: t5_onnx_small_quantized (external data format: False)
Configuration saved in t5_onnx_small_quantized/ort_config.json
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.

现在我们已经准备好了量化模型。让我们再次检查一下时间消耗：

from optimum.pipelines import pipeline
from transformers import T5Tokenizer
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from timeit import timeit

model = ORTModelForSeq2SeqLM.from_pretrained('t5_onnx_small_quantized')
tokenizer = T5Tokenizer.from_pretrained('t5-small')
pipe = pipeline(task='text2text-generation',model=model, tokenizer=tokenizer)

globals = dict(map(
    lambda x: (x, eval(x)),
    dir()
))

time_taken = timeit(
    globals = globals,
    number = 5,
    stmt="pipe('Translate English to French: Hi, How are you ?')",
)

print(f"It took about {time_taken} seconds")

我们可以看到，我们的逻辑大部分保持不变，唯一的变化是，我们不再加载由 PyTorch 支持的 hugging face transformers T5 模型，而是加载来自 optimal.onnxruntime 的 ONNX 等效模型。

(venv-meta) tarun@Taruns-MacBook-Pro ML % /Users/tarun/ML/venv-meta/bin/python /Users/tarun/ML/
t5_onnx.py
It took about 0.09925400000000018 seconds

只需很少的代码改动，性能就能提升 5 倍。

同样，让我们​​看看内存方面的表现如何：

from memory_profiler import profile
from optimum.pipelines import pipeline
from transformers import T5Tokenizer
from optimum.onnxruntime import ORTModelForSeq2SeqLM

model = ORTModelForSeq2SeqLM.from_pretrained('t5_onnx_small_quantized')
tokenizer = T5Tokenizer.from_pretrained('t5-small')
pipe = pipeline(task='text2text-generation',model=model, tokenizer=tokenizer)

@profile
def perform_inference():
    return pipe('Translate English to French: Hi, How are you ?')

if __name__ == '__main__':
    perform_inference()

除模型初始化部分外，代码再次看起来大致相似：

(venv-meta) tarun@Taruns-MacBook-Pro ML % python -m memory_profiler t5_onnx_quantized_memory_profile.py
Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    10    933.9 MiB    933.9 MiB           1   @profile
    11                                         def perform_inference():
    12    938.7 MiB      4.9 MiB           1       return pipe('Translate English to French: Hi, How are you ?')

我们可以清楚地看到，我们最终消耗的内存少了很多。事实上，它的内存消耗比以前少了近 2.2 倍。正如一位聪明的程序员曾经说过的：

“当谈到扩展时，优化就是游戏的名称”😎

这个故事的所有代码示例都可以在我的 github 页面上找到。

---

原文链接：ONNX 模型的量化 - BimAnt