asnumpy:让你的 NumPy 代码零改动跑在昇腾 NPU 上

有个同事之前跟我说,他写了一套数据预处理的 pipeline,全是 NumPy 写的,后来要迁移到昇腾 NPU 上跑,“感觉天都要塌了”——几千行 NumPy 代码,难道要全部重写?

后来他发现了 asnumpy,这个东西基本上解决了这个问题。不需要改你的 NumPy 代码,直接换一下 import,数据搬上 NPU 跑,原来的语法一个不动。


先搞清楚 asnumpy 是什么

asnumpy 不是 NumPy 的替代品,它是昇腾 NPU 对 NumPy API 的一层兼容封装。你写的是 NumPy 代码,但底层跑在了昇腾 NPU 的向量计算单元上。

这里有个容易混淆的地方:asnumpy 不等于 NumPy 本身。NumPy 是 CPU 上的 Python 数值计算库,asnumpy 是昇腾 NPU 的同名实现,API 签名几乎一致,但底层硬件完全不同。如果你看到有人说"用 asnumpy 替换 NumPy",这个说法是不准确的——正确理解是"用 asnumpy 的接口来写代码,享受 NPU 的加速"。

asnumpy 属于 ops-tensor 仓库的一部分,和 ops-math、ops-nn 这些算子库并列。同一个父项目下,不同的算子能力,asnumpy 主要负责向量级别的数值计算支持。


为什么要用它

标准 NumPy 在 CPU 上跑,数据要先从 NPU 的显存里搬到 CPU 内存,算完再搬回去。这个来回搬运的过程,在数据量大的时候开销非常可观。

举个例子,你有一段图像预处理的代码:

import numpy as np

# 原始 NumPy 版本(CPU)
def preprocess(image):
    image = image.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = (image - mean) / std
    return image

如果你想在昇腾 NPU 上跑这个预处理,通常的做法是把数据从 NPU 搬到 CPU,用 NumPy 算完再搬回去——等于白加速了预处理部分。

asnumpy 的思路是:让这部分计算直接在 NPU 上完成,不用搬运

import ascendnumpy as np  # 直接替换 import

# 同一个函数,一个字不用改,数据全程在 NPU 上
def preprocess(image):
    image = image.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = (image - mean) / std
    return image

这样预处理和模型推理都在 NPU 上,数据零拷贝。这个"零改动"的体验是 asnumpy 最核心的价值。


常用操作速查

asnumpy 的 API 覆盖了 NumPy 的常用子集,不是 100% 全覆盖,但主流操作基本都有。我列几个最常用的:

import ascendnumpy as np

# 张量创建
a = np.zeros((1024, 1024), dtype=np.float32)
b = np.ones((512, 512))
c = np.random.randn(256, 256).astype(np.float16)

# 数学运算
d = np.matmul(a, b)         # 矩阵乘法,昇腾 NPU 硬件加速
e = np.sum(c, axis=0)       # 按列求和
f = np.mean(c, axis=1)      # 按行求均值
g = np.clip(a, 0, 1)         # 截断

# 数据类型转换,asnumpy 特有
h = a.asnpu()   # 转为 NPU 显存格式,提交给模型推理用
i = h.ascpu()   # 从 NPU 拿回 CPU,结果导出用

这里要特别说一说 asnpu()ascpu() 这一对方法。这是 asnumpy 和标准 NumPy 最大的差异——你写的代码可以全程跑在 NPU 上,但总有需要把结果拿回 CPU 的时候,比如保存文件、打印日志、写磁盘。这两个方法就是连接 NPU 和 CPU 数据的桥。

# 典型场景:预处理在 NPU,算完拿回 CPU 存图
result = preprocess(image_npu)
result_cpu = result.ascpu()  # 这步才是真正有开销的,但只跑一次

# 相反方向:CPU 数据批量喂给 NPU
images = [load_image(f) for f in image_paths]  # CPU 上读取
images_npu = [img.asnpu() for img in images]  # 批量上传 NPU

asnumpy 和 catlass 的区别

有人会问:asnumpy 也能算矩阵乘法,catlass 也是做矩阵运算的,它们是什么关系?

这个问题问得好。简单说:asnumpy 是给数据处理用的,catlass 是给算子开发用的

asnumpy 的矩阵乘法是封装好的高层接口,你调一下 np.matmul,它内部帮你把数据切块、调度硬件资源,你不需要关心底层怎么实现的。catlass 是昇腾算子模板库,提供了 GEMM(通用矩阵乘法)的底层开发框架,你需要自己写 tiling 逻辑、SRAM 使用策略,然后编译成可调用的算子。

一个面向应用开发者,一个面向算子开发者。日常写预处理脚本、做数据增强,asnumpy 够用了;如果你要开发新的融合算子、优化特定算子的性能上限,那得看 catlass。

昇腾异构计算架构里,这两层的定位是不同的:

AscendCL(应用层接口)
  └─ asnumpy(数据处理,API 友好)
catlass(算子开发层,底层灵活)
  └─ opbase(基础组件,上面两个都依赖它)

几个容易踩的坑

坑一:不是所有 NumPy API 都支持。 asnumpy 是 NumPy 的子集实现,不支持 FFT(用 ops-fft)、稀疏矩阵、某些特殊的线性代数操作。如果你的代码里用到了,先跑一遍看看报不报错。昇腾 CANN 的生态覆盖是逐步完善的,缺的功能可以提 Issue 到社区。

坑二:ascpu() 是隐性的性能杀手。 这个方法会触发一次完整的数据从 NPU 到 CPU 的搬运,有些场景下这一步会把前面的加速全部抵消。如果你的 pipeline 里频繁调用 ascpu(),比如每处理一张图就拿回 CPU 存一次,实际上并没有省到搬运时间。正确的做法是批量处理完再统一导出,或者干脆不导出直接喂给下游模型。

# 低效写法:每张图都搬运一次
for img in images:
    result = preprocess(img)
    save(result.ascpu())   # 每次都搬运,慢了

# 高效写法:批量处理,最后统一导出
batch = np.stack([img.asnpu() for img in images])
results = preprocess(batch)       # 全程 NPU
final = results.ascpu()           # 只搬运一次

坑三:dtype 精度问题。 asnumpy 的计算精度和 NumPy 有细微差异,主要体现在 float16 和 bfloat16 上。昇腾 NPU 的向量单元对 float16 有原生支持,asnumpy 默认会使用硬件加速;如果你明确需要 float32 的精度,需要显式指定 dtype,但这样可能会触发额外的精度转换开销。这个坑在图像处理里不明显,在科学计算场景里需要留意。


结尾

asnumpy 这个东西上手门槛很低,你甚至不需要懂昇腾 NPU 的编程模型,会 NumPy 就能用。但它解决的问题很实在——数据预处理不再需要绕道 CPU,NPU 加速从头到尾覆盖到。

如果你在昇腾 NPU 上跑 PyTorch 推理,可以把数据预处理部分用 asnumpy 重写,preprocess + 推理统一在 NPU 上完成,吞吐会比"CPU 预处理 + NPU 推理"的串行方式高出不少。

ops-tensor 仓库里除了 asnumpy,还有 tensorapi 和 Blaze 等其他张量工具,可以一起看看。源码在 https://atomgit.com/cann/ops-tensor。

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐