作者：昇腾实战派

知识地图：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

背景概述

2020 年 DeepMind 推出 AlphaFold2 后，蛋白质结构预测领域迎来了深度学习驱动的革命。紧接着在 2021 年，华盛顿大学蛋白质设计研究所（IPD）的 David Baker 团队推出了 RoseTTAFold——一种基于**三轨神经网络（3-track network）**的蛋白质结构预测模型，发表于 Science。RoseTTAFold 与 AlphaFold2 同期独立开发，在架构上采用了创新的三轨并行设计，同时处理序列进化信息、残基对相互作用和三维结构坐标，通过 SE(3) 等变 Transformer 保证了对旋转和平移的物理不变性。

RoseTTAFold 的核心优势在于：

• 三轨并行架构：同时利用 1D 序列、2D 残基对、3D 结构三类关键信息，三轨之间持续信息交互
• SE(3) 等变层：保证对输入结构做旋转/平移时，输出坐标做相同变换，符合蛋白质结构的物理不变性
• 端到端预测：从 MSA 输入直接产出 3D 原子坐标和置信度（pLDDT）
• 复合物建模：支持蛋白质复合物和蛋白-蛋白相互作用（PPI）预测

本文介绍 RoseTTAFold 模型的昇腾 Ascend NPU 适配版本——在保持原始模型精度的前提下，将 CUDA/DGL 依赖全面替换为 NPU 兼容实现，实现在昇腾 AI 平台上的端到端推理。

模型介绍

RoseTTAFold 概述

RoseTTAFold 是一个用于蛋白质结构预测的深度学习模型，输入蛋白质的多序列比对（MSA）和结构模板信息，输出蛋白质的三维原子坐标。模型的核心创新是三轨架构（3-track architecture），相比 AlphaFold2 的二轨设计多出一个结构轨道，使得结构信息能在网络的早期阶段就参与特征提取。

整体架构

三轨并行设计

轨道	信息类型	表示形状	功能
1D 序列轨	MSA 进化与上下文信息	B × N × L × d_msa	提取序列保守性、共进化信号
2D 残基对轨	残基间距离与相互作用约束	B × L × L × d_pair	建模残基对的空间关系
3D 结构轨	三维原子坐标	B × L × 3 × 3	生成与细化蛋白质骨架坐标

网络组成

1. Embedding 层

   • MSA Embedding：将 MSA 编码为 1D 序列表示（d_msa=384）
   • Pair Embedding：将序列对和模板信息编码为 2D 残基对表示（d_pair=288）
   • Template Embedding：处理结构模板的距离和方向信息（d_templ=64）

2. 主三轨迹模块（8 个 two-track 层 + 4 个 three-track 层）

   每个 Block 包含：

   • 1D 注意力：MSA 内行注意力（12 头）+ 列注意力，使用 Performer 高效注意力
   • 2D 注意力：残基对的双轴注意力（8 头）
   • 3D SE(3) 等变层：使用 SE(3)-Transformer 更新原子坐标
   • 跨轨信息交互：1D↔2D outer product，2D→3D 距离约束，3D→2D 坐标编码

3. 距离预测头（DistanceNetwork）

   • 从 2D 残基对表示预测 6D 距离和方向分布（Cβ-Cβ 距离、ω/θ/φ 角度）

4. 结构精修模块（Refine Module，4 层）

   • 使用 SE(3) 等变网络对初步 3D 结构做迭代优化
   • 修正键长、键角，消除原子碰撞

5. TRFold 模块

   • 基于预测的距离/方向分布，通过梯度下降优化最终 3D 坐标
   • 输出最终的蛋白质骨架 PDB 坐标

SE(3) 等变 Transformer

SE(3) 等变层是 RoseTTAFold 结构轨道的核心。它保证：对输入做任意 3D 旋转或平移时，输出会做完全相同的旋转或平移，完美符合蛋白质结构的物理不变性。具体实现基于 SE(3)-Transformer，使用球谐函数和 Wigner-D 矩阵实现等变特征传播。

残基框架表示

RoseTTAFold 采用刚性残基框架建模：每个残基用 (r, z_Cα) 表示——旋转矩阵 r 和 Cα 平移向量。通过 Gram-Schmidt 正交化从 C、N、Cα 原子坐标计算旋转矩阵，实现原子坐标与刚性框架的双向转换。

推理流程

输入 MSA (.a3m) + 模板 (.hhr/.atab)
    ↓
MSA/Pair/Template Embedding
    ↓
8 × Two-track 层（1D + 2D 注意力 + 跨轨交互）
    ↓
4 × Three-track 层（+ SE(3) 等变结构更新）
    ↓
距离/方向分布预测
    ↓
4 × 结构精修模块
    ↓
TRFold 梯度优化
    ↓
输出 PDB 坐标 + pLDDT 置信度

应用场景

• 单体蛋白结构预测：从序列/MSA 直接预测 3D 结构
• 蛋白质复合物建模：预测多链蛋白复合物结构
• 蛋白-蛋白相互作用筛选：使用 2-track 快速版本进行 PPI 筛选
• 结构驱动的蛋白质设计：为 RFdiffusion 等生成式设计提供基座模型

昇腾 NPU 适配版本

迁移动机

原始 RoseTTAFold 深度依赖 CUDA 和 DGL（Deep Graph Library）进行图神经网络计算。DGL 不支持 NPU 设备（dgl.graph().to(device) 无法识别 NPU 设备类型），因此需要对 SE(3) 等变注意力模块进行完整重写，使用 PyTorch 原生操作替代所有 DGL 图消息传递功能。

仓库结构

RoseTTAFold_npu/
├── README.md
├── change.txt                          # 迁移改动详细记录
├── requirements.txt                    # PyTorch + NPU 依赖
├── requirements_pyg.txt                # PyG 扩展依赖
├── network/
│   ├── predict_e2e.py                  # 端到端推理入口（NPU 适配）
│   ├── predict_complex.py              # 复合物推理入口
│   ├── predict_pyRosetta.py            # PyRosetta 版推理
│   ├── RoseTTAFoldModel.py             # 主模型定义
│   ├── Attention_module_w_str.py       # 三轨注意力模块
│   ├── Embeddings.py                   # MSA/Pair/Template Embedding
│   ├── DistancePredictor.py            # 距离预测头
│   ├── Refine_module.py                # 结构精修模块
│   ├── SE3_network.py                  # SE(3) 等变网络
│   ├── Transformer.py                  # Transformer 层
│   ├── performer_pytorch.py            # Performer 高效注意力
│   ├── trFold.py                       # TRFold 坐标优化
│   ├── equivariant_attention/
│   │   ├── modules.py                  # ⭐ 完整重写：DGL → PyTorch scatter
│   │   ├── fibers.py                   # Fiber 表示（view→reshape 适配）
│   │   └── from_se3cnn/               # SO(3) 表示与球谐函数
│   └── utils/
├── DAN-msa/                            # DeepAccNet MSA 质量评估
├── example/                            # 示例输入输出
│   ├── input.fa
│   ├── end-to-end/                     # e2e 预测示例
│   ├── pyrosetta/                      # PyRosetta 版示例
│   └── complex_modeling/               # 复合物建模示例
└── weights/                            # 模型权重（需下载）

核心迁移改动

本次迁移涉及 11 个文件的修改，其中 equivariant_attention/modules.py 为完整重写（905 行）。

1. DGL → PyTorch 原生操作（核心改动）

modules.py 完整重写，使用 scatter_add_ 替代所有 DGL 图消息传递：

原 DGL 操作	替换实现
import dgl	from torch_scatter import scatter_mean, scatter_sum（含 fallback）
dgl.DGLGraph	自定义 SimpleGraph 类（src, dst, num_nodes, edata）
G.update_all(msg_fn, fn.mean)	src_idx 索引 → matmul 计算消息 → scatter_add_ 聚合
G.apply_edges(msg_fn)	直接用 src_idx/dst_idx 索引节点特征计算边输出
edge_softmax(G, e)	自定义 scatter_softmax() 函数
AvgPooling / MaxPooling	tensor.mean() / tensor.max()

2. CUDA → NPU 自动迁移

所有推理入口文件顶部添加：

import torch_npu
from torch_npu.contrib import transfer_to_npu

3. autocast 适配

# 原：
@torch.cuda.amp.autocast(enabled=False)
# 改为：
@torch.amp.autocast("npu", enabled=False)

4. float64 精度问题修复

NPU 不支持 float64（double）运算，会静默降级为 float32。修复方法：

# 原：
ang = torch.linspace(0, np.pi, 19, device=self.device)[:-1]
# 改为：
ang = torch.linspace(0, float(np.pi), 19, device=self.device)[:-1]

5. NPU 异步执行修复

NPU 采用异步执行模式，fold() 返回时计算可能尚未完成，导致 PDB 文件全为 NaN：

result = bb[minidx].permute(1,0,2)
if result.device.type == 'npu':
    torch.npu.synchronize()
return result

6. 内存连续性修复

NPU 上 tensor 内存可能非连续，.view() 会报错，全部替换为 .reshape()。

7. NumPy 废弃 API 修复

np.bool → np.bool_（NumPy 1.24+ 移除了旧别名）

关键注意事项

• GNormBias 必须使用 self.bias（nn.ParameterDict）而非 self.transform（nn.ModuleDict），否则权重加载会报 Missing key
• GAttentiveSelfInt 必须保持原版 _build_net 多层网络结构，不能简化为单层 Linear，否则 state_dict 键名不匹配
• get_basis 中使用 torch.einsum 替代 torch.matmul（规避 aarch64 上 PyTorch oneDNN 的小维度 bug）

版本信息

软件	版本
HDK	25.3.rc1.2
CANN	8.3.RC1
Python	3.10
PyTorch	2.10.0
torch_npu	2.10.0

环境配置

创建 Conda 环境

conda create --name RoseTTAFold python=3.10 -y
conda activate RoseTTAFold

克隆代码

git clone https://atomgit.com/AI4Science/RoseTTAFold.git
cd RoseTTAFold

安装依赖

export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/
pip install -r requirements.txt
pip install -r requirements_pyg.txt --no-build-isolation

requirements.txt 主要包含：

• torch==2.10.0、torch_npu==2.10.0
• numpy==1.26.4、scipy、biopython、pandas、scikit-learn
• CANN 运行时依赖：decorator、attrs、psutil、absl-py

requirements_pyg.txt 包含：

• torch-scatter、torch-sparse、torch-cluster
• torch_geometric、lie-learn

验证安装

python3 -c "import torch;import torch_npu; a = torch.randn(3, 4).npu(); print(a + a);"

若输出类似以下信息，说明安装成功：

tensor([[-0.6066,  6.3385,  0.0379,  3.3356],
        [ 2.9243,  3.3134, -1.5465,  0.1916],
        [-2.1807,  0.2008, -1.1431,  2.1523]], device='npu:0')

若报错，排查顺序：

1. set_env.sh 是否已 source
2. decorator 等运行时依赖是否安装
3. CANN 与 torch_npu 版本是否匹配

模型权重下载

wget https://files.ipd.uw.edu/pub/RoseTTAFold/weights.tar.gz
tar xfz weights.tar.gz

权重包含：

• RoseTTAFold_e2e.pt：端到端模型
• RF2t.pt：2-track PPI 筛选模型
• 及其他相关检查点

注意：权重使用 Rosetta-DL Software License（非商业用途）。

推理命令

端到端结构预测

python network/predict_e2e.py \
    -m ./weights \
    -i example/end-to-end/t000_.msa0.a3m \
    -o ./output/t000_.e2e

输出文件：

• t000_.e2e.pdb：预测的 PDB 结构文件（B-factor 列为残基级 CA-lDDT 置信度）
• t000_.e2e.npz：预测的距离/方向分布

复合物结构预测

python network/predict_complex.py \
    -i paired.a3m \
    -o complex \
    -Ls 218 310

PPI 筛选（2-track 快速版本）

python network_2track/predict_msa.py \
    -msa input.a3m \
    -npz complex.npz \
    -L1 218

迁移适配要点总结

改动文件清单

文件	改动类型	说明
network/predict_e2e.py	3 类标准改动	import torch_npu + np.bool_ + autocast
network/predict_complex.py	同上	相同 3 类改动
network/predict_pyRosetta.py	同上	相同 3 类改动
network/equivariant_attention/modules.py	⭐ 完整重写	DGL → PyTorch scatter（905 行）
network/Attention_module_w_str.py	DGL → SimpleGraph	make_graph() + autocast
network/Refine_module.py	autocast 适配	cuda → cpu
network/SE3_network.py	autocast 适配	cuda → npu（2 处）
network/equivariant_attention/fibers.py	view → reshape	6 处内存连续性修复
network/trFold.py	float64 + 异步	5 处 float 修复 + npu.synchronize
network/utils/utils_data.py	重写	DGL → SimpleGraph 兼容
from_se3cnn/representations.py	设备适配	cuda → npu（测试块）

迁移策略总结

问题类别	解决方案
DGL 不支持 NPU	用 PyTorch scatter_add_ 替代所有图消息传递
CUDA API 调用	torch_npu + transfer_to_npu 自动迁移
torch.cuda.amp.autocast	替换为 torch.amp.autocast("npu")
float64 不支持	所有 numpy 常量用 float() 包裹
.view() 非连续报错	全部替换为 .reshape()
NPU 异步执行	关键位置添加 torch.npu.synchronize()
np.bool 废弃	替换为 np.bool_

已知限制

• SE(3) 等变层计算量大，NPU 上的推理速度受限于 scatter 操作的效率
• 原始 PyRosetta 折叠步骤仍需 CPU 运行（PyRosetta 不支持 NPU）
• DAN-msa 错误预测模块使用 TensorFlow checkpoint，仍在 CPU 上运行
• 权重使用 Rosetta-DL 非商业许可证

参考文献

• Baek, M., et al. Accurate prediction of protein structures and interactions using a three-track neural network. Science (2021). https://www.science.org/doi/10.1126/science.abj8754
• Humphreys, I.R., et al. Computed structures of core eukaryotic protein complexes. Science (2021). https://www.science.org/doi/10.1126/science.abm4805
• 上游代码仓库：https://github.com/RosettaCommons/RoseTTAFold
• 昇腾 NPU 适配版：https://atomgit.com/AI4Science/RoseTTAFold