概述

Protenix 是由字节跳动开发的全原子共折叠与结构预测模型，是 AlphaFold 3 架构的开源复现实现。该模型支持蛋白-蛋白、蛋白-核酸、蛋白-小分子配体等多种复合物类型的全原子结构预测，采用 Pairformer + 扩散模型架构，在精度上对标 AlphaFold 3，同时提供完全开源的代码和权重。

• 参考实现：

  url=https://github.com/bytedance/Protenix.git
  commit_id=bd54a05d047b8925a241056f36d619700604068a
  

• 适配昇腾 AI 处理器的实现：

  url=https://gitcode.com/AI4Science/LifeScience.git
  code_path=PyTorch/Protenix
  

组件版本

hdk：25.5.0
cann：8.3.RC1
python：3.11
torch：2.7.1
torch-npu：2.7.1

环境准备

a. 新建conda环境
conda create --name protenix python=3.11
conda activate protenix

b. 安装torch
pip install torch==2.7.1

c. 安装torch_npu
AArch64架构:
wget https://gitcode.com/Ascend/pytorch/releases/download/v26.0.0-pytorch2.7.1/torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_aarch64.whl
pip3 install torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_aarch64.whl

X86_64架构:
wget https://gitcode.com/Ascend/pytorch/releases/download/v26.0.0-pytorch2.7.1/torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_x86_64.whl
pip3 install torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_x86_64.whl

如果遇到证书不可信，需在命令结尾添加 --no-check-certificate

d. 安装Cann补丁包：
如果要跑训练，必须在8.3.RC1的Cann版本上安装补丁包。
git clone https://gitcode.com/Ascend-SACT/CANN-FAG-ops.git
cd CANN-FAG-ops
cd op_run_packages_on_${device}
chmod +x cann-opbase_8.5.0.alpha001_linux-aarch64.run
# install-path是cann的toolkit的安装路径 请自行修改成您的安装路径
./cann-opbase_8.5.0.alpha001_linux-aarch64.run --full  --install-path=${install_path}/ascend-toolkit
chmod +x cann-ops-transformer-custom_linux-aarch64.run
./cann-ops-transformer-custom_linux-aarch64.run

拉取模型代码

git clone https://gitcode.com/AI4Science/LifeScience.git
# 和LifeScience同层级拉取Protenix代码
git clone https://github.com/bytedance/Protenix.git
cd Protenix
git checkout bd54a05d047b8925a241056f36d619700604068a
git apply --ignore-whitespace ../LifeScience/PyTorch/Protenix/patch/Protenix.patch
#安装依赖
pip install -r requirements.txt

数据集准备

a. 进入Protenix目录
cd Protenix/scripts/database
b. 在终端窗口执行下载数据命令
# 默认下载推理数据集
bash download_protenix_data.sh
# 完整下载全量数据集（训练+推理）
bash download_protenix_data.sh --full

推理

cd Protenix
bash inference_demo.sh

训练

cd Protenix
bash train_demo.sh

训练结果对比

芯片	卡数	性能-平均单步迭代耗时(s)
竞品A	1p	22
Atlas 800T A2	1p	33