在这里插入图片描述

前言

CANN 体系比较庞大,一口气全学会效率很低。下面是一条可操作的学习路径,分三个阶段,每个阶段都有对应的代码练习。


第一阶段:跑通一个模型(1-2 周)

目标:能把自己训练的 PyTorch 模型编译成 .om,在 NPU 上跑推理。

1. 环境安装

# 安装 CANN 工具包
wget https://ascend-repo.obs.cn-north-4.myhuaweicloud.com/CANN/7.0.RC1/Ascend-cann-toolkit_7.0.RC1_linux-x86_64.run
chmod +x Ascend-cann-toolkit_7.0.RC1_linux-x86_64.run
./Ascend-cann-toolkit_7.0.RC1_linux-x86_64.run --install

# 安装 torch_npu
pip install torch_npu

2. 第一个推理程序

import torch
import torch_npu

# 加载模型
model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)
model = model.npu().eval()

# 准备输入
input_tensor = torch.randn(1, 3, 224, 224).npu()

# 推理
with torch.no_grad():
    output = model(input_tensor)

print(f"Output shape: {output.shape}")
print(f"Top-1 class: {output.argmax()}")

3. 导出 ONNX 并编译

# 导出 ONNX
dummy_input = torch.randn(1, 3, 224, 224).npu()
torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    opset_version=11
)

# 编译成 .om
!atc --model=resnet18.onnx \
    --framework=5 \
    --output=resnet18 \
    --input_shape="input:1,3,224,224"

4. 用 Ascend CL 做推理

import acl

# 初始化
acl.init()
acl.rt.set_device(0)

# 加载模型
model_id = acl.mdl.load_from_file("resnet18.om")

# 准备输入
input_data = input_tensor.cpu().numpy()
input_dataset = acl.mdl.create_dataset()
input_buffer = acl.util.numpy_to_ptr(input_data)
acl.mdl.add_dataset_buffer(input_dataset, input_buffer)

# 推理
output_dataset = acl.mdl.create_dataset()
acl.mdl.execute(model_id, input_dataset, output_dataset)

# 获取输出
output_data = acl.util.ptr_to_numpy(
    acl.mdl.get_dataset_buffer(output_dataset, 0),
    (1, 1000)
)
print(f"Top-1 class: {output_data.argmax()}")

阶段目标:能完整走通 PyTorch → ONNX → .om → 推理 的流程。


第二阶段:优化模型性能(2-4 周)

目标:让模型跑得更快,显存占用更少。

1. 开启混合精度

from torch_npu.contrib import transfer_to_npu

# 开启自动混合精度
model = model.npu()
model = transfer_to_npu(model)

# 或者用 AMP
scaler = torch.npu.amp.GradScaler()
with torch.npu.amp.autocast():
    output = model(input_tensor)

2. 用 AOE 做自动调优

# 调优 ResNet18
aoe --model=resnet18.onnx \
    --framework=5 \
    --job_type=2 \
    --mode=rl \
    --output=resnet18_optimized

调优后推理速度通常能提升 15-30%。

3. 算子融合验证

# 对比融合前后的节点数量
import onnx

model_before = onnx.load("resnet18.onnx")
print(f"融合前节点数: {len(model_before.graph.node)}")

model_after = onnx.load("resnet18_optimized.onnx")
print(f"融合后节点数: {len(model_after.graph.node)}")

4. 动态 Shape 支持

# 编译支持动态 batch 的模型
!atc --model=resnet18.onnx \
    --framework=5 \
    --output=resnet18_dynamic \
    --dynamic_batch_size="1,2,4,8,16"
# 推理时用不同 batch size
model = torch.jit.load("resnet18_dynamic.om", map_location="npu:0")

for batch in [1, 4, 8, 16]:
    input_tensor = torch.randn(batch, 3, 224, 224).npu()
    output = model(input_tensor)
    print(f"Batch {batch}: output shape = {output.shape}")

阶段目标:能用 AOE 调优、开启混合精度、支持动态 shape。


第三阶段:自定义算子开发(4-8 周)

目标:能给 CANN 贡献自定义算子,或者优化现有算子。

1. 用 Ascend C 写算子

// custom_op.cpp
#include "kernel_operator.h"

class CustomOp {
public:
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int32_t len) {
        xGm.SetGlobalBuffer((__gm__ half*)x);
        yGm.SetGlobalBuffer((__gm__ half*)y);
        len_ = len;
    }

    __aicore__ inline void Process() {
        // 把数据从 GM 搬到 UB
        xLocal = xGm.GetTensor<int32_t>(0, len_);
        
        // 在 UB 里做计算
        for (int i = 0; i < len_; i++) {
            yLocal[i] = xLocal[i] * 2.0f;  // 示例:乘 2
        }
        
        // 把结果写回 GM
        yGm.SetTensor<int32_t>(yLocal, 0, len_);
    }

private:
    GlobalTensor<half> xGm, yGm;
    LocalTensor<half> xLocal, yLocal;
    int32_t len_;
};

extern "C" __global__ void custom_op(GM_ADDR x, GM_ADDR y, int32_t len) {
    CustomOp op;
    op.Init(x, y, len);
    op.Process();
}

2. 编译算子

# 用 CANN 的算子编译工具
ccec --cce=custom_op.cpp \
    --o=custom_op.o \
    --arch=davinci

3. 在 PyTorch 里调用自定义算子

import torch
import torch_npu
import ctypes

# 加载自定义算子库
lib = ctypes.CDLL("./custom_op.so")

# 注册算子
@torch.npu.custom_op("custom::my_op")
def my_op(input: torch.Tensor) -> torch.Tensor:
    output = torch.empty_like(input)
    lib.custom_op(
        input.data_ptr(),
        output.data_ptr(),
        input.numel()
    )
    return output

# 使用
input_tensor = torch.randn(1024, 1024).npu()
output = my_op(input_tensor)

4. 算子性能验证

import time

def benchmark(op, input_tensor, iterations=100):
    # 预热
    for _ in range(10):
        op(input_tensor)
    
    torch.npu.synchronize()
    start = time.time()
    
    for _ in range(iterations):
        op(input_tensor)
    
    torch.npu.synchronize()
    end = time.time()
    
    return (end - start) / iterations * 1000  # ms

# 对比自定义算子和原生算子
custom_time = benchmark(my_op, input_tensor)
native_time = benchmark(lambda x: x * 2, input_tensor)

print(f"自定义算子: {custom_time:.2f} ms")
print(f"原生算子: {native_time:.2f} ms")

阶段目标:能独立开发 Ascend C 算子,并在 PyTorch 里调用。


第四阶段:生产部署(持续)

目标:把模型部署到生产环境,处理真实流量。

1. 服务化部署

# 用 TorchServe 部署
import torchserve

# 打包模型
torch-model-archiver \
    --model-name=resnet18 \
    --version=1.0 \
    --serialized-file=resnet18.pt \
    --handler=image_classifier \
    --export-path=models

# 启动服务
torchserve --start \
    --model-store=models \
    --models resnet18=resnet18.mar

2. 性能监控

import psutil
import torch

def monitor_npu():
    """监控 NPU 使用率"""
    # 用 Ascend CL 查询 NPU 状态
    info = torch.npu.mem_get_info()
    print(f"显存使用: {info['used']/1024/1024:.1f} MB")
    print(f"显存总量: {info['total']/1024/1024:.1f} MB")
    
    # 用 psutil 监控 CPU
    print(f"CPU 使用率: {psutil.cpu_percent()}%")

# 推理时监控
for i in range(100):
    output = model(input_tensor)
    if i % 10 == 0:
        monitor_npu()

3. 批量推理优化

# 批量推理
def batch_inference(model, dataloader, batch_size=16):
    results = []
    
    for batch in dataloader:
        # 拼成 batch
        batch_tensor = torch.stack(batch).npu()
        
        # 一次推理
        with torch.no_grad():
            outputs = model(batch_tensor)
        
        # 收集结果
        results.extend(outputs.cpu())
    
    return results

学习资源推荐

官方文档

# 必读文档列表
docs = [
    "CANN 开发指南",          # 基础概念
    "torch_npu API 参考",     # PyTorch 适配层
    "Ascend C 算子开发指南",  # 自定义算子
    "GE 图优化原理",          # 深入理解
    "AOE 调优指南",          # 性能优化
]

代码仓库

# 必看代码仓库
repos = [
    "cann/samples",                    # 样例代码
    "cann/ops-adv",                   # 优化算子
    "cann/ascend-transformer-boost",  # Transformer 优化
    "ascend/pytorch",                 # torch_npu 源码
]

练习项目

  1. 入门:把 ResNet50 跑在 NPU 上,对比 GPU 速度
  2. 进阶:用 AOE 调优 YOLOv5,记录调优前后的性能
  3. 实战:给 ops-adv 贡献一个自定义算子
  4. 生产:用 TorchServe 部署一个图像分类服务

参考资源

  • CANN 官方文档:https://www.hiascend.com/document/detail/zh/CANN/
  • torch_npu 仓库:https://gitee.com/ascend/pytorch
  • Ascend C 教程:https://www.hiascend.com/document/detail/zh/CANN/
  • CANN 样例代码:https://atomgit.com/cann/samples
  • 昇腾社区:https://www.hiascend.com/community

总结

CANN 的学习路径分四个阶段:第一阶段跑通模型(PyTorch → ONNX → .om),第二阶段优化性能(混合精度、AOE 调优、算子融合),第三阶段开发自定义算子(Ascend C),第四阶段生产部署(服务化、监控、批量推理)。每个阶段都有对应的代码练习,边学边练效率最高。学习资源以官方文档为主,代码仓库为辅,遇到具体问题再查文档。

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐