简介

概述

MindIE Torch是昇腾推理引擎组件中针对Pytorch框架模型的推理加速插件。Pytorch框架上训练的模型利用MindIE Torch提供的简易C++/Python接口，少量代码即可完成模型迁移，实现高性能推理。

支持的特性

• TorchScript路线
  • 支持对torch.jit.trace/script导出的TorchScript模型进行编译优化，生成可直接在昇腾NPU设备加速推理的TorchScript模型。
  • 支持静态输入和动态输入，动态输入分为动态Dims和ShapeRange两种模式。
  • 编译优化时支持混合精度、FP32以及FP16精度策略。
  • 支持异步推理和异步数据拷贝。
  • 支持自动分图功能，自动识别当前不支持的算子并将其fallback至CPU或torch_npu执行。
  • 支持手动分图功能，可将指定算子fallback至CPU或torch_npu执行。
  • 支持与torch_npu配套使用。
  • 支持多语言API（C++、Python）。
• torch.export路线
  • 支持对torch.export导出的ExportedProgram进行编译优化，生成可直接在昇腾NPU设备加速推理的nn.Module模型。
  • 支持静态输入和动态ShapeRange输入。
  • 编译优化时支持混合精度、FP32和FP16精度策略。
  • 支持异步推理和异步数据拷贝。
  • 支持自动分图功能，自动识别当前不支持的算子并将其fallback至CPU或torch_npu执行。
  • 支持手动分图功能，可将指定算子fallback至CPU或torch_npu执行。
  • 支持与torch_npu配套使用。
  • 支持Python API。
• torch.compile路线
  • 提供了名为mindie的编译后端，支持在推理时对torch.compile生成的GraphModule进行即时编译优化，使得模型可以在昇腾NPU设备上加速推理。
  • 支持torch.compile静态和动态两种场景。
  • 编译优化时支持混合精度、FP32和FP16精度策略。
  • 支持异步推理和异步数据拷贝。
  • 支持自动分图功能，自动识别当前不支持的算子并将其fallback至torch_npu执行。
  • 支持手动分图功能，可将指定算子fallback至torch_npu执行。
  • 支持Python API。
    

    由于torch.compile的即时编译特点，当前版本下使用MindIE Torch提供的编译后端时，需配套torch_npu一起使用以获得最佳体验。