模型编译、推理和保存样例

为了确保cpp文件可以正确找到Torch，可参考如下方式设置：

在cmake命令的DCMAKE_PREFIX_PATH参数中添加Torch的安装路径。
在CMakeLists.txt中添加find_package(Torch REQUIRED)。

Python开发环境

编译推理TorchScript模型的样例：

# 请务必先导入Torch，再导入MindIE Torch
import torch 
import mindietorch

# 1. 加载原始Torchscript(ts)模型
model = torch.jit.load(model_path)
model.eval() 

# 2. 构造编译配置
# 2.1 纯静态shape，单输入
inputs = [mindietorch.Input((batchsize, 3, 224, 224))] # batchsize 自行写入

# 2.2 纯静态shape，多输入
inputs = [mindietorch.Input((batchsize, 3, 224, 224)), mindietorch.Input((batchsize, 3, 224, 224))] 

# 2.3 动态分档
inputs = []
inputs_gear_1 = [mindietorch.Input((1, 3, 224, 224))]  # 1档
inputs_gear_2 = [mindietorch.Input((8, 3, 224, 224))]  # 2档
inputs_gear_3 = [mindietorch.Input((32, 3, 224, 224))]  # 3档
inputs.append(inputs_gear_1)
inputs.append(inputs_gear_2)
inputs.append(inputs_gear_3) 

# 2.4 input shape range
min_shape = (1, 3, 224, 224)
max_shape = (32, 3, 224, 224)
inputs = []
inputs.append(mindietorch.Input(min_shape = min_shape, max_shape= max_shape)) # inputs 作为最终的模型输入设置。

# 3. 编译优化
compiled_module = mindietorch.compile(    
     model,
     inputs=inputs,
     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=[],
     soc_version="Ascend310xxx",  # Ascend310xxx为昇腾AI处理器类型，根据服务器设备类型配置。
     optimization_level=0)     

# 4. 模型推理
mindietorch.set_device(0)
input_data = torch.ones([1, 3, 224, 224]).to("npu")
results = compiled_module.forward(input_data) 

# 5. ts模型保存加载
compiled_module.save("CompiledModel.ts")
reload_module = torch.jit.load("CompiledModel.ts") 

# 6. 保存OM离线模型，框架推理插件不支持OM的推理；推理实现详情请参见《CANN  AscendCL应用软件开发指南 (Python)》中“快速入门”章节。
output_engine = mindietorch.export_engine(
     model,
     "forward",
     inputs=inputs,
     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=[],
     soc_version="Ascend310xxx",  # Ascend310xxx为昇腾AI处理器类型，根据服务器设备类型配置。
     optimization_level=0)
with open('engine.om', 'wb') as file:     
    file.write(output_engine)

编译推理torch.export路线模型的样例：

# 请务必先导入Torch，再导入MindIE Torch
import torch
from torch._export import export, dynamic_dim 
import mindietorch

# 1. 加载原始PyTorch模型(nn.Module)
model = torch.load(model_path)
model.eval() 

# 2. 构造编译配置
# 2.1 纯静态shape，单输入
inputs = [mindietorch.Input((batchsize, 3, 224, 224))] 

# 2.2 纯静态shape，多输入
inputs = [mindietorch.Input((batchsize, 3, 224, 224)), mindietorch.Input((batchsize, 3, 224, 224))] 

# 2.3 input shape range
min_shape = (1, 3, 224, 224)
max_shape = (32, 3, 224, 224)
inputs = []
inputs.append(mindietorch.Input(min_shape = min_shape, max_shape= max_shape)) 

# 3. 编译优化
# 3.1 使用MindIE Torch进行编译之前可以选择提前将模型导出为ExportedProgram
input_data = torch.ones([32, 3, 224, 224])
# 使用PyTorch导出模型时，假如输入动态的，则还需要构造constraints
constraints = [
    dynamic_dim(input_data, 0) <= 32,
    dynamic_dim(input_data, 0) >= 1,
]
# export的更多详细用法请参考PyTorch的官方文档
model_ep = export(model, args=(input_data,), constraints=constraints)

# 3.2 调用MindIE Torch的compile方法编译
compiled_module = mindietorch.compile(    
     model_ep,
     inputs=inputs,
     precision_policy=mindietorch.PrecisionPolicy.FP16,
     soc_version="Ascend310xxx",    # Ascend310xxx为昇腾AI处理器类型，根据服务器设备类型配置
     ir="dynamo")     

# 4. 模型推理
mindietorch.set_device(0)
results = compiled_module(input_data.to("npu")) 

# 5. 编译后模型支持通过torch.save和torch.load进行保存和加载
torch.save(compiled_module, "./compiled.pt")        
reloda_compiled_model = torch.load("./compiled.pt")

编译推理torch.compile路线模型的样例：

# 请务必先导入Torch和torch_npu后，再导入MindIE Torch
import torch
import torch_npu
import mindietorch

# 设置所使用的npu
mindietorch.set_device(0)

# 1. 加载原始PyTorch模型(nn.Module)并将其放到npu设备上
model = torch.load(model_path)
model = model.eval().to("npu")

# 2. 对原始模型进行包装
# 2.1 可以选择以torch.compile作为入口
# 构造一个字典来表示MindIE Torch编译模型时的设置。字典的key为具体设置的名称，val为该设置的取值。
# key的命名请参考mindietorch.compile章节，val的取值请参考mindietorch.compile章节中除module、inputs和ir参数之外的参数。
backend_kwargs = {
    "precision_policy": mindietorch.PrecisionPolicy.FP16,
    "soc_version": "Ascend310xxx",   # Ascend310xxx为昇腾AI处理器类型，根据服务器设备类型配置
}
opt_model = torch.compile(
    model,
    backend="mindie",
    dynamic=False, # 根据模型是否需要支持动态输入设置
    options=backend_kwargs,
)

# 2.2 也可以选择以mindietorch.compile作为入口
opt_model = mindietorch.compile(    
    model,
    precision_policy=mindietorch.PrecisionPolicy.FP16,
    soc_version="Ascend310xxx",      # Ascend310xxx为昇腾AI处理器类型，根据服务器设备类型配置
    ir="torch_compile",
    enable_dynamic_torch_compile=False, # 根据模型是否需要支持动态输入设置
)     

# 3. 模型推理的时候会触发模型的编译
input_data = torch.ones([1, 3, 224, 224])
input_data = input_data.to("npu")
results = opt_model(input_data) 

# 4. 由于PyTorch的dynamo的限制，torch.compile路线编译后的模型不支持落盘和加载，若有需要保存原始模型即可。
torch.save(model, "./model.pt")

C++开发环境

编译推理TorchScript模型的样例：

#include <torch/torch.h> 
#include <torch/script.h> 
#include "torch/csrc/jit/api/module.h"  
#include "torch_aie.h"  

// 1. 加载原始ts模型 
torch::jit::script::Module model = torch::jit::load(model_path); 
model.eval();  

// 2. 构造编译配置 
// 2.1 纯静态
shape std::vector<torch_aie::Input> inputs; 
inputs.emplace_back(torch_aie::Input(dims_1, torch_aie::DataType::INT32, torch_aie::TensorFormat::ND));  // input1 
inputs.emplace_back(torch_aie::Input(dims_2, torch_aie::DataType::INT32, torch_aie::TensorFormat::ND));  // input2 
torch_aie::torchscript::CompileSpec compile_spec(inputs);  

// 2.2 动态分档 
std::vector<torch_aie::InputProfile> inputs_gear; 
torch_aie::InputProfile inputs_gear_1;  // 1档 
torch_aie::InputProfile inputs_gear_2;  // 2档 
torch_aie::InputProfile inputs_gear_3;  // 3档 
inputs_gear_1.emplace_back(torch_aie::Input({1, 3, 640, 640}, torch_aie::DataType::FLOAT, torch_aie::TensorFormat::NCHW));
inputs_gear_2.emplace_back(torch_aie::Input({8, 3, 640, 640}, torch_aie::DataType::FLOAT, torch_aie::TensorFormat::NCHW)); 
inputs_gear_3.emplace_back(torch_aie::Input({32, 3, 640, 640}, torch_aie::DataType::FLOAT, torch_aie::TensorFormat::NCHW)); 
inputs_gear.emplace_back(inputs_gear_1); 
inputs_gear.emplace_back(inputs_gear_2); 
inputs_gear.emplace_back(inputs_gear_3);  

torch_aie::torchscript::CompileSpec compileSpec(inputs_gear);  

// 2.3 input shape range 
std::vector<int64_t> minShape = { 1, 3, 640, 640 }; 
std::vector<int64_t> maxShape = { 32, 3, 640, 640 }; 
std::vector<torch_aie::Input> inputs; 
inputs.emplace_back(torch_aie::Input(minShape, maxShape, torch_aie::DataType::FLOAT, torch_aie::TensorFormat::NCHW)); 
torch_aie::torchscript::CompileSpec compileSpec(inputs);  

// 3. 编译参数配置
compileSpec.precision_policy = torch_aie::PrecisionPolicy::PREF_FP32;
compileSpec.truncateLongAndDouble = true;
compileSpec.requireFullCompilation = false;
compileSpec.soc_version = "Ascend310xxx";   //Ascend310xxx为昇腾AI处理器类型，根据服务器设备类型配置
compileSpec.allowTensorRepalceInt = false;
compileSpec.minBlockSize = 1;
// 以add算子为例
compileSpec.torchExecutedOps = {"aten::add"};
compileSpec.optimiazationLevel = 0;
// 4. 编译优化 
auto compiled_module = compile(model, compile_spec);  

// 5. 模型推理 
torch_aie::set_device(1);  // 设置运行ID，不设置即为编译时的DeviceId  

std::vector<torch::jit::IValue> inputs_ivalues; 
inputs_ivalues.emplace_back(at::randn(shape_1, dtype)); 
inputs_ivalues.emplace_back(at::randn(shape_2, dtype)); 
auto results = compiled_module.forward(inputs_ivalues);  

// 5. ts模型保存加载 
compiled_module.save("CompiledModel.ts"); 
torch::jit::script::Module reload_module = torch::jit::load("CompiledModel.ts");  

// 6. 保存OM离线模型，框架推理插件不支持OM的推理；推理实现详情请参见《CANN AscendCL应用软件开发指南 (C&C++)》中“快速入门”章节。
export_engine(model, "forward", compile_spec, path);  

// 7. 去初始化 
torch_aie::finalize(); // 程序退出前需显式调用去初始化接口，否则可能会导致程序退出异常

父主题： 样例参考