-
ESM-2(Evolutionary Scale Modeling 2)当下最先进的预训练蛋白质语言模型之一， 由Facebook AI
Research开发，最新版本使用了48层Transformer编码器架构，有150亿参数。 -
蛋白质语言模型(PLM)可以用于理解和预测蛋白质序列的特性，包括它们的结构、功能等。
以下是使用ESM-2的详细教程：
# pip install esm
# pip install torch
# pip安装的esm不包括模型，size is small，but torch is big, more than 1G.
import torch
import esm
import os
我们先指定一个新的目录路径，然后写入环境变量字典，设置TORCH_HOME环境变量
-
为了方便管理模型，我们可以设置TORCH_HOME环境变量，将模型下载到我们指定的目录；
-
在运行esm.pretrained.esm2_t33_650M_UR50D()时，PyTorch会检查这个目录，如果模型esm2_t33_650M_UR50D已经存在，它会从那里加载模型，否则它会从网上下载模型并保存在这个目录下。
-
例如：当环境变量设置为’D:\Desktop\model’时，模型的下载地址为’D:\Desktop\model\hub\checkpoints\esm2_t33_650M_UR50D.pt’
-
设置 TORCH_HOME 环境变量后，所有 PyTorch 相关的库（比如 torch.hub 或 transformers）在下载模型和数据集时，都会使用这个目录作为下载位置。
-
注：
PyTorch和ESM都是Facebook的产品；
os.environ 返回一个代表当前环境变量的字典对象。
new_dir = 'D:\Desktop\model'
os.environ['TORCH_HOME'] = new_dir
下载模型到我们上面指定的目录，或者从指定的目录加载模型；alphabet代表模型使用的字母表，它定义了模型能够处理的字符集合。
-
Size of esm2_t33_650M_UR50D is very big, about 2.4G, 这个模型使用了33层Transformer编码器架构，有650百万(6.5亿参数)，使用UniRef50作为训练集;
-
关于UniRef100、UniRef90和UniRef50的知识，请参考：https://pubmed.ncbi.nlm.nih.gov/17379688/
-
esm包含好几个版本的蛋白质语言预训练模型，可以通过esm.pretrained.xxx指定使用不同版本：esm.pretrained.esm2_t36_3B_UR50D…
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
构建数据格式转换器
batch_converter = alphabet.get_batch_converter()
将模型设置为评估模式，这会关闭dropout等训练特有的行为
-
在神经网络中，“dropout” 是一种正则化技术，用于防止或减少模型的过拟合。Dropout通过在训练过程中随机"丢弃"（即暂时移除）网络中的一些神经元（包括它们所有的连接），来减少神经元之间复杂的共适应关系，从而促进模型的泛化能力。
-
在训练神经网络时启用dropout，以减少过拟合。
-
在模型评估或预测时禁用dropout，确保所有神经元都参与工作。
model.eval()
demo数据
data = [("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),("protein2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE")
]
使用数据格式转换器将数据转换为模型可以理解/处理的tokens，并将这些tokens填充为相同的长度，然后批量计算每个序列的长度
-
batch_labels：存储每个序列的标签或ID，可能用于后续的监督学习任务；
-
batch_strs：存储原始的蛋白质序列字符串，可能用于调试或显示目的；
-
batch_tokens：是一个二维Tensor张量，一行代表一个蛋白质序列，存储转换后的tokens，这些tokens是原始序列中氨基酸的整数索引(0-20，20种氨基酸)；
-
alphabet.padding_idx是一个填充矩阵；
-
batch_tokens != alphabet.padding_idx 生成一个布尔矩阵，实际有氨基酸的位置是True，填充的位置是False；
-
.sum(1) 对这个布尔矩阵沿着维度1求和，即生成每个序列的长度；
-
这样做的目的是批量求每个序列的长度。
batch_labels, batch_strs, batch_tokens = batch_converter(data)
batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)
提取蛋白质序列中每个氨基酸残基的 token_representations，获取模型对输入序列的深层次理解
-
with torch.no_grad(): 禁用PyTorch中的梯度计算。在模型评估或预测阶段，我们不需要进行反向传播，禁用梯度计算可以减少内存消耗并加速计算。使用with表明该操作是暂时的。
-
repr_layers=[33] 表示使用Transformer架构第33层的输出作为特征representation；
-
return_contacts=True 表示获取模型预测的氨基酸残基之间的接触图，这在蛋白质结构预测中是一个有用的特征；
-
results[“representations”][33] 表示从模型输出results中提取第33层的representations；
-
results是一个字典，"representations"是这个字典的一个键(key)，该键对应的值(value)也是一个字典，存放着每一层的token_representations。字典中嵌套字典。
with torch.no_grad():results = model(batch_tokens, repr_layers=[33], return_contacts=True)
token_representations = results["representations"][33]
生成每一个序列的 representations
- token 0是序列开始标记, 所以第一个氨基酸残基是token 1；最后一个token是序列结束标记；
- token_representations 是一个三维Tensor张量(可以理解为三维数组)；
- token_representations[i, 1:tokens_len-1] 使用切片取出每条蛋白质序列的token_representations(二维，根据氨基酸数量取行，列全取)；
- .mean(0) 使用每条蛋白质序列的二维token_representations的列平均值作为该蛋白质序列的representations(sequence_representations)
sequence_representations = []
for i, tokens_len in enumerate(batch_lens): # 例如，氨基酸token为65个，则batch_lens=67sequence_representations.append(token_representations[i, 1:tokens_len-1].mean(0))
生成使用无监督学习方法预测的蛋白质内部残基间的接触图
import matplotlib.pyplot as plt
for (ID, seq), tokens_len, attention_contacts in zip(data, batch_lens, results["contacts"]):plt.matshow(attention_contacts[: tokens_len, : tokens_len])plt.title(ID)path = os.path.join(os.getcwd(), ID)plt.savefig(path, bbox_inches="tight")
参考：https://github.com/facebookresearch/esm
