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 训练参考

相关推荐
m沐沐14 天前
【计算机视觉】OpenCV 模板匹配银行卡数字识别---下
人工智能·python·opencv·计算机视觉·pycharm·numpy
云和数据.ChenGuang16 天前
大模型厂商常用的数据库有哪些?
数据库·人工智能·pytorch·深度学习·numpy
MATLAB代码顾问17 天前
Python NumPy数值计算核心指南
开发语言·python·numpy
FBI HackerHarry浩17 天前
解决pip 安装 numpy 时元数据生成失败
numpy·pip
m沐沐19 天前
【计算机视觉】OpenCV 模板匹配银行卡数字识别---上
人工智能·后端·python·opencv·计算机视觉·pycharm·numpy
iRayCheung21 天前
virtualbox安装的ubuntu系统跑numpy报错
linux·ubuntu·numpy
SilentSamsara22 天前
scikit-learn 工作流工程化:Pipeline、ColumnTransformer 与自定义转换器
开发语言·人工智能·python·机器学习·青少年编程·numpy·scikit-learn
SilentSamsara23 天前
NumPy 进阶:广播机制、ufunc 与向量化计算的工程实践
开发语言·python·青少年编程·性能优化·numpy
DogDaoDao23 天前
【第 04 篇】列表与元组 —— 序列类型核心详解
人工智能·python·深度学习·神经网络·机器学习·conda·numpy
zyl8372123 天前
Python NumPy 学习
python·学习·numpy