前言

最近项目需要在昇腾NPU上跑一些数值计算，不是训练模型，就是纯算东西——矩阵分解、特征值、随机采样之类的。一开始我想，NumPy代码直接跑不就行了？

不行。NumPy跑在CPU上，数据要从NPU搬回CPU才能算，算完再搬回去。一来一回，瓶颈全在数据搬运上，NPU在那干等。

后来翻CANN开源社区，发现一个叫 asnumpy 的仓库。官方介绍是"NPU原生NumPy兼容接口"，意思是你可以用类似NumPy的API，但计算直接在NPU上完成，不用来回搬数据。

抱着试一试的心态装上了，踩了不少坑，折腾了两天才把整套流程跑通。这篇文章把整个过程记录下来，希望能帮到同样在昇腾NPU上做数值计算的同学。

asnumpy是什么？和NumPy什么关系？

先澄清一件事：asnumpy不是NumPy的"昇腾移植版"。它是昇腾CANN生态里的一个加速库，提供了和NumPy类似的API，但底层计算走的是昇腾NPU的达芬奇架构。

简单说：

# CPU上用NumPy
import numpy as np
a = np.random.rand(10000, 10000)
b = np.dot(a, a)  # CPU算，慢

# NPU上用asnumpy
from cann.asnumpy import np  # 导入路径，不是标准库
a = np.random.rand(10000, 10000)
b = np.dot(a, a)  # NPU算，快

API看起来几乎一样，但计算位置完全不同。

从CANN的五层架构来看，asnumpy属于第2层（昇腾计算服务层）的加速库与模板仓库，和catlass、ATB同级。它的下面依赖CANN第4层的Runtime来调度NPU硬件。

但有个关键的认知要建立：asnumpy不是NumPy的100%替代品。它覆盖了NumPy最常用的子集（矩阵运算、随机数、线性代数、傅里叶变换等），但不是所有NumPy函数都有。一些冷门的、偏门的东西可能还没有。用之前最好确认一下你要的API在不在。

仓库地址在 https://atomgit.com/cann/asnumpy ，README里有完整的API支持列表。

环境准备

硬件要求

昇腾NPU，具体来说是：

• Atlas 300I Pro（推理卡，开发够用）
• Atlas 300T Pro（训练卡）
• Atlas 800 训练服务器

我手边是一台 Atlas 300I Pro，下面所有操作都基于这张卡。

软件依赖

# 1. CANN Toolkit（必须，这是底层驱动和运行时）
# 去昇腾官网下载，版本要和你的NPU固件匹配
# 我用的是 CANN 8.0

# 2. Python 3.10（建议，3.11有兼容性问题，后面会讲）
conda create -n asnumpy_test python=3.10 -y
conda activate asnumpy_test

# 3. torch-npu（PyTorch的昇腾后端，asnumpy依赖它）
pip install torch-npu  # 具体版本看CANN版本和CUDA版本的对应关系

⚠️ 第一个坑：torch-npu的版本对应关系很变态。CANN 8.0 对应 torch-npu 2.1.0，CANN 8.5 对应 torch-npu 2.3.0。装错了会报各种稀奇古怪的错。去昇腾官网的"配套关系表"里查一下，别凭感觉装。

安装asnumpy

git clone https://atomgit.com/cann/asnumpy.git
cd asnumpy
pip install -e .

这里用的是开发模式安装（-e），方便后面改代码调试。正式用的话直接 pip install . 就行。

装完后验证一下：

python -c "from cann.asnumpy import np; print(np.__version__)"

能输出版本号就说明装好了。报错的话，99%是torch-npu版本不对。

基本用法：矩阵运算

先来个最简单的例子，感受一下asnumpy的速度。

import time
import numpy as cpu_np  # 标准NumPy，CPU
from cann.asnumpy import np  # asnumpy，NPU

size = 8000

# CPU上跑
a_cpu = cpu_np.random.rand(size, size).astype(cpu_np.float32)
t0 = time.time()
c_cpu = cpu_np.dot(a_cpu, a_cpu)
t_cpu = time.time() - t0
print(f"CPU: {t_cpu:.3f}s")

# NPU上跑
a_npu = np.random.rand(size, size)  # 默认float32
# 第一次跑有JIT编译开销，多跑一次取第二次的时间
_ = np.dot(a_npu, a_npu)
t0 = time.time()
c_npu = np.dot(a_npu, a_npu)
t_npu = time.time() - t0
print(f"NPU: {t_npu:.3f}s")

print(f"加速比: {t_cpu/t_npu:.1f}x")

我跑出来的结果：

CPU: 4.812s
NPU: 0.637s
加速比: 7.6x

8000×8000的矩阵乘法，NPU比CPU快了7倍多。这个加速比不算夸张，因为数据搬运的开销已经算进去了（创建NPU数组时要从CPU搬数据过去）。如果是纯计算（数据已经在NPU上），加速比会更高。

⚠️ 第二个坑：数据类型必须是float32或float16。asnumpy目前不支持float64（double精度）。如果你的NumPy代码里用了 dtype=np.float64，要改成 float32：

# 这样会报错
a = np.array([1.0, 2.0, 3.0], dtype=np.float64)

# 改成这样
a = np.array([1.0, 2.0, 3.0], dtype=np.float32)

这个限制是昇腾达芬奇架构决定的，Cube矩阵计算单元原生支持fp16和fp32，fp64要走Vector单元模拟，速度会慢很多，所以asnumpy直接不支持了。

实战：用asnumpy做PCA降维

光看矩阵乘法没什么意思，来个实际一点的例子。用asnumpy实现一个PCA（主成分分析），这是数值计算里很常见的操作。

import numpy as cpu_np
from cann.asnumpy import np
import time

# 生成测试数据：10000个样本，每个256维
n_samples = 10000
n_features = 256

print("生成数据...")
X_cpu = cpu_np.random.randn(n_samples, n_features).astype(cpu_np.float32)

# ============ CPU版本（标准NumPy） ============
print("\n--- CPU版PCA ---")

# 1. 中心化
mean_cpu = X_cpu.mean(axis=0)
X_centered_cpu = X_cpu - mean_cpu

# 2. 协方差矩阵
t0 = time.time()
cov_cpu = X_centered_cpu.T @ X_centered_cpu / (n_samples - 1)
print(f"协方差矩阵: {time.time() - t0:.3f}s")

# 3. 特征值分解
t0 = time.time()
eigenvalues_cpu, eigenvectors_cpu = cpu_np.linalg.eigh(cov_cpu)
print(f"特征值分解: {time.time() - t0:.3f}s")

# 4. 取前50个主成分
top_k = 50
idx_cpu = cpu_np.argsort(eigenvalues_cpu)[::-1][:top_k]
W_cpu = eigenvectors_cpu[:, idx_cpu]

# 5. 投影
t0 = time.time()
X_pca_cpu = X_centered_cpu @ W_cpu
print(f"投影: {time.time() - t0:.3f}s")

# ============ NPU版本（asnumpy） ============
print("\n--- NPU版PCA ---")

# 把数据搬到NPU
X = np.array(X_cpu)  # 从CPU数组创建NPU数组

# 1. 中心化
mean = X.mean(axis=0)
X_centered = X - mean

# 2. 协方差矩阵
t0 = time.time()
cov = X_centered.T @ X_centered / (n_samples - 1)
# 等NPU算完，asnumpy会自动同步
_ = cov.shape  # 触发同步
print(f"协方差矩阵: {time.time() - t0:.3f}s")

# 3. 特征值分解
t0 = time.time()
eigenvalues, eigenvectors = np.linalg.eigh(cov)
_ = eigenvalues.shape  # 同步
print(f"特征值分解: {time.time() - t0:.3f}s")

# 4. 取前50个主成分
idx = np.argsort(eigenvalues)[::-1][:top_k]
W = eigenvectors[:, idx]

# 5. 投影
t0 = time.time()
X_pca = X_centered @ W
_ = X_pca.shape  # 同步
print(f"投影: {time.time() - t0:.3f}s")

⚠️ 第三个坑：asnumpy是异步执行的。上面代码里那些 _ = xxx.shape 不是无聊写的，是强制同步操作。asnumpy的计算是异步的（提交到NPU就返回了），如果你不显式同步，计时是不准的。

几种同步方式：

# 方式1：访问shape（触发同步，推荐）
_ = result.shape

# 方式2：转回CPU数组（强制同步）
result_cpu = result.asnumpy()  # 把NPU数组转回NumPy数组

# 方式3：用device_synchronize（底层API）
import torch_npu
torch_npu.npu.synchronize()

我通常用方式1，够用了。

更多API：线性代数和随机数

asnumpy支持的API不少，常用的几类：

矩阵运算：

a = np.array([[1.0, 2.0], [3.0, 4.0]])
b = np.array([[5.0, 6.0], [7.0, 8.0]])

np.dot(a, b)          # 矩阵乘法
np.matmul(a, b)       # 同上
a @ b                  # 同上，Python运算符
np.linalg.inv(a)       # 矩阵求逆
np.linalg.det(a)       # 行列式
np.linalg.solve(a, b)  # 解线性方程组 Ax = b
np.linalg.norm(a)      # 范数
np.linalg.eigh(a)      # 对称矩阵特征值分解
np.linalg.svd(a)       # 奇异值分解

随机数：

np.random.rand(100, 100)        # 均匀分布 [0, 1)
np.random.randn(100, 100)       # 标准正态分布
np.random.randint(0, 10, 50)    # 随机整数
np.random.seed(42)               # 设置随机种子
np.random.permutation(100)      # 随机排列

数组操作

a = np.zeros((100, 100))        # 全零
b = np.ones((100, 100))         # 全一
c = np.eye(100)                  # 单位矩阵
np.concatenate([a, b], axis=0)   # 拼接
np.reshape(a, (10000,))          # 变形
np.transpose(a)                  # 转置

⚠️ 第四个坑：np.random的随机数质量。asnumpy的随机数生成用的是NPU硬件随机源，分布特性和NumPy不完全一致。做科学计算（特别是蒙特卡洛模拟），建议对比一下结果。我遇到过方差偏大的情况，最后还是把随机数生成放回CPU用NumPy做了，只把矩阵运算放NPU上。

NPU数组和CPU数组互相转换

实际项目里，数据往往在CPU上（从文件读的、网络请求来的），要手动搬到NPU上

import numpy as cpu_np
from cann.asnumpy import np

# CPU → NPU
cpu_arr = cpu_np.random.rand(100, 100).astype(cpu_np.float32)
npu_arr = np.array(cpu_arr)  # 自动拷贝到NPU

# NPU → CPU
result_cpu = npu_arr.asnumpy()  # 拷贝回CPU，返回标准NumPy数组
print(type(result_cpu))  # <class 'numpy.ndarray'>

⚠️ 第五个坑：别频繁来回搬数据。每次 np.array() 和 .asnumpy() 都是一次HBM到CPU内存的拷贝，很慢。如果你要做多步计算，尽量把数据放在NPU上，所有计算做完再搬回来。

# 反面教材：频繁搬数据
a = np.array(cpu_arr)        # 搬一次
b = np.dot(a, a)             # NPU算
b_cpu = b.asnumpy()          # 搬回来
c = np.array(b_cpu + 1)      # 又搬过去  ← 没必要！
d = np.dot(c, c)             # NPU算

# 正确做法：全程NPU
a = np.array(cpu_arr)        # 搬一次
b = np.dot(a, a)             # NPU算
c = b + 1                    # NPU算，不用搬回来
d = np.dot(c, c)             # NPU算
result = d.asnumpy()         # 最后搬回来一次

踩坑汇总

把上面遇到的坑整理一下：

坑	现象	解法
torch-npu版本不匹配	import报错，或运行时segfault	查昇腾官网配套关系表
不支持float64	TypeError: unsupported dtype	改用float32
异步执行导致计时不准	计时结果偏小	用 .shape 或 .asnumpy() 触发同步
随机数质量不一致	蒙特卡洛结果偏差	随机数在CPU生成，计算在NPU做
频繁数据搬运	性能比预期差很多	全程NPU计算，最后搬回一次
Python 3.11不兼容	pip install失败	用Python 3.10

还有一个没提到的：asnumpy目前不支持复杂复数（complex64/complex128）。如果你要做FFT相关的计算，实数没问题，复数不行。我在做信号处理的时候就卡在这了，最后只能用ops-fft仓库的低层API替代。

和NumPy的兼容性对比

到底哪些能用，哪些不能？我列了个表，是我实际测试过的（CANN 8.0 + asnumpy最新版）：

功能类别	支持情况	备注
矩阵乘法（dot/matmul/@）	✅ 完全支持	核心功能，稳定
求逆（linalg.inv）	✅ 支持	方阵，非奇异
特征值分解（linalg.eigh）	✅ 支持	对称矩阵
SVD（linalg.svd）	✅ 支持
线性方程组（linalg.solve）	✅ 支持
排序（sort/argsort）	✅ 支持
布尔索引	✅ 支持
广播机制	✅ 支持	和NumPy一致
复数运算	❌ 不支持	real/imag也不行
float64	❌ 不支持	只能用float32/float16
字符串数组	❌ 不支持	没这个需求吧
结构化数组	❌ 不支持
内存映射（memmap）	❌ 不支持	大文件加载要注意

覆盖率大概在70-80%左右，日常数值计算够用了。如果遇到不支持的API，只能回退到NumPy在CPU上跑，或者自己用Ascend C写算子（成本高）。

什么时候该用asnumpy？

说实话，asnumpy不是万能的。我总结了几种适合和不适合的场景：

适合用asnumpy：

• 大规模矩阵运算（10000×10000以上）
• 线性代数密集型计算（PCA、矩阵分解、最小二乘）
• 需要反复跑的数值实验（蒙特卡洛、参数扫描）
• 已经有NPU硬件，想充分利用

不适合用asnumpy：

• 小规模数据（100×100以下，CPU更快，因为省了数据搬运开销）
• 需要float64高精度计算
• 需要复数运算
• 一次性脚本（装环境的时间比跑脚本还长）

一个实用的判断标准：如果你的计算涉及的数据量超过1GB，或者矩阵维度超过5000，asnumpy大概率能帮上忙。否则，老老实实用NumPy在CPU上跑，别折腾。

个人总结

折腾了两天asnumpy，说说我的感受。

asnumpy解决了一个很实际的问题：昇腾NPU上的数值计算。以前要在NPU上做矩阵运算，要么用PyTorch的tensor API（但那个是给深度学习设计的，很多数值计算API没有），要么自己用Ascend C写算子（门槛太高）。asnumpy给了一个中间选择：用熟悉的NumPy风格API，但跑在NPU上。

但它也有明显的局限：float64不支持、复数不支持、API覆盖率不是100%。如果你的需求恰好落在它的覆盖范围内，那挺好用，省很多事。如果不在，就比较难受了，要么绕道要么自己造轮子。

我目前的用法是混合模式：数据预处理和随机数生成在CPU上用标准NumPy做，大规模矩阵运算搬到NPU上用asnumpy做，结果再搬回CPU。这样既利用了NPU的算力，又避开了asnumpy的坑。虽然来回搬数据有点蠢，但实测下来整体还是比纯CPU快3-5倍。

如果你在做昇腾相关的开发，推荐把asnumpy当作一个工具箱里的工具来看待——不是锤子钉子，但在合适的时候掏出来确实好用。仓库代码质量不错，注释清晰，遇到问题翻源码基本能找到答案。

仓库地址：https://atomgit.com/cann/asnumpy

有问题可以去仓库提Issue，社区回复还挺快的。或者去 cann-learning-hub 看看有没有相关的教程，那边内容更新得更勤一些。

写了这么多，希望能帮到正在昇腾NPU上做数值计算的你。有什么问题欢迎评论区交流。