protpardelle

论文《全原子蛋白质生成模型》的代码。

该代码正在积极开发中，我们欢迎贡献、功能请求、问题报告、修正和任何问题！对于使用或改编自他人的代码，我们已尽力给予适当的归属，但如有任何需要更正的地方，请告诉我们。

环境和设置

要设置conda环境，请运行conda env create -f configs/environment.yml，然后conda activate delle。您还需要将ProteinMPNN仓库克隆到包含protpardelle/仓库的同一目录中。您可能还需要在使用的配置中将home_dir变量设置为包含protpardelle/目录的路径。

在WebApp和Pymol中使用

您可以直接在由Gradio支持的便捷HuggingFace WebApp中使用protpardelle。

或者，您可以直接在Pymol中设计蛋白质。为此，请将protpardelle_pymol.py文件下载到您的计算机上。

然后打开Pymol，并导航到文件所在的目录。

pwd (查看当前所在目录)
cd 文件路径

然后加载文件。如果是首次启动，这将在Pymol使用的Python安装中安装gradio_client。这可能需要几秒钟，Pymol窗口可能会出现冻结。

load protpardelle_pymol.py

条件设计

要运行条件设计，首先加载一个结构，例如

fetch 1pga

然后选择一些残基（并可选择命名该选择）。要重新采样名为sele的选择中的残基并生成5个样本，在Pymol控制台中运行以下命令：

protpardelle 1pga, sele, 5

无条件设计

要在长度50到60之间以步长5采样蛋白质，并为每个采样长度生成1个样本，使用以下命令：

protpardelle_uncond 50, 60, 5, 1

要使用仅骨架模型，使用：

protpardelle_uncond 50,60,5,1,backbone

推理

采样的入口点是draw_samples.py，它是sampling.py和model.sample()函数的便捷包装器。有许多参数可以用来控制模型检查点、采样配置和采样蛋白质的长度。我们提供了Protpardelle的仅骨架和全原子版本的模型权重。这两个版本都是无条件训练的；我们将在后续更新中发布条件模型。以下是一些绘制样本的示例。

用于绘制全原子样本的默认命令（例如，对于range(80, 100, 5)中的每个长度绘制2个蛋白质）：

python draw_samples.py --type allatom --minlen 80 --maxlen 100 --steplen 5 --perlen 2

要从仅骨架模型中为range(70, 150, 5)中的长度绘制每个长度8个样本，并使用100个去噪步骤而不是默认值，请运行以下命令。（这只是说明性的；我不期望用100个去噪步骤获得高质量的样本。）

python draw_samples.py --type backbone --param n_steps --paramval 100 --minlen 70 --maxlen 150 --steplen 5 --perlen 8

我们还添加了提供输入PDB文件和PDB文件中要条件的（从零开始索引的）索引列表的功能。请注意，当前模型仅支持单链，因此多链PDB将被视为单链（我们计划在后续更新中发布多链模型）。我们可以期望它在不同问题上表现更好或更差（在简单问题如修复上表现更好，在困难问题如不连续支架上表现更差）。使用以下命令重新采样my_pdb.pdb的前25个和第71至第80个残基。

python draw_samples.py --input_pdb my_pdb.pdb --resample_idxs 0-25,70-80

要对采样过程进行更多控制，包括调整采样超参数和更具体的条件方法，您可以直接与model.sample()函数交互；我们在sampling.py中提供了如何配置和运行这些命令的示例。

训练

注意（2023年9月）：实验室决定收集有关对训练自己版本的Protpardelle感兴趣的人的使用统计数据（用于资金和其他目的）。要获取包含训练代码的存储库副本，请填写此Google表单 - 您将收到一个Google Drive压缩文件的链接，其中包含带有训练代码的存储库。发布后，计划直接将完整的训练代码包含在此存储库中。

我们提供了预训练的模型权重，但如果您有兴趣训练自己的模型，我们提供了训练代码以及一些基本的在线评估。您需要创建一个Weights & Biases账户。

数据集可以从CATH下载，使用的训练/验证/测试拆分可以通过以下命令下载

wget http://people.csail.mit.edu/ingraham/graph-protein-design/data/cath/chain_set_splits.json

这些拆分中的一些PDB现在已经过时；我们手动将这些PDB替换为更新/新PDB ID的文件。数据加载器期望在数据集目录中有名为'train_pdb_keys.list'、'eval_pdbs_keys.list'和'test_pdb_keys.list'的文本文件，这些文件列出了与每个数据集拆分相关的文件名。这与PDB文件目录一起足以满足数据加载器的需求。

主要入口点是train.py；有一些参数可以控制计算、实验等。模型特定的训练代码与训练基础设施分开，由runners.py中的runner类处理；模型相关的超参数由配置文件处理。使用configs/backbone.yml训练仅骨架模型；configs/allatom.yml训练全原子模型，configs/seqdes.yml训练迷你MPNN模型。一些例子：

默认命令（用于生成保存的模型权重）：

python train.py --project protpardelle --train --config configs/allatom.yml --num_workers 8

对于迷你MPNN模型的简单调试运行：

python train.py --config configs/seqdes.yml

要使用8个数据加载工作线程过拟合到100个数据示例，对于具有2层的裁剪条件骨架模型，在configs/backbone.yml中将train.crop_conditional和model.crop_conditional更改为True，然后运行：

python train.py --train --config configs/backbone.yml --overfit 100 --num_workers 8

使用DDP进行训练稍微复杂一些，需要使用torch.distributed。请注意，配置中的批量大小成为每设备批量大小。要在单个节点上的2个GPU上使用DDP训练全原子模型，请运行：

python -m torch.distributed.run --standalone --nnodes=1 --nproc_per_node=2 --master_port=$RANDOM train.py --config configs/allatom.yml --train --n_gpu_per_node 2 --use_ddp --num_workers 8

引用

如果您觉得我们的工作有帮助，请引用

@article {chu2023allatom,
    author = {Alexander E. Chu and Lucy Cheng and Gina El Nesr and Minkai Xu and Po-Ssu Huang},
    title = {An all-atom protein generative model},
    year = {2023},
    doi = {10.1101/2023.05.24.542194},
    URL = {https://www.biorxiv.org/content/early/2023/05/25/2023.05.24.542194},
    journal = {bioRxiv}
}