昇腾原生图机器学习库 dgl-ascend 0.0.3 版本发布
北邮 GAMMA Lab
2026年10月21日 18:35
北京
去阅读
一、研究背景
图神经网络作为机器学习的一个重要分支，被广泛应用于知识图谱、社交网络、流体仿真、蛋白质设计、分子属性预测、世界模型等场景。依托强大的拓扑结构解析与建模能力，图神经网络正借助 AI for Science 浪潮持续驱动科研范式革新。
当前，行业主流的图神经网络库是 DGL（Deep Graph Library）和 PyTorch Geometric（PyG），二者已深度落地各类 AI4S 研究场景，是全球科研团队开展生物大分子建模、分子筛选、药物仿真等实验的核心技术底座。长期以来，DGL/PyG 等原生框架未提供昇腾 NPU 适配能力，图神经网络模型难以直接在昇腾 AI 软硬件平台上高效运行，限制了昇腾 AI 软硬件平台在图科学计算领域的应用落地。
二、项目概况
为进一步拓宽图神经网络算力生态的适配边界，充分释放昇腾 AI 基础软硬件在基础科研场景中的算力效能，北京邮电大学石川教授团队联合华为昇腾生态发展部、昇腾计算基础软件开发部、小巧灵突击队，正式启动图机器学习适配昇腾专项攻坚项目。项目将完成 DGL、PyG 等框架与昇腾 AI 基础软硬件平台的深度适配、算子调优与稳定性打磨，实现各类图神经网络模型在昇腾 NPU 上高效、稳定、规模化地部署运行，为 AI4S 前沿科研创新提供更强的算力支撑。
当前，dgl-ascend 项目已适配约 50 个算子，支持 LightGCN、GraphSAGE、HERec 等主流图神经网络，以及 DeepChem、DGL-LifeSci 等套件中的 100+ 模型，并于今天发布 0.0.3 版本。北京邮电大学石川教授在今天的 CCF CNCC 2026《图基础模型研究进展与应用》专题讲座中，介绍了昇腾图机器学习平台及相关项目的最新进展。
昇腾图机器学习平台与生态
三、0.0.3 版本新特性
dgl-ascend 0.0.3 新版本特性清单如下：
- 支持 dgl-ascend 在 950PR + CANN 9.1.0 上编译运行；
- 新增 concat、UnaryElewise、COOLineGraph 等 10+ 算子；
- 修复 2969 个 nn 测试用例，DGL 内置的 60+ 基础模型（如 GraphConv 等）已完成适配；
- 修复部分算子 bug。
四、快速上手：在昇腾上使用 dgl-ascend
4.1 环境准备
依赖包：torch\_npu 2.8.0 / CANN 9.0.0
4.2 安装
方式一：pip 安装（推荐，简单快速）
pip install dgl-ascend==0.0.3 -i https://mirrors.aliyun.com/pypi/simple/
方式二：源码编译安装
适合需要定制、跟进主线或适配特定 CANN/芯片版本的场景。
# 克隆源码
git clone https://gitcode.com/AI4Science/dgl-ascend.git
cd dgl-ascend
# 初始化子模块
git submodule update --init --recursive
# 编译
bash script/build_dgl_ascend.sh
# 编译完成后按脚本提示安装
pip install -e ./python
4.3 环境验证
python -c \"import dgl; print(dgl.__version__); print(dgl.backend.backend_name)\"
# 期望输出：0.0.3 / pytorch
4.4 示例：在昇腾 NPU 上训练 GCN
以下示例构造一个小图并定义两层 GCN，通过 .to(\"npu\") 将图和模型迁移到昇腾 NPU 上完成训练。为避免联网下载数据集（如 Cora），示例使用自构造图。
import torch
import torch.nn as nn
import torch.nn.functional as F
import dgl
import dgl.nn as dglnn
# ---------- 1. 构造一个小图 ----------
num_nodes = 64
src = torch.arange(num_nodes)
dst = (torch.arange(num_nodes) + 1) % num_nodes
g = dgl.graph((src, dst))
g = dgl.add_self_loop(g) # 消除 0 入度节点
g.ndata[\"feat\"] = torch.randn(num_nodes, 16) # 特征维度 16
g.ndata[\"label\"] = torch.tensor([i % 2for i in range(num_nodes)], dtype=torch.long) # 二分类标签
# ---------- 2. 定义 GCN ----------
class GCN(nn.Module):
def __init__(self, in_feats, hidden, nclass):
super().__init__()
self.conv1 = dglnn.GraphConv(in_feats, hidden, allow_zero_in_degree=True)
self.conv2 = dglnn.GraphConv(hidden, nclass, allow_zero_in_degree=True)
def forward(self, g, feat):
h = F.relu(self.conv1(g, feat))
return self.conv2(g, h)
# ---------- 3. 迁移到 NPU 并训练 ----------
device = torch.device(\"npu\")
model = GCN(16, 32, 2).to(device) # 模型迁移到 NPU
g = g.to(device) # 图数据迁移到 NPU
feat, label = g.ndata[\"feat\"], g.ndata[\"label\"]
opt = torch.optim.Adam(model.parameters(), lr=0.01)
model.train()
for epoch in range(5):
logits = model(g, feat) # 前向在 NPU 上完成
loss = F.cross_entropy(logits, label)
opt.zero_grad()
loss.backward() # 反向在 NPU 上完成
opt.step()
print(f\"epoch {epoch} | loss={loss.item():.4f} | dev={logits.device}\")
print(\"GCN on NPU OK\")
4.5 运行结果
实测输出（Ascend 910B3）：
epoch 0 | loss=0.7343 | dev=npu:0
epoch 1 | loss=0.7102 | dev=npu:0
epoch 2 | loss=0.7048 | dev=npu:0
epoch 3 | loss=0.7043 | dev=npu:0
epoch 4 | loss=0.7020 | dev=npu:0
GCN on NPU OK
dev=npu:0 表示图与模型已成功在昇腾 NPU 上执行。
五、下一步计划
后续，我们将继续联合华为昇腾生态发展部、昇腾计算基础软件开发部、小巧灵突击队，推进 DGL、PyG 算子的 NPU 迁移工作，逐步完成图机器学习各类核心算子在昇腾平台上的适配与优化，进一步提升图计算任务在国产 AI 硬件上的执行能力，并构建昇腾图机器学习社区。同时，我们也将持续扩展模型支持范围，覆盖更多图神经网络、推荐系统模型和异构图学习模型，推动昇腾图机器学习在更多实际业务场景中的落地应用。
六、实习招募
围绕 dgl-ascend、CANN/ops-gnn 等开源项目，我们长期招募对图机器学习、机器学习系统和 AI 基础软件感兴趣的同学参与开发实习。无论你是大一、大二的低年级本科生，还是大三准备保研或出国深造的同学，我们都将根据你的兴趣、基础和发展阶段，安排合适的项目任务与培养方案。
在老师和学长学姐的指导下，你将参与真实开源项目与科研项目开发，在实践中系统提升工程能力、科研能力和团队协作能力。
我们希望你具备
- 开发能力：实习期间至少能够独立完成一个 DGL/PyG 算子的开发、迁移或适配工作，具体任务将根据个人意愿与能力进行安排。
- 学习与责任心：对新技术保持好奇心和学习热情，责任心强，愿意在项目实践中不断打磨与提升自己的工程能力。
- 时间投入：能够保证每周有一定投入时间，至少可以持续数个月的实习周期，欢迎长期合作。
- 合作精神：勤奋踏实、细致严谨，具备良好的沟通能力和团队协作意识。
- 技术基础：具备 Python 和 C 语言编程基础，了解深度学习相关基础知识；熟悉 PyTorch、TensorFlow 等深度学习框架者优先。
你将收获
- 开源贡献与署名：所有开发成果将合入 GAMMA Lab 的开源项目，在 GitHub/GitCode 等相关仓库与文档中署名，真实参与大型开源项目的维护与迭代。
- 科研与工程实践：在真实科研项目中系统学习图神经网络、图基础模型等前沿技术，积累一线工程和科研实践经验。
- 后续发展机会：对于在项目中表现优秀、配合度高的同学，将组织参与科研论文实习；后续在保研本组时优先考虑，并在保研外校、申请海外 PhD 等方面提供支持与帮助。
本次项目主要面向对图机器学习、机器学习等方向有兴趣，并具备一定编程与机器学习基础的同学。期待你在实际项目中与我们一起学习、开发和成长。
项目与资料
- dgl-ascend：https://gitcode.com/AI4Science/dgl-ascend
- ops-gnn：https://gitcode.com/cann/ops-gnn
- 相关学习资料：https://my.feishu.cn/docx/Um8edaBkjoyWo2x33Bgc4madnK9?from=from\_copylink
联系方式
有意者请邮件联系：
- 邮箱：[email protected]
- 邮件主题：姓名 + 学院/专业 + 年级 + GAMMA Lab 开发实习
- 邮件内容：个人简介 + 个人简历 PDF
版权声明
本站仅做备份收录，仅供研究与教学参考之用。
读者将信息用于其他用途的，全部法律及连带责任由读者自行承担，本站不承担任何责任。
评论