蛋白质嵌入评估任务（TAPE）

用于在训练好的蛋白质嵌入模型上运行 TAPE 基准测试的数据、权重和代码。我们提供了预训练语料库、五个有监督下游任务、预训练语言模型权重以及基准测试代码。该代码现已更新为使用 PyTorch，因此之前的预训练模型权重和代码将无法使用。旧的 TensorFlow 版本 TAPE 仓库仍可在 https://github.com/songlab-cal/tape-neurips2019 找到。

本仓库并非旨在与原始论文保持最大程度的兼容性和可重复性，而是为了便于使用及未来的开发工作（无论是对我们自身还是对社区）。尽管我们提供了许多相同的功能，但我们并未对所有模型和下游任务的训练进行全方位测试，并且也做出了一些有意的改动。因此，如果您希望复现我们论文中的结果，请使用原始代码。

我们的论文可在 https://arxiv.org/abs/1906.08230 获取。

部分文档尚不完整。我们会逐步补充，但如果您对某些内容需要解释，请提交一个问题，以便我们明确优先改进的方向！

更新：2020年9月26日： 我们不再建议直接使用 TAPE 的训练代码进行训练。虽然它在一段时间内可能仍然有效，但不会针对未来版本的 PyTorch 进行更新。在内部，我们一直在使用不同的框架进行训练（特别是 PyTorch Lightning 和 Fairseq）。我们强烈建议使用类似这些框架的工具，因为它们可以减轻维护与 PyTorch 版本兼容性的负担。TAPE 模型将继续提供，如果当前代码对您有效，您可以继续使用。然而，我们将不再修复训练过程中出现的多 GPU 错误、内存不足错误等问题。

目录

• 安装
• 示例
  • Hugging Face API 加载预训练模型
  • 使用预训练模型嵌入蛋白质
  • 训练语言模型
  • 评估语言模型
  • 训练下游模型
  • 评估下游模型
  • 模型和任务列表
  • 添加新模型和任务
• 数据
  • LMDB 数据
  • 原始数据
• 排行榜
  • 二级结构
  • 接触预测
  • 远程同源性检测
  • 荧光
  • 稳定性
• 引用指南

安装

我们建议您使用以下命令将 tape 安装到 Python 的 虚拟环境 中：

$ pip install tape_proteins

示例

Hugging Face API 加载预训练模型

我们基于优秀的 Hugging Face 库，并将其用作定义模型和提供预训练模型的 API。通过此 API，预训练模型将在需要时自动下载，并缓存以供后续使用。

import torch
from tape import ProteinBertModel, TAPETokenizer
model = ProteinBertModel.from_pretrained('bert-base')
tokenizer = TAPETokenizer(vocab='iupac')  # iupac 是 TAPE 模型的词汇表，UniRep 模型则使用 unirep
sequence = 'GCTVEDRCLIGMGAILLNGCVIGSGSLVAAGALITQ'
token_ids = torch.tensor([tokenizer.encode(sequence)])
output = model(token_ids)
sequence_output = output[0]
pooled_output = output[1]

# 注意：pooled_output 并未针对 Transformer 模型进行训练，未经微调请勿使用。目前更好的做法是直接取 sequence_output 的平均值。

目前可用的预训练模型包括：

• bert-base（Transformer 模型）
• babbler-1900（UniRep 模型）
• xaa、xab、xac、xad、xae（trRosetta 模型）

如果您希望使用特定的预训练模型，请提交一个问题，我们将尽力添加！

使用预训练模型嵌入蛋白质

给定一个输入的 FASTA 文件，您可以通过 tape-embed 命令生成包含蛋白质嵌入的 .npz 文件。

假设这是我们的输入 FASTA 文件：

>seq1
GCTVEDRCLIGMGAILLNGCVIGSGSLVAAGALITQ
>seq2
RTIKVRILHAIGFEGGLMLLTIPMVAYAMDMTLFQAILLDLSMTTCILVYTFIFQWCYDILENR

那么我们可以使用 UniRep 的 babbler-1900 模型对其进行嵌入：

tape-embed unirep my_input.fasta output_filename.npz babbler-1900 --tokenizer unirep

无需手动下载预训练模型——如果需要，它会自动下载。此外，请注意将分词器更改为 unirep 分词器。UniRep 使用不同的词汇表，因此需要使用此分词器。如果遇到 cublas 运行时错误，请务必检查是否正确切换了分词器。

嵌入功能已完全批处理，并会根据机器可用的 GPU 数量自动分配任务。在 Titan Xp 上，每秒大约可以处理 200 条序列。

获得输出文件后，我们可以使用 NumPy 加载它：

arrays = np.load('output_filename.npz', allow_pickle=True)

list(arrays.keys())  # 将输出 FASTA 文件中各条序列的名称（或未命名时为 '0'、'1' 等）。
arrays[<protein_id>]  # 返回一个字典，包含 'pooled' 和 'avg' 键（或使用 --full_sequence_embed 标志时为 'seq'）。

默认情况下，为节省内存，TAPE 会返回序列嵌入的平均值以及通过池化函数生成的池化嵌入。对于某些模型（如 UniRep），池化嵌入经过训练，可以直接使用。而对于其他模型（如 Transformer），池化嵌入未经训练，因此应使用平均嵌入。未来，我们将探索对所有模型的池化嵌入进行自监督训练的方法。

如果您希望获取完整的嵌入而不是平均嵌入，可以在运行 tape-embed 时添加 --full_sequence_embed 标志。

训练语言模型

Tape 提供了两个训练命令：tape-train 和 tape-train-distributed。第一个命令使用标准的 PyTorch 数据分发方式，在所有可用的 GPU 上进行分布式训练。第二个命令则采用类似于 torch.distributed.launch 的多进程机制，按照指定的 GPU 数量进行分布式训练（也可以用于跨多个节点的分布式训练）。我们通常建议使用第二个命令，因为它可以带来 10% 到 15% 的速度提升，但两者都可以正常工作。

例如，要对 Transformer 模型进行掩码语言建模训练，可以运行以下命令：

tape-train-distributed transformer masked_language_modeling --batch_size BS --learning_rate LR --fp16 --warmup_steps WS --nproc_per_node NGPU --gradient_accumulation_steps NSTEPS

在训练过程中，Tape 使用了多种特性：

• 多进程分布式训练
• 半精度训练
• 梯度累积
• 累积后的梯度 All-Reduce 操作
• 根据序列长度自动调整批次大小

你最可能需要的第一个特性是 gradient_accumulation_steps。默认情况下，Tape 设置了一个相对较高的批次大小（1024）。这个批次大小是指每次反向传播时使用的批次大小。实际每块 GPU 上处理的样本数量会等于总批次大小除以 GPU 数量再除以梯度累积步数。因此，如果批次大小为 1024、使用 2 块 GPU 且梯度累积步数为 1，则每块 GPU 将处理 512 个样本。如果你遇到显存不足的情况（很可能发生），Tape 会给出明确的错误提示，并建议你增加梯度累积步数。

此外还有其他一些未在此提及的功能。有关所有命令的详细信息，请参阅 tape-train-distributed --help。

评估语言模型

训练好语言模型后，你会在 results 文件夹中找到预训练权重文件。要评估该模型，你可以采取两种方法。一种是直接评估语言建模的准确率或困惑度。tape-train 会在每个 epoch 结束时报告训练集和验证集上的困惑度。然而，根据经验，我们认为语言建模的准确率和困惑度并不是衡量下游任务性能的良好指标。因此，为了更全面地评估语言模型，我们强烈建议你在我们提供的一个或多个下游任务上继续训练你的模型。

训练下游模型

使用 tape-train 命令同样可以训练下游任务模型。只需沿用与训练语言模型相同的语法，并添加 --from_pretrained <path_to_your_saved_results> 标志即可。例如，要在一个预训练的 Transformer 模型上进行二级结构预测训练，可以运行如下命令：

tape-train-distributed transformer secondary_structure \
	--from_pretrained results/<path_to_folder> \
	--batch_size BS \
	--learning_rate LR \
	--fp16 \
  	--warmup_steps WS \
  	--nproc_per_node NGPU \
  	--gradient_accumulation_steps NSTEPS \
  	--num_train_epochs NEPOCH \
  	--eval_freq EF \
  	--save_freq SF

在训练下游模型时，你可能需要不断尝试不同的超参数组合，以获得最佳效果（不同任务和不同模型的最佳超参数可能会有所不同）。需要考虑的参数包括：

* 批次大小
* 学习率
* 预热步数
* 训练轮数

这些参数都会显著影响模型性能，默认设置则是为了最大化语言建模任务的表现，而非下游任务。此外，eval_freq 和 save_freq 参数也很有用，它们分别用于减少验证频率和保存模型的频率。由于下游任务的训练轮数通常较短（而且往往需要更多轮次），适当提高这两个参数可以使训练过程更加高效。

评估下游模型

为了评估你的下游任务模型，我们提供了 tape-eval 命令。该命令会输出你的模型预测结果以及你指定的一组评估指标。目前支持均方误差 (mse)、平均绝对误差 (mae)、斯皮尔曼等级相关系数 (spearmanr) 和准确率 (accuracy)。L/5 精确度指标也将很快加入。

命令的语法如下：

tape-eval MODEL TASK TRAINED_MODEL_FOLDER --metrics METRIC1 METRIC2 ...

例如，要评估一个经过二级结构训练的 Transformer 模型，可以运行：

tape-eval transformer secondary_structure results/<path_to_trained_model> --metrics accuracy

这将报告整体准确率，并在训练好的模型目录中生成一个 results.pkl 文件，供你进一步分析。

trRosetta

我们最近重新实现了 Yang 等人（2020）提出的 trRosetta 模型。原始仓库的链接在这里，它是本次重实现的基础。我们提供了基于 PyTorch 的实现及数据集，方便大家探索和实验该模型。数据可从这里下载。这些数据与原论文中的数据相同，但我们额外提供了训练集和验证集的划分文件，以便你可以复现并训练自己的模型。使用该模型的方法如下：

from tape import TRRosetta
from tape.datasets import TRRosettaDataset

# 下载数据并将其放置在 `<data_path>/trrosetta` 目录下

train_data = TRRosettaDatset('<data_path>', 'train')  # 会对 MSA 进行采样
valid_data = TRRosettaDatset('<data_path>', 'valid')  # 不会对 MSA 进行采样

model = TRRosetta.from_pretrained('xaa')  # 可选值为 'xaa', 'xab', 'xac', 'xad', 'xae'，分别对应不同的集成模型。

batch = train_data.collate_fn([train_data[0]])
loss, predictions = model(**batch)

预测结果可以保存为 .npz 文件，然后输入到 Yang 实验室提供的结构建模脚本中进行进一步处理。

模型与任务列表

当前可用的模型包括：

• transformer（提供预训练模型）
• resnet
• lstm
• unirep（提供预训练模型）
• onehot（无需预训练）
• trrosetta（提供预训练模型）

当前可用的标准任务包括：

• language_modeling
• masked_language_modeling
• secondary_structure
• contact_prediction
• remote_homology
• fluorescence
• stability
• trrosetta（仅适用于 trrosetta 模型）

具体的模型和任务列表可以在 tape/datasets.py 和 tape/models/modeling*.py 中找到。

添加新模型和新任务

我们已努力使新的代码库更易于理解和扩展。请参阅 examples 文件夹，其中包含如何向 TAPE 添加新模型和新任务的示例。如果你希望看到更多示例，或者认为现有示例有所欠缺，请提交一个问题。

数据

数据应放置在 ./data 文件夹中，不过您也可以根据需要指定其他数据目录。

有监督数据压缩后约为 120MB，解压后约为 2GB。 无监督 Pfam 数据集压缩后约为 7GB，解压后约为 19GB。用于训练的数据托管在 AWS 上。默认情况下，我们以 LMDB 格式提供数据——有关加载数据的示例，请参阅 tape/datasets.py。如果您希望下载 TAPE 的所有数据，可以运行 download_data.sh 脚本来完成下载。我们还在下方提供了每个数据集的 LMDB 和 JSON 格式链接。

LMDB 数据

预训练语料（Pfam） | 二级结构 | 接触预测（ProteinNet） | 远程同源性 | 荧光 | 稳定性

原始数据

原始数据文件以 JSON 格式存储，以实现最大的可移植性。这些数据已被转换为 JSON 格式，因此移除了某些结构（尤其是 NumPy 数组）。因此，它们无法直接加载到提供的 PyTorch 数据集中（尽管只需添加对 np.array 的调用即可轻松完成转换）。

预训练语料（Pfam） | 二级结构 | 接触预测（ProteinNet） | 远程同源性 | 荧光 | 稳定性

排行榜

我们很快将推出一个排行榜，用于跟踪 TAPE 核心五项任务的进展，请稍后再查看此处的链接。有关当前性能水平的参考信息，请参阅我们论文中的主要表格。排行榜的发布将以遵守以下引用规范为前提。

在此期间，以下是各任务的临时排行榜。该排行榜上报告的所有模型均使用了无监督预训练。

二级结构

排名	模型	准确率（3 类）
1.	One Hot + Alignment	0.80
2.	LSTM	0.75
2.	ResNet	0.75
4.	Transformer	0.73
4.	Bepler	0.73
4.	Unirep	0.73
7.	One Hot	0.69

接触预测

排名	模型	L/5 中长程
1.	One Hot + Alignment	0.64
2.	Bepler	0.40
3.	LSTM	0.39
4.	Transformer	0.36
5.	Unirep	0.34
6.	ResNet	0.29
6.	One Hot	0.29

远程同源性检测

排名	模型	Top 1 准确率
1.	LSTM	0.26
2.	Unirep	0.23
3.	Transformer	0.21
4.	Bepler	0.17
4.	ResNet	0.17
6.	One Hot + Alignment	0.09
6.	One Hot	0.09

荧光

排名	模型	斯皮尔曼等级相关系数
1.	Transformer	0.68
2.	LSTM	0.67
2.	Unirep	0.67
4.	Bepler	0.33
5.	ResNet	0.21
6.	One Hot	0.14

稳定性

排名	模型	斯皮尔曼等级相关系数
1.	Transformer	0.73
1.	Unirep	0.73
1.	ResNet	0.73
4.	LSTM	0.69
5.	Bepler	0.64
6.	One Hot	0.19

引用指南

如果您觉得 TAPE 有用，请引用我们的相关论文。此外，任何使用 TAPE 中提供的数据集的人员都必须说明并引用其所使用的全部数据集组件。生成这些数据需要耗费大量时间和资源，我们坚持要求所有 TAPE 用户对此予以认可。为方便起见，data_refs.bib 文件中包含了所有必要的引用信息。我们也在下方提供了每项数据集的具体引用。

TAPE（我们的论文）：

@inproceedings{tape2019,
author = {Rao, Roshan and Bhattacharya, Nicholas and Thomas, Neil and Duan, Yan and Chen, Xi and Canny, John and Abbeel, Pieter and Song, Yun S},
title = {Evaluating Protein Transfer Learning with TAPE},
booktitle = {Advances in Neural Information Processing Systems}
year = {2019}
}

Pfam（预训练）：

@article{pfam,
author = {El-Gebali, Sara and Mistry, Jaina and Bateman, Alex and Eddy, Sean R and Luciani, Aur{\'{e}}lien and Potter, Simon C and Qureshi, Matloob and Richardson, Lorna J and Salazar, Gustavo A and Smart, Alfredo and Sonnhammer, Erik L L and Hirsh, Layla and Paladin, Lisanna and Piovesan, Damiano and Tosatto, Silvio C E and Finn, Robert D},
doi = {10.1093/nar/gky995},
file = {::},
issn = {0305-1048},
journal = {Nucleic Acids Research},
keywords = {community,protein domains,tandem repeat sequences},
number = {D1},
pages = {D427--D432},
publisher = {Narnia},
title = {{The Pfam protein families database in 2019}},
url = {https://academic.oup.com/nar/article/47/D1/D427/5144153},
volume = {47},
year = {2019}
}

SCOPe：（远程同源性和接触）-

@article{scop,
  title={SCOPe: Structural Classification of Proteins—extended, integrating SCOP and ASTRAL data and classification of new structures},
  author={Fox, Naomi K and Brenner, Steven E and Chandonia, John-Marc},
  journal={Nucleic acids research},
  volume={42},
  number={D1},
  pages={D304--D309},
  year={2013},
  publisher={Oxford University Press}
}

PDB：（二级结构和接触）

@article{pdb,
  title={The protein data bank},
  author={Berman, Helen M and Westbrook, John and Feng, Zukang and Gilliland, Gary and Bhat, Talapady N and Weissig, Helge and Shindyalov, Ilya N and Bourne, Philip E},
  journal={Nucleic acids research},
  volume={28},
  number={1},
  pages={235--242},
  year={2000},
  publisher={Oxford University Press}
}

CASP12：（二级结构和接触）

@article{casp,
author = {Moult, John and Fidelis, Krzysztof and Kryshtafovych, Andriy and Schwede, Torsten and Tramontano, Anna},
doi = {10.1002/prot.25415},
issn = {08873585},
journal = {Proteins: Structure, Function, and Bioinformatics},
keywords = {CASP,community wide experiment,protein structure prediction},
pages = {7--15},
publisher = {John Wiley {\&} Sons, Ltd},
title = {{Critical assessment of methods of protein structure prediction (CASP)-Round XII}},
url = {http://doi.wiley.com/10.1002/prot.25415},
volume = {86},
year = {2018}
}

NetSurfP2.0：（二级结构）

@article{netsurfp,
  title={NetSurfP-2.0: Improved prediction of protein structural features by integrated deep learning},
  author={Klausen, Michael Schantz and Jespersen, Martin Closter and Nielsen, Henrik and Jensen, Kamilla Kjaergaard and Jurtz, Vanessa Isabell and Soenderby, Casper Kaae and Sommer, Morten Otto Alexander and Winther, Ole and Nielsen, Morten and Petersen, Bent and others},
  journal={Proteins: Structure, Function, and Bioinformatics},
  year={2019},
  publisher={Wiley Online Library}
}

ProteinNet：（接触）

@article{proteinnet,
  title={ProteinNet: a standardized data set for machine learning of protein structure},
  author={AlQuraishi, Mohammed},
  journal={arXiv preprint arXiv:1902.00249},
  year={2019}
}

荧光：

@article{sarkisyan2016,
  title={Local fitness landscape of the green fluorescent protein},
  author={Sarkisyan, Karen S and Bolotin, Dmitry A and Meer, Margarita V and Usmanova, Dinara R and Mishin, Alexander S and Sharonov, George V and Ivankov, Dmitry N and Bozhanova, Nina G and Baranov, Mikhail S and Soylemez, Onuralp and others},
  journal={Nature},
  volume={533},
  number={7603},
  pages={397},
  year={2016},
  publisher={Nature Publishing Group}
}

稳定性：

@article{rocklin2017,
  title={Global analysis of protein folding using massively parallel design, synthesis, and testing},
  author={Rocklin, Gabriel J and Chidyausiku, Tamuka M and Goreshnik, Inna and Ford, Alex and Houliston, Scott and Lemak, Alexander and Carter, Lauren and Ravichandran, Rashmi and Mulligan, Vikram K and Chevalier, Aaron and others},
  journal={Science},
  volume={357},
  number={6347},
  pages={168--175},
  year={2017},
  publisher={American Association for the Advancement of Science}
}

TAPE 快速上手指南

TAPE (Tasks Assessing Protein Embeddings) 是一个用于评估蛋白质嵌入模型的基准工具包，提供预训练语料库、五个监督下游任务、预训练语言模型权重及基准测试代码。本项目已迁移至 PyTorch 框架。

环境准备

• 系统要求：Linux 或 macOS（Windows 支持未明确提及，建议优先使用 Linux）。
• Python 版本：建议使用 Python 3.6+。
• 前置依赖：
  • PyTorch (最新稳定版)
  • CUDA (如需使用 GPU 加速)
  • 建议创建独立的虚拟环境以避免依赖冲突。

安装步骤

推荐使用 pip 在虚拟环境中安装官方发布的包：

# 创建并激活虚拟环境 (可选但推荐)
python -m venv tape_env
source tape_env/bin/activate  # Linux/macOS
# tape_env\Scripts\activate   # Windows

# 安装 tape_proteins
pip install tape_proteins

注意：国内用户若下载缓慢，可尝试添加清华或阿里镜像源： pip install tape_proteins -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

1. 加载预训练模型并生成嵌入 (Python API)

TAPE 集成了 Hugging Face Transformers API，可自动下载并缓存预训练模型。以下示例展示如何加载 bert-base 模型并对蛋白质序列进行编码：

import torch
from tape import ProteinBertModel, TAPETokenizer

# 加载预训练模型 (自动下载 'bert-base')
model = ProteinBertModel.from_pretrained('bert-base')

# 初始化分词器 (TAPE 模型使用 'iupac' 词汇表)
tokenizer = TAPETokenizer(vocab='iupac')

# 示例蛋白质序列
sequence = 'GCTVEDRCLIGMGAILLNGCVIGSGSLVAAGALITQ'

# 编码序列并转换为 Tensor
token_ids = torch.tensor([tokenizer.encode(sequence)])

# 前向传播获取输出
output = model(token_ids)
sequence_output = output[0]  # 序列级输出
pooled_output = output[1]    # 池化输出

# 注意：对于 Transformer 模型，pooled_output 未经过专门训练。
# 建议直接使用 sequence_output 的均值作为序列表示，除非进行了微调。

可用预训练模型列表：

• bert-base: Transformer 模型
• babbler-1900: UniRep 模型
• xaa, xab, xac, xad, xae: trRosetta 模型

2. 批量生成蛋白质嵌入 (命令行工具)

若需处理 FASTA 文件中的大量序列，可使用 tape-embed 命令生成 .npz 格式的特征文件。

输入文件 (input.fasta)：

>seq1
GCTVEDRCLIGMGAILLNGCVIGSGSLVAAGALITQ
>seq2
RTIKVRILHAIGFEGGLMLLTIPMVAYAMDMTLFQAILLDLSMTTCILVYTFIFQWCYDILENR

执行命令： 以下命令使用 UniRep (babbler-1900) 模型进行嵌入。注意 UniRep 需指定 --tokenizer unirep。

tape-embed unirep input.fasta output_filename.npz babbler-1900 --tokenizer unirep

• 模型会自动下载（如需）。
• 支持多 GPU 自动并行加速。
• 默认输出包含平均嵌入 (avg) 和池化嵌入 (pooled)。如需完整序列嵌入，请添加 --full_sequence_embed 标志。

读取结果：

import numpy as np

arrays = np.load('output_filename.npz', allow_pickle=True)
print(list(arrays.keys()))  # 查看序列 ID (如 'seq1', 'seq2')

# 获取特定蛋白质的嵌入字典 (包含 'pooled' 和 'avg')
data = arrays['seq1'] 

3. 下游任务微调示例

使用预训练模型在下游任务（如二级结构预测）上进行微调：

tape-train-distributed transformer secondary_structure \
	--from_pretrained results/<path_to_saved_model> \
	--batch_size 32 \
	--learning_rate 1e-4 \
	--fp16 \
  	--warmup_steps 100 \
  	--nproc_per_node 2 \
  	--gradient_accumulation_steps 4 \
  	--num_train_epochs 10

• --from_pretrained: 指向预训练权重路径。
• --gradient_accumulation_steps: 显存不足时增大此值以模拟更大 batch size。
• --fp16: 启用半精度训练以节省显存并加速。