ONNX 模型的量化
我们都希望从代码中榨取更多的性能，对吧？
在现代，充斥着需要大量计算资源的复杂机器学习算法，因此，榨取每一点性能至关重要。
传统上，机器学习算法是在具有支持大量并行计算能力的 GPU 上进行训练的。但是，当部署我们训练过的模型进行推理时，将所有东西都部署在能够访问这种高端 GPU 的机器上是不切实际的
1、实际示例
让我们举一个例子来更好地理解我们正在解决的问题。
假设我们正在为神经机器翻译做一个机器学习组件。虽然你可以用任何机器学习模型来做到这一点，但我们以谷歌的 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 页面上找到。
原文链接：Blazing Fast Inference with Quantized ONNX Models
BimAnt翻译整理，转载请标明出处