山东大学2019级软件工程应用与实践——基于人工智能的多肽药物分析问题（七）

2021SC@SDUSC

基于人工智能的多肽药物分析问题

主题：蛋白质预训练模型

学习论文：

《ProtTrans: Towards Cracking the Language of Life’s Code Through Self-Supervised Learning》

文章地址：

https://www.biorxiv.org/content/10.1101/2020.07.12.199554v3

源码地址：

https://github.com/agemagician/ProtTrans

摘要：

计算生物学和生物信息学为蛋白质序列提供了巨大的数据金矿，是NLP语言模型的理想选择。这些语言模型以较低的推理成本达到了新的预测前沿。在这里，我们根据包含3930亿氨基酸的UniRef和BFD的数据训练了两个自回归模型（Transformer XL、XLNet）和四个自动编码器模型（BERT、Albert、Electra、T5）。语言模型在Summit超级计算机上使用5616 GPU和多达1024核的TPU Pod进行训练。

降维显示，未标记数据中的原始蛋白质语言模型嵌入捕获了蛋白质序列的一些生物物理特征。我们验证了使用嵌入作为几个后续任务的独占输入的优势。第一个是蛋白质二级结构的per-residue预测（三态精度Q3=81%-87%）；第二个是蛋白质亚细胞定位的per-protein预测（十态精度：Q10=81%）和膜与水溶性（两态精度Q2=91%）。对于per-residue预测，在不使用进化信息的情况下，信息量最大的嵌入（ProtT5）的传输首次超过了最新技术，从而绕过了昂贵的数据库搜索。总之，研究结果表明蛋白质语言模型学习了一些生活语言的语法。

代码分析：

ProtTrans/Benchmark/ProtAlbert.ipynb

加载必要的库

pip install -q transformers

import torch
from transformers import AlbertModel
import time
from datetime import timedelta
import os
import requests
from tqdm.auto import tqdm

设置ProtAlbert和词汇表文件 (vocabulary file) 的url位置

modelUrl = 'https://www.dropbox.com/s/gtajtmege43ec7k/pytorch_model.bin?dl=1'
configUrl = 'https://www.dropbox.com/s/me7zsqrnpiz043v/config.json?dl=1'
tokenizerUrl = 'https://www.dropbox.com/s/60mg00r361vth4t/albert_vocab_model.model?dl=1'

下载ProtAlbert模型和词汇表文件

downloadFolderPath = 'models/ProtAlbert/'

modelFolderPath = downloadFolderPath
modelFilePath = os.path.join(modelFolderPath, 'pytorch_model.bin')
configFilePath = os.path.join(modelFolderPath, 'config.json')
tokenizerFilePath = os.path.join(modelFolderPath, 'spm_model.model')

if not os.path.exists(modelFolderPath):
    os.makedirs(modelFolderPath)
    
def download_file(url, filename):
  response = requests.get(url, stream=True)
  with tqdm.wrapattr(open(filename, "wb"), "write", miniters=1,
                    total=int(response.headers.get('content-length', 0)),
                    desc=filename) as fout:
      for chunk in response.iter_content(chunk_size=4096):
          fout.write(chunk)
          
if not os.path.exists(modelFilePath):
    download_file(modelUrl, modelFilePath)

if not os.path.exists(configFilePath):
    download_file(configUrl, configFilePath)

if not os.path.exists(tokenizerFilePath):
    download_file(tokenizerUrl, tokenizerFilePath)

加载 ProtAlbert 模型

model = AlbertModel.from_pretrained(modelFolderPath)

将模型加载到GPU（如果有），并切换到推理模式

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

model = model.to(device)
model = model.eval()

基准配置

min_batch_size = 8
max_batch_size = 32
inc_batch_size = 8

min_sequence_length = 64
max_sequence_length = 512
inc_sequence_length = 64

iterations = 10

开始基准测试

device_name = torch.cuda.get_device_name(device.index) if device.type == 'cuda' else 'CPU'

with torch.no_grad():
    print((' Benchmarking using ' + device_name + ' ').center(80, '*'))
    print(' Start '.center(80, '*'))
    for sequence_length in range(min_sequence_length,max_sequence_length+1,inc_sequence_length):
        for batch_size in range(min_batch_size,max_batch_size+1,inc_batch_size):
            start = time.time()
            for i in range(iterations):
                input_ids = torch.randint(1, 20, (batch_size,sequence_length)).to(device)
                results = model(input_ids)[0].cpu().numpy()
            end = time.time()
            ms_per_protein = (end-start)/(iterations*batch_size)
            print('Sequence Length: %4d \t Batch Size: %4d \t Ms per protein %4.2f' %(sequence_length,batch_size,ms_per_protein))
        print(' Done '.center(80, '*'))
    print(' Finished '.center(80, '*'))

基准测试结果