基于昇腾适配蛋白质序列模型ProteinMPNN

一 ProteinMPNN介绍

ProteinMPNN(Protein Message Passing Neural Network)是一种基于深度学习的蛋白质序列设计模型,核心目标是解决“逆向折叠问题”(inverse folding problem),即根据给定的蛋白质三维结构,设计出能够折叠成该结构的氨基酸序列。ProteinMPNN在计算和实验测试中都有出色的性能表现,不同位置的氨基酸序列可以在单链或多链之间偶联,从而广泛的应用于当前蛋白质设计上。ProteinMPNN不仅在天然蛋白质序列恢复率上面性能要高于传统的Rosetta方法,并且可以恢复先前设计失败的蛋白质。通过前沿AI技术突破科学研究的效率瓶颈,对于蛋白质工程、药物设计、酶设计等领域有极其重要的意义。

二 整体架构

2.1 编码器

2.1.1 图构建

将蛋白质结构表示为图结构,图的节点代表氨基酸残基,边代表残基之间的空间或序列关系(如距离、接触、氢键等)。

2.1.2 节点嵌入

使用可学习的嵌入层将每个残基的类型(氨基酸种类)、位置信息编码为初始节点特征。

是节点v在第l层的嵌入向量。

是节点v的邻居节点集合,是从节点u到节点v的消息。

2.1.3 边嵌入

对边进行编码,边的初始特征包括残基之间对距离、角度的几何变换,或残基间的相互作用特征。

e_{u v}^{(l)}是边(u, v)在第l层的嵌入向量,分别是节点u和节点v的嵌入向量。

2.2 消息传递网络

模型通过多层堆叠的消息传递和节点更新操作迭代更新节点状态

2.2.1 多层图卷积

在每一层图卷积中,节点v接收邻居节点u的消息Muv,由边特征euv和节点特征hu计算得到:

其中是第l层的可学习函数(如MLP),e_{u v} 包含距离、角度等几何信息。

2.2.2 SE3-等变注意力机制

注意力分数计算

计算节点v对邻居u的注意力分数 e_{u v},结合系欸但特征和边特征:

其中Wa, Wb, Wc是可学习权重矩阵,\beta是层间参数,softmax确保归一化。

消息聚合(SE3-等变)

消息传递需满足SE3群的对称性,通常通过张量积实现:

其中 ,是可学习的变换矩阵,k是SE3的表示维度。

2.3 解码器

2.3.1 结构预测头

用于预测三位坐标,二面角或距离矩阵。

\sigma^{(coord)}是激活函数,W_{c}^{(l)} 是坐标回归权重。

2.3.2 序列设计头

自回归生成(交叉熵损失) 逐个生成氨基酸序列,条件于当前节点特征和已生成序列:

是时间t的氨基酸类型, 为权重矩阵。

损失函数(序列恢复)

交叉熵损失对比预测序列与天然序列:

多任务学习

同时优化结构预测(如坐标误差)和序列设计(如恢复天然序列)任务,通过联合损失函数平衡不同目标。

三 实验

3.1 设备

Atlas 800T A2

3.2 组件版本

hdk:24.1.RC3

cann:8.0.RC3

python:3.10

torch:2.1.0

torch-npu:2.1.0.post8

3.3 环境准备

官方文档对依赖包没有版本要求,仅需保证python的版本高于等于3.0即可。

新建conda环境

conda create --name mlfold python=3.10

安装torch_npu和一些必要的依赖包,下载torch_npu

wget https://gitee.com/ascend/pytorch/releases/download/v6.0.rc3-pytorch2.1.0/torch_npu-2.1.0.post8-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl --no-check-certificate

pip3 install torch_npu-2.1.0.post8-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

注:该环境已经提前安装过了,所以不会提示安装成功

3.4 数据集下载

下载原始数据集

wget https://files.ipd.uw.edu/pub/training_sets/pdb_2021aug02.tar.gz

注:原始pdb_2021aug02数据集压缩包大小为17GB,解压之后为92GB,须预留足够的空间。下载之后使用tar -xzf pdb_2021aug02.tar.gz进行解压。

下载用于测试的数据集样本

如果没有足够的磁盘空间,可下载该数据集,压缩包为48MB,解压之后为255MB。

wget https://files.ipd.uw.edu/pub/training_sets/pdb_2021aug02_sample.tar.gz

3.5 推理

ProteinMPNN在example目录下面有很多样例,均可用来验证推理。

在training目录下也有test_inference.sh用于推理验证。

3.5.1 激活conda环境

conda activate mlfold

3.5.2 修改推理脚本

Slurm参数、用于输入的.PDB文件、权重路径、权重名称

3.5.3 修改protein_mpnn_run.py

3.5.4 执行推理

3.5.5 查看推理结果

3.6 训练

3.6.1 修改training.py文件

注:training.py脚本内,使用的cuda api不需要手动修改。torch_npu会自动将cuda api更改为torch_npu的api。例:

3.6.2 修改sh脚本

3.6.3 source cann

3.6.4 验证npu环境是否可用

3.6.5 开始训练

在training目录下面执行bash submit_exp_020.sh

查看npu状态,显示0卡正在被进程占用

注:可以通过环境变量export ASCEND_RT_VISIBLE_DEVICES来指定用哪张卡训练。

3.6.6 训练完成

训练完成之后,训练日志和权重会存放在exp_020目录下。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值