asnumpy:NPU 原生的 NumPy 体验

如果你用 Python 做 AI 开发,大概率写过类似这样的代码:用 NumPy 预处理数据,把结果转成 PyTorch Tensor,推理完再转回 NumPy 做后处理。这个过程里数据在 CPU 和 NPU 之间来回拷贝了一次,拖慢了整体速度。

CANN 的 asnumpy 仓库提供了直接在昇腾NPU 上运行的 NumPy 兼容接口。你在 NPU 显存里有一个 Tensor,不需要拷回 CPU 就能做切片、统计、类型转换------asnumpy 会把这些操作翻译成 CANN Runtime 上的算子调用,在 NPU 上执行完再返回结果。


asnumpy 到底是什么

asnumpy 不是 NumPy 的替代品,也不是 CANN 的 Python 封装。它的定位更精确:NPU 张量的 NumPy 兼容视图

当你有一个 CUDA Tensor 在 GPU 上,想拿到它的统计信息,你得先拷到 CPU 再用 NumPy 算。asnumpy 的思路是在 NPU 上直接算------把 np.mean() 映射为 CANN 的 ReduceMean 算子,在 NPU 上执行,只返回一个标量结果给 CPU。

这段代码展示了 asnumpy 的用法:

python 复制代码
import numpy as np
from cann import asnumpy as asnp

# 创建一个 NPU Tensor(内容在 NPU 显存里)
tensor = asnp.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)

# 这些计算在 NPU 上执行,不会拷回 CPU
mean = asnp.mean(tensor, axis=1)
std = asnp.std(tensor)
mask = asnp.where(tensor > 3, tensor, 0)

# 需要结果时再拷回 CPU
result = asnp.to_numpy(mask)
print(result)  # [[0, 0, 0], [4, 5, 6]]

对标的不是 NumPy 的全部 API------asnumpy 覆盖了最常用的操作:创建、索引、切片、数学统计、类型转换、形状变换。不常用的高级功能(像线性代数分解、傅里叶变换)走 CANN 的原生接口。


为什么不能直接等同于 NumPy

asnumpy 的 API 模仿 NumPy,但底层的执行语义不一样。

NumPy 的操作是 CPU 上立即执行的------np.add(a, b) 调用 CPU 的 SIMD 指令当场算完。asnumpy 的操作是异步提交到 NPU Stream 上的------asnp.add(a, b) 把算子提交到 CANN Runtime 的 Stream 队列,函数立即返回,真正的计算在 NPU 上排队执行。

这个差异在你串行多个 asnumpy 操作时感受不到,但跨设备交互时需要注意:

python 复制代码
x = asnp.array([1, 2, 3])
y = asnp.square(x)        # 异步提交到 NPU Stream
z = np.sqrt(asnp.to_numpy(y))  # 等待 NPU 算完,拷回 CPU,再做 NumPy sqrt

to_numpy 是一个隐式的同步点------它会等 NPU 把 square 算完才拷数据。如果不需要立即拿到 CPU 结果,尽量推迟 to_numpy 调用,让更多操作在 NPU 上串成一条流水线。


NPU 张量计算如何加速

asnumpy 的性能收益来自"省掉一次数据搬运"。数据在 NPU 上生成并在 NPU 上消费,全程不碰 CPU 内存。

一个具体的场景:推理结果的 top-k 后处理

python 复制代码
# 传统做法:结果从 NPU 拷回 CPU,用 NumPy 做 top-k
raw_output = aclmdl_get_output(model_id)  # 在 NPU 上
cpu_output = raw_output.to_cpu()          # 拷回 CPU → 约 0.5ms
indices = np.argsort(cpu_output, axis=1)[:, -5:]  # CPU 上算

# asnumpy 做法:在 NPU 上直做
raw_output = aclmdl_get_output(model_id)
indices = asnp.argsort(raw_output, axis=1)[:, -5:]  # NPU 上算
result = asnp.to_numpy(indices)  # 只拷 5 个值回 CPU

传统做法搬运了整个输出 Tensor(几 MB 到几十 MB)。asnumpy 做法在 NPU 上做完筛选后只拷 5 个整数(20 字节)。后处理延迟从毫秒级降到微秒级。


典型应用场景

预处理链在 NPU 上完成。 图像 resize → normalize → HWC→CHW → float 转换,传统做法用 NumPy 在 CPU 上一条龙做完再拷到 NPU。asnumpy 允许你在 NPU 显存中完成预处理,只需要拷原始图像过去一次。

不卸 NPU 的中间分析。 调试阶段经常需要打印 Tensor 的形状、均值、方差来判断有没有 NaN。asnumpy 让这些分析在 NPU 上做,只有分析结论拷回来。

CI 测试中的数据验证。 多框架(PyTorch、MindSpore)在昇腾上的推理结果对齐验证------用 asnumpy 在 NPU 上比较两个 Tensor 的差异,不需要拷回 CPU 再比对。


小结

asnumpy 是 CANN 对 Python 生态的亲和力补强。它没有发明新的计算模型,而是把 NumPy 开发者熟悉的操作搬到了 NPU 上执行。不写 CUDA Kernel、不调 AscendCL API,用普通的 asnp.mean() 就能直接在 NPU 上算 Tensor 统计量------减少了跨设备数据搬运的次数。

asnumpy 仓库

PyTorch NPU 适配


asnumpy 不支持的场景

asnumpy 覆盖了 NumPy 常用功能的 70% 左右。以下场景当前不支持或性能不佳:

  • 复杂广播:高维 Tensor 的复杂广播规则可能退化为逐元素循环,建议手动 reshape 后再操作
  • In-place 修改tensor[0] = 1 这种 in-place 赋值在 NPU 上开销大,asnumpy 会走 Copy-on-Write 路径
  • 混合精度统计:当前只在 FP16 和 FP32 上做过验证,BF16 支持正在开发

这些限制也符合 asnumpy 的定位------覆盖常规场景,极端场景走 CANN 原生 API。


社区态度

asnumpy 在 CANN 开源社区里的活跃度在上升。2025 年以来新增了随机数生成、高级索引和 FFT 操作的支持。社区讨论区也经常看到用户贴出 asnumpy 的使用示例------预处理 10 万张图像,用 asnumpy 在 NPU 上做归一化比 CPU NumPy 快了 8 倍。这类分享驱动了更多开发者在自己的项目里尝试 NPU 原生的数据处理。

如果你在写 Python 推理脚本,asnumpy 是值得一试的依赖------它不会让你扔掉 NumPy 习惯,但能帮你省掉大量 CPU ↔ NPU 之间的数据搬运时间。

参考仓库

asnumpy 仓库
torchtitan-npu 训练参考

相关推荐
毋语天12 小时前
NumPy 完全入门指南:核心数据结构与高效数值计算
数据结构·numpy
Omics Pro1 天前
填补蛋白质组深度学习预处理教学空白
人工智能·python·深度学习·plotly·numpy·pandas·scikit-learn
bloxed1 天前
【AI大模型--NumPy-02】-数组创建与高级索引完全指南
人工智能·numpy
沉下去,苦磨练!3 天前
python的数据分析numpy
python·数据分析·numpy
深兰科技5 天前
深兰科技签约乌兹别克斯坦智慧城市项目,推动中国AI出海规模化
人工智能·beautifulsoup·numpy·智慧城市·fastapi·matplotlib·深兰科技
Cloud_Shy6186 天前
Python 数据分析基础入门:《Excel Python:飞速搞定数据分析与处理》学习笔记系列(第九章 Excel 自动化 上篇)
python·数据分析·excel·numpy·pandas
Cloud_Shy6186 天前
Python 数据分析基础入门:《Excel Python:飞速搞定数据分析与处理》学习笔记系列(第八章 使用读写包操作 Excel 文件 下篇)
python·数据分析·excel·numpy·pandas
Cloud_Shy6187 天前
Python 数据分析基础入门:《Excel Python:飞速搞定数据分析与处理》学习笔记系列(第九章 Excel 自动化 下篇)
python·数据分析·excel·numpy·pandas
石山代码7 天前
Python 数据分析三大库:NumPy + Pandas + Matplotlib
python·数据分析·numpy