PSP - 解决 ESMFold 推理长序列蛋白质结构的显存溢出问题

欢迎关注我的CSDN：https://spike.blog.csdn.net/
本文地址：https://spike.blog.csdn.net/article/details/134709211

使用 ESMFold 推理长序列 (Seq. Len. > 1500) 时，导致显存不足，需要设置 chunk_size 参数，实现长序列蛋白质的结构预测，避免显存溢出。

ESMFold：https://github.com/facebookresearch/esm

测试 ESM 单条 Case，序列长度 1543 较长，即：

python -u myscripts/esmfold_infer.py \
-f fasta_446/7WY5_R1543.fasta \
-o mydata/test_gpcr/

A100 显存溢出：

Tried to allocate 54.74 GiB (GPU 0; 79.32 GiB total capacity; 73.53 GiB already allocated; 3.94 GiB free; 74.24 GiB reserved in total by PyTorch)

解决显存问题，参考：Out of memory - upper limit on sequence length?

关键参数：chunk-size

Chunks axial attention computation to reduce memory usage from O(L^2) to O(L). Equivalent to running a for loop over chunks of of each dimension. Lower values will result in lower memory usage at the cost of speed. Recommended values: 128, 64, 32. Default: None.将轴向注意力计算分块 (Chunks) ，将内存使用量从 O(L^2) 减少到 O(L)。 相当于在每个维度的块上运行 for 循环。 较低的值将导致内存使用量降低，但代价是速度。 建议值：128、64、32。默认值：无。

关键参数：max-tokens-per-batch，即 max_tokens_per_batch

Maximum number of tokens per gpu forward-pass. This will group shorter sequences together for batched prediction. Lowering this can help with out of memory issues, if these occur on short sequences.每个 GPU 前向传递的最大令牌数。 这会将较短的序列分组在一起以进行批量预测。 如果内存不足问题发生在短序列上，降低此值可以帮助解决这些问题。

将 chunk-size 设置成 128，问题解决，即：

max_len = 1200
# A100 最多支持 1200 长度的序列
if len(seq) > max_len:chunk_size = 128print(f"[Warning] seq length is too long! {len(seq)} > {max_len}, chunk_size: {chunk_size}")self.model.set_chunk_size(chunk_size)
else:self.model.set_chunk_size(None)with torch.no_grad():output = self.model.infer_pdb(seq)

推理脚本：

#!/usr/bin/env python
# -- coding: utf-8 --
"""
Copyright (c) 2022. All rights reserved.
Created by C. L. Wang on 2023/7/5
"""
import argparse
import os
import sys
import time
from pathlib import Pathimport torch
from tqdm import tqdmimport esmp = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if p not in sys.path:sys.path.append(p)from myutils.protein_utils import get_seq_from_fasta
from myutils.project_utils import time_elapsed, mkdir_if_not_exist, traverse_dir_filesclass EsmfoldInfer(object):"""ESMFold的推理类"""def __init__(self):print("[Info] 开始加载 ESMFold 模型!")s_time = time.time()model = esm.pretrained.esmfold_v1()self.model = model.eval().cuda()print(f"[Info] vocab: {self.model.esm_dict.to_dict()}")# 耗时: 00:01:13.264272print(f"[Info] 完成加载 ESMFold 模型! 耗时: {time_elapsed(s_time, time.time())}")def predict_seq(self, seq, out_path, is_log=True):"""预测序列"""print(f"[Info] seq_len: {len(seq)}")max_len = 1200# A100 最多支持 1200 长度的序列if len(seq) > max_len:chunk_size = 128print(f"[Warning] seq length is too long! {len(seq)} > {max_len}, chunk_size: {chunk_size}")self.model.set_chunk_size(chunk_size)else:self.model.set_chunk_size(None)s_time = time.time()with torch.no_grad():output = self.model.infer_pdb(seq)seq_len = len(seq)if is_log:print(f"[Info] 完成推理，链长 {seq_len}, 耗时: {time_elapsed(s_time, time.time())}, "f"平均序列耗时: {(time.time() - s_time) / seq_len}")with open(out_path, "w") as f:f.write(output)if is_log:print(f"[Info] 输出: {output}")def predict_fasta_dir(self, input_path, output_dir):"""预测 FASTA 文件夹"""print(f"[Info] input_path: {input_path}")print(f"[Info] output_dir: {output_dir}")assert os.path.isfile(input_path) or os.path.isdir(input_path)mkdir_if_not_exist(output_dir)if os.path.isdir(input_path):path_list = traverse_dir_files(input_path, ext="fasta")elif os.path.isfile(input_path):path_list = [input_path]else:raise Exception(f"Error input: {input_path}")print(f"[Info] Fasta 数量: {len(path_list)}")s_time = time.time()for path in tqdm(path_list, desc="[Info] fasta"):fasta_name = os.path.basename(path).split(".")[0]output_fasta_dir = os.path.join(output_dir, fasta_name)mkdir_if_not_exist(output_fasta_dir)pdb_name = os.path.basename(path).replace("fasta", "pdb")output_pdb_path = os.path.join(output_fasta_dir, pdb_name)if os.path.exists(output_pdb_path):print(f"[Info] 已预测完成: {output_pdb_path}")continueseqs, _ = get_seq_from_fasta(path)seq = seqs[0]self.predict_seq(seq, output_pdb_path, is_log=False)print(f"[Info] 全部运行完成: {output_dir}, 耗时: {time_elapsed(s_time, time.time())}")def main():parser = argparse.ArgumentParser()parser.add_argument("-f","--fasta-input",type=Path,required=True,)parser.add_argument("-o","--output-dir",type=Path,required=True)args = parser.parse_args()fasta_input = str(args.fasta_input)output_dir = str(args.output_dir)mkdir_if_not_exist(output_dir)ei = EsmfoldInfer()ei.predict_fasta_dir(fasta_input, output_dir)if __name__ == '__main__':main()