配套Torch_NPU使用样例

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))]  # batchsize 需要手动设置

# 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 = min_shape, max_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="Ascend310xxx",  #Ascend310xxx为昇腾AI处理器类型，根据服务器设备类型配置
     optimization_level=0)      

# 4. 模型推理 
input_data = torch.ones([1, 3, 224, 224]).to("npu") # 将输入数据加载到NPU上 
results = compiled_module.forward(input_data)  

# 5. 若通过Torch_NPU在NPU上执行了算子，则编译后模型的保存和加载在当前Torch_NPU版本下暂不支持（没有fallback算子时可以正常执行模型的保存和加载）。

#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的推理；推理实现详情请参见《CANN AscendCL应用软件开发指南 (C&C++)》中“快速入门”章节
export_engine(model, "forward", compileSpec, path);  

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