ONNX

Polar_bear
413 阅读6分钟

ONNX开放神经网络交换

简要说明：

  • 在计算机视觉中，通常使用TensorRT优化模型推断过程，尤其在NVIDIA GPU上，可以显著提高推断性能；在NLP领域，使用TensorRT并不是一个常见的做法，这是因为NLP模型的推断通常有CPU或者GPU完成，而不像CV那样直接利用TensorRT进行硬件加速，NLP模型的架构和需求也不同。
  • NLP中，常见的部署方法包括将整个模型转换为ONNX格式，然后使用ONNX Runtime进行推断。
  • NLP部署通常更注重模型的端到端流程、跨平台兼容性和易用性，而不像CV那样专注于硬件加速。

0 参考：

python >> onnx: zhuanlan.zhihu.com/p/498425043

[模型部署]：深度学习模型部署（已更Pytorch篇）：blog.csdn.net/qq_41897558…

深度学习模型部署篇——从0部署深度学习分类模型（一）：juejin.cn/post/726720…

netron 神经网络可视化教程：zhuanlan.zhihu.com/p/431445882

1 Python-ONNX-TensorRT：

  1. 训练阶段：

    • 使用 PyTorch 框架在支持 NVIDIA GPU 的系统上进行训练，并使用 CUDA 来加速训练过程。
    • 训练完成后，你可以选择将模型保存为 PyTorch 模型（包含 GPU 参数）。

  2. 推断阶段：

    • 在推断阶段，可以加载训练好的 PyTorch 模型，并使用 CUDA 进行 GPU 加速的推断。这样可以充分利用 GPU 的并行计算能力，提高推断性能。

  3. 模型保存：

    • 可以选择将训练好的模型以两种方式保存：
      • 以 GPU 训练的方式保存：这样保存的模型将包含 GPU 参数。
      • 以 CPU 训练的方式保存：这样保存的模型将只包含 CPU 参数。

  4. 转换为 ONNX 格式：

    • 使用 PyTorch 的 torch.onnx.export

  5. TensorRT 优化：

    • 使用 NVIDIA TensorRT 将 ONNX 格式的模型转换为 TensorRT 格式。这一步骤会应用 TensorRT 的硬件优化技术，提高推断性能。

Key

  • 将模型转换成onnx格式时，模型形式是GPU还是CPU的都行，不影响
  • tensorrt只用于推理阶段，不用于训练阶段
  • CUDA是invida提供的并行计算平台和编程模型，cuDNN时nvidia提供的深度学习库，深度学习框架调cuDNN来执行底层深度学习计算，cuDNN是CUDA的一部分，提供了深度学习所需的一些优化功能
  • pytorch在训练阶段通常会充分利用动态图计算和自动微分的优势，TensorRT并没有针对这些训练框架做直接的支持，它可以进行精确度的降低，减少推断时的计算量，从而提高推断性能。
import torch 
import onnx 
import onnxruntime as ort 
print('PyTorch 版本', torch.**version**) 
print('ONNX 版本', onnx.**version**) 
print("CUDA 是否能用？ ", torch.cuda.is\_available()) 
print("GPU 型号： ", torch.cuda.get\_device\_name(0)) 
print("torch version: ", torch.**version**) 
print("CUDA version:", torch.version.cuda) 
print("cuDNN version:", torch.backends.cudnn.version())

2 python 转换为onnx示例：

  • 其中model表示要转换的pytorch模型；x是模型的任意一组输入，注意这个输入的shape要和模型输入相匹
import torch 
from transformers import BertTokenizer, BertForSequenceClassification 
device = torch.device("cuda" if torch.cuda.is\_available() else "cpu")) # 加载预训练的 BERT 模型和分词器 
model = BertForSequenceClassification.from\_pretrained('bert-base-uncased') 
tokenizer = BertTokenizer.from\_pretrained('bert-base-uncased') # 输入文本 
text = "Hello, how are you doing today?" # 使用分词器进行 tokenization，并转换为模型输入格式 inputs = tokenizer(text, return\_tensors='pt') # 导出 ONNX 模型 
torch.onnx.export( 
        model, # PyTorch 模型 
        (inputs\['input\_ids'], inputs\['attention\_mask']), # 模型输入 "bert\_model.onnx", # 输出文件路径 
        verbose=True, 
        input\_names=\['input\_ids', 'attention\_mask'], # 模型输入名称 output\_names=\['output'], # 模型输出名称 
        dynamic\_axes={'input\_ids': {0: 'batch\_size'}, 'attention\_mask': {0: 'batch\_size'}, 'output': {0: 'batch\_size'}}, # 动态轴 
        )

3 torch.onnx.export()参数详解：

torch.onnx.export(
       model, # 要导出的 PyTorch 模型 
       args, # 模型输入参数，可以是单个 Tensor 或包含多个 Tensor 的元组或列表 
       f, # 导出的 ONNX 文件路径或文件对象 
       export\_params=True, # 是否导出模型的权重参数 
       verbose=False, # 是否输出详细信息 
       training=False, # 是否导出为训练模式 
       input\_names=None, # 模型输入的名称，可以是字符串或包含多个字符串的列表
       output\_names=None, # 模型输出的名称，可以是字符串或包含多个字符串的列表 
       aten=False, # 是否使用 ATen 运算符，如果为 False，则使用标准 ONNX 运算符 
       export\_raw\_ir=False, # 是否导出未优化的中间表示 
       operator\_export\_type=None, # 导出运算符的类型，可以是 OperatorExportTypes 枚举的值 
       opset\_version=None, # 使用的 ONNX Opset 版本号 
       do\_constant\_folding=True, # 是否对模型进行常量折叠优化 
       dynamic\_axes=None, # 动态轴的字典，用于指定哪些维度是动态的 
       example\_outputs=None, # 示例输出，用于确定导出的模型的形状 
       keep\_initializers\_as\_inputs=None
       ) # 是否将模型的初始化器作为输入

image.png

4 Netron可视化：

!pip install netron 
!netron path/to/your/bert\_model.onnx
  • 可以发现图中没有特征图的维度，只有输入数据的维度
  • 在netron中，如果想看到特征图的维度，需要在导出为onnx的时候，同时加上特征图维度信息。
# 增加维度信息 
model\_file = 'model.onnx' 
onnx\_model = onnx.load(model\_file) onnx.save(onnx.shape\_inference.infer\_shapes(onnx\_model), model\_file)

image.png

image.png

5 train()模式和eval()模式：

  • pytorch中，模型有两种模式，训练模式（model.train())和评估模式（model.eval())。两种模式主要影响到一些特定的层：Batch Normalization 和 Dropout
  • 在训练模式下，BN会使用当前batch的均值和方差进行归一化，而评估模式下，BN会使用整个训练集的均值和方差进行归一化，Dropout也不会执行随机丢弃
  • 在模型导出到ONNX格式时，最好将模式转换为评估模式，可以避免由于训练模式下的特殊操作而导致的不一致性

6 ONNX Runtime & TensorRT 推理引擎：

  1. 支持的框架和模型格式:

    • ONNX Runtime 是一个开放式的深度学习推理引擎，它可以支持多种深度学习框架（如PyTorch、TensorFlow、MXNet等）导出的 ONNX 模型。
    • TensorRT 是 NVIDIA 推出的用于深度学习推理的库，主要支持 NVIDIA GPU。它可以直接优化 TensorFlow 或 PyTorch 模型，并且也能够从 ONNX 模型进行转换。

  2. 平台支持:

    • ONNX Runtime 具有较好的跨平台支持，可以在多种硬件平台上运行，包括 CPU、GPU 和边缘设备。
    • TensorRT 主要针对 NVIDIA GPU 进行了优化，因此在 NVIDIA GPU 上表现更为出色。

  3. 部署和集成:

    • ONNX Runtime 更灵活，因为它是一个通用的深度学习推理引擎，适用于多种硬件和场景。
    • TensorRT 专注于 NVIDIA GPU 上的优化，适用于需要充分利用 NVIDIA GPU 性能的场景。

  4. 性能和优化:

    • TensorRT 在 NVIDIA GPU 上进行了更深入的优化，通常在这个硬件平台上表现更优。
    • ONNX Runtime 也有一些优化策略，但可能不如 TensorRT 针对 NVIDIA GPU 那么深入。

  5. 部署环境:

    • ONNX Runtime 更适合跨平台和多硬件环境的部署，因此可以更灵活地应用于不同的场景。
    • TensorRT 更适合专注于 NVIDIA GPU 的部署，特别是在需要充分利用 GPU 性能的场合。
import onnxruntime 
ort\_session = onnxruntime.InferenceSession('bert\_model.onnx') 
# ONNX Runtime 输入 
ort\_inputs = {'input': input\_tensor} 
# ONNX Runtime 输出 
pred\_logits = ort\_session.run(\['output'], ort\_inputs)\[0] 
pred\_logits = torch.tensor(pred\_logits) 

""" 这里我们把input\_tensor先构建成一个字典，该字典的键input要和
torch.onnx.export函数中输入名称一   致。 ort\_session.run方法
的参数有两个，第一个是输出张量名称，也就是output，这个也和我们前面
torch.onnx.export函数中输出名称一致。 """
avatar
文章
阅读
粉丝