用途:将 DiffSBDD(基于等变扩散模型的结构化药物设计项目)从 CUDA 迁移到华为 Ascend NPU,覆盖环境搭建、依赖安装、torch_scatter 源码编译、代码适配与 de novo 推理验证。
前置条件:Ascend910 系列(至少 1 卡);openEuler / Ubuntu / KylinOS(aarch64 或 x86_64);CANN ≥ 8.2(推荐 8.2.RC1 或 8.3.RC1);Python 3.10;PyTorch 2.5.1;torch_npu 2.5.1。
主要流程:克隆 DiffSBDD 仓库并创建 Conda 环境;配置 CANN 环境与华为镜像源;安装 PyTorch、torch_npu 及 numpy==1.26.4 等依赖;安装 pytorch-lightning、wandb、biopython、rdkit、openbabel 等项目依赖;从源码编译 torch_scatter(设置 FORCE_CUDA=0,禁止使用 PyG 预编译 wheel);在 generate_ligands.py、test.py、inpaint.py、optimize.py、train.py、lightning_modules.py 六个入口文件顶部注入 torch_npu 与 transfer_to_npu 实现自动迁移;补充 analysis/ 与 analysis/SA_Score/ 的 init.py;下载预训练模型后执行 de novo 推理验证。
验证命令:python generate_ligands.py checkpoints/crossdocked_fullatom_cond.ckpt --pdbfile example/3rfm.pdb --ref_ligand example/3rfm_B_CFF.sdf --outfile example/3rfm_generated.sdf --n_samples 2 --num_nodes_lig 44。通过标准为程序正常退出(exit code 0)且输出 SDF 文件非空、格式正确。
特殊适配点:numpy 必须为 1.26.4,conda 安装 rdkit 后可能被覆盖为 2.x,需手动恢复;conda 可能残留 numpy-base 2.x,需 conda remove numpy-base --force;torch_scatter 必须源码编译;wandb 原始 0.13.1 不兼容 numpy 1.26.4,需升级;Ascend910 不支持 double,自动降级为 fp32,扩散模型比较分布指标即可;transfer_to_npu 自动处理 CUDA→NPU 设备映射。
配套脚本:python scripts/validate_diffsbdd_env.py 用于依赖与 NPU 运行时预检。参考资料见 references/dependency-notes.md。