配套Torch_NPU使用时,将特定算子fallback至Torch_NPU执行的功能目前仅编译TorchScript模型时支持。
Python开发环境样例如下所示:
import torch
#导入PTA,并设置NPU设备
import torch_npu
torch_npu.npu.set_device("npu:0")
import mindietorch
# 1. 加载原始TorchScript(ts)模型,并将模型加载到NPU上
model = torch.jit.load(model_path)
model.eval().to("npu")
# 2. 构造编译配置
# 2.1 纯静态shape
input = [mindietorch.Input((batchsize, 3, 224, 224))]
# 2.2 动态分档
inputs = []
inputs_gear_1 = mindietorch.Input((1, 3, 640, 640)) # 1档
inputs_gear_2 = mindietorch.Input((8, 3, 640, 640)) # 2档
inputs_gear_3 = mindietorch.Input((32, 3, 640, 640)) # 3档
inputs.append(inputs_gear_1)
inputs.append(inputs_gear_2)
inputs.append(inputs_gear_3)
# 2.3 input shape range
min_shape = (1, 3, 640, 640)
max_shape = (32, 3, 640, 640)
inputs = []
inputs.append(mindietorch.Input(min_shape, max_shape))
# 3. 编译优化
compiled_module = mindietorch.compile(
model,
inputs=input,
precision_policy=mindietorch.PrecisionPolicy.FP16,
truncate_long_and_double=True,
require_full_compilation=False,
allow_tensor_replace_int=False,
min_block_size=3,
torch_executed_ops=[ unsupported operator ], # 将框架插件不支持的算子添加到列表中
soc_version="Ascend310P3",
optimization_level=0)
# 4. 模型推理
input_data = torch.ones([1, 1, 5, 5]).to("npu") # 将输入数据加载到NPU上
results = compiled_module.forward(input_data)
# 5. 若通过Torch_NPU在NPU上执行了算子,则编译后模型的保存和加载在当前Torch_NPU版本下暂不支持(没有fallback算子时可以正常执行模型的保存和加载)。
C++开发环境样例如下所示:
#include <torch/torch.h> #include <torch/script.h> #include "torch/csrc/jit/api/module.h" #include "torch_npu/library_npu.h" // include torch_npu #include "torch_aie.h" // 1. 构造npu device,并初始化 auto device = at::Device(at_npu::key::NativeDeviceType, 0); torch_npu::init_npu(device); // 2. 加载原始ts模型,并将模型加载到NPU上 torch::jit::script::Module model = torch::jit::load(modelPath); model.eval().to(device); // 3. 构造编译配置 // 与不使用Torch_NPU用法相同,略 // 4. 编译优化 auto CompiledModel = compile(model, compileSpec); // 5. 模型推理 std::vector<torch::jit::IValue> inputsIvalues; inputs_ivalues.emplace_back(at::randn(shape1, dtype).to(device)); // 输入数据为"npu" tensor inputs_ivalues.emplace_back(at::randn(shape2, dtype).to(device)); auto results = CompiledModel.forward(inputsIvalues); // 6. ts模型保存加载 compiledModule.save("CompiledModel.ts"); torch::jit::script::Module reloadModule = torch::jit::load("CompiledModel.ts"); // 若通过Torch_NPU在NPU上执行了算子,则编译后模型的保存和加载在当前Torch_NPU版本下暂不支持 //(没有fallback算子时可以正常执行模型的保存和加载) // 7. 保存OM离线模型,框架推理插件不支持OM的推理 export_engine(model, "forward", compileSpec, path); // 8. 去初始化 torch_aie::finalize(); // 程序退出前需显式调用去初始化接口,否则可能会导致程序退出异常