PSP - 蛋白质结构预测 OpenFold Multimer 重构训练模型的数据加载

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

OpenFold Multimer 在训练过程的数据加载时，需要将 MSA 与 Template 信息转换成 Feature，再进行训练，这样速度较慢。通过修改数据集类 OpenFoldSingleMultimerDataset 的 __getitem__ 方法，可以加速训练过程。

---

1. 准备训练数据

在训练过程中，需要读取 mmcif_cache.json 文件，数据结构如下：

{
    "4ewn": {
        "release_date": "2012-12-05",
        "chain_ids": [
            "D"
        ],
        "seqs": [
            "MLAKRI..."
        ],
        "no_chains": 1,
        "resolution": 1.9
    },
    "5m9r": {
        "release_date": "2017-02-22",
        "chain_ids": [
            "A",
            "B"
        ],
        "seqs": [
            "MQDNS...",
            "MQDNS..."
        ],
        "no_chains": 2,
        "resolution": 1.44
    },
#...
}  

当前的训练数据格式，例如 train_200_mini.csv，如下：

pdb_id,chain_id,resolution,release_date,seq,len,chain_type,filepath
7m5z,"A,B",3.06,2021-10-06,"LEDVV...,QNKLE...","263,264","protein,protein",[pdb_path]/structures/m5/pdb7m5z.ent.gz
7k05,"A,B",1.85,2021-10-06,"MSFPP...,MSFPP...","200,200","protein,protein",[pdb_path]/structures/k0/pdb7k05.ent.gz
# ...

同时需要将 feature 的路径，也加入到训练文件 mmcif_cache.json 中，进而，通过预读文件，进行特征抽取，即：

[your folder]/multimer_train/features

使用特征文件夹中，已经预处理之后的特征 features.pkl，进行训练即可：

# 单个文件夹内容
chain_id_map.json
features.pkl
sequences.fasta

训练文件的转换命令，如下：

python openfold_scripts/main_mmcif_cache_transfer.py -i data/train_200_mini.csv -f [your folder]/multimer_train/features -o mydata/openfold/mmcif_cache_mini.json

源码如下：

#!/usr/bin/env python
# -- coding: utf-8 --
"""
Copyright (c) 2022. All rights reserved.
Created by C. L. Wang on 2023/8/31
"""
import argparse
import json
import os
import sys
from pathlib import Path

import pandas as pd
from tqdm import tqdm

p = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if p not in sys.path:
    sys.path.append(p)


class MmcifCacheTransfer(object):
    """
    训练 CSV 转换成 OpenFold 的 mmcif_cache.json 格式
    """
    def __init__(self):
        pass

    @staticmethod
    def process(input_path, feature_dir, output_path):
        print(f"[Info] 输入文件: {input_path}")
        print(f"[Info] 特征文件夹: {feature_dir}")
        print(f"[Info] 输出文件: {output_path}")
        assert os.path.isfile(input_path)
        df = pd.read_csv(input_path)
        print(f"[Info] 输入样本: {len(df)}")
        mmcif_cache_dict = dict()
        # pdb_id,chain_id,resolution,release_date,seq,len,chain_type,filepath
        for _, row in tqdm(df.iterrows(), "[Info] pdb"):
            pdb_id = row["pdb_id"]
            release_date = row["release_date"]
            chain_ids = row["chain_id"].split(",")
            seqs = row["seq"].split(",")
            no_chains = len(chain_ids)
            resolution = float(row["resolution"])
            feature_folder = os.path.join(feature_dir, pdb_id[1:3], f"pdb{pdb_id}_{''.join(chain_ids)}")
            pdb_dict = {
                "release_date": str(release_date),
                "chain_ids": chain_ids,
                "seqs": seqs,
                "no_chains": no_chains,
                "resolution": resolution,
                "feature_folder": feature_folder
            }
            mmcif_cache_dict[pdb_id] = pdb_dict
        with open(output_path, "w") as fp:
            fp.write(json.dumps(mmcif_cache_dict, indent=4))
        print(f"[Info] 全部处理完成: {output_path}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-i",
        "--input-path",
        help="the input file path.",
        type=Path,
        required=True,
    )
    parser.add_argument(
        "-f",
        "--feature-dir",
        help="the preprocess feature dir.",
        type=Path,
        required=True
    )
    parser.add_argument(
        "-o",
        "--output-path",
        help="the output file path.",
        type=Path,
        required=True
    )

    args = parser.parse_args()

    input_path = str(args.input_path)
    feature_dir = str(args.feature_dir)
    output_path = str(args.output_path)
    assert os.path.isfile(input_path)

    # from root_dir import ROOT_DIR, DATA_DIR
    # input_path = os.path.join(ROOT_DIR, "data", "train_200_mini.csv")
    # output_path = os.path.join(DATA_DIR, "openfold", "mmcif_cache_mini.json")
    mct = MmcifCacheTransfer()
    mct.process(input_path, feature_dir, output_path)


if __name__ == '__main__':
    main()

---

2. 加载训练数据

OpenFold Multimer 的特征读取逻辑，在 openfold/data/data_modules.py#OpenFoldSingleMultimerDataset() 中，即：

if self.mode == 'train' or self.mode == 'eval':
    path = os.path.join(self.data_dir, f"{mmcif_id}")
    ext = None
    for e in self.supported_exts:
        if os.path.exists(path + e):
            ext = e
            break
    if ext is None:
        raise ValueError("Invalid file type")
    # TODO: Add pdb and core exts to data_pipeline for multimer
    path += ext
    if ext == ".cif":
        data = self._parse_mmcif(
            path, mmcif_id, self.alignment_dir, alignment_index)
    else:
        raise ValueError("Extension branch missing")
else:
    path = os.path.join(self.data_dir, f"{mmcif_id}.fasta")
    data = self.data_pipeline.process_fasta(
        fasta_path=path,
        alignment_dir=self.alignment_dir)

修改成直接加载 Feature 的形式，即：

if self.mode == 'train' or self.mode == 'eval':
    # 训练或评估时，使用预处理的特征
    feat_folder = self.mmcif_data_cache[mmcif_id]['feature_folder']
    feat_path = os.path.join(feat_folder, "features.pkl")
    # logger.info(f"[Info] feat_path: {feat_path}")
    data = {}
    with open(feat_path, "rb") as f:
        feat_dict = pickle.load(f)
    data.update(feat_dict)
    # logger.info(f"[Info] data: {data.keys()}")
else:
    path = os.path.join(self.data_dir, f"{mmcif_id}.fasta")
    data = self.data_pipeline.process_fasta(
        fasta_path=path,
        alignment_dir=self.alignment_dir)

同时，还需要修改训练数据总数：

def __len__(self):
    # 数据部分都由 mmcif_data_cache 提供
    # return len(self._chain_ids)
    return len(self.mmcif_data_cache.keys)

---

3. 配置模型训练

模型训练的参数，如下：

python3 train_openfold.py \
    --train_data_dir [your folder]/af2-data-v230/pdb_mmcif/mmcif_files/ \
    --train_alignment_dir mydata/alignment_dir/ \
    --train_mmcif_data_cache_path [your folder]/multimer_train/openfold_cache/mmcif_cache_mini.json \
    --template_mmcif_dir [your folder]/af2-data-v230/pdb_mmcif/mmcif_files/ \
	--output_dir mydata/output_dir/ \
    --max_template_date "2021-10-10" \
    --config_preset "model_1_multimer_v3" \
    --template_release_dates_cache_path mmcif_cache.json \
    --precision bf16 \
    --gpus 1 \
    --replace_sampler_ddp=True \
    --seed 42 \
    --deepspeed_config_path deepspeed_config.json \
    --checkpoint_every_epoch \
    --obsolete_pdbs_file_path [your folder]/af2-data-v230/pdb_mmcif/obsolete.dat

模型训练占用显存较多，V100 目前无法支持，调低 crop_size 与 num_workers，降低资源占用，配置位于 openfold/config.py 中，即：

# crop_size
elif "multimer" in name:
    c.update(multimer_config_update.copy_and_resolve_references())
    c.data.train.crop_size = 64  # TODO: 用于测试

# num_workers
"data_module": {
    "use_small_bfd": False,
    "data_loaders": {
        "batch_size": 1,
        # "num_workers": 16,
        "num_workers": 2,  # TODO: 用于测试
        "pin_memory": True,
    },
},

其中，crop_size = 64 占用显存约是 5141MiB

训练日志，如下：

Epoch 0:   0%|                                 | 0/199 [00:00<?, ?it/s]INFO:openfold/data/data_modules.py:mmcif_id is: 7poc, idx: 148 and has 4 chains
INFO:openfold/data/data_modules.py:mmcif_id is: 7u49, idx: 97 and has 3 chains
INFO:openfold/data/data_modules.py:mmcif_id is: 7z7h, idx: 114 and has 6 chains
INFO:openfold/data/data_modules.py:mmcif_id is: 7nup, idx: 111 and has 4 chains
cum_loss: tensor([84.1698], device='cuda:0', dtype=torch.float64, grad_fn=<MulBackward0>) losses: {'distogram': tensor(4.1562, device='cuda:0', dtype=torch.float64), 'experimentally_resolved': tensor(0.6914, device='cuda:0'), 'fape': tensor(1.6598, device='cuda:0', dtype=torch.float64), 'plddt_loss': tensor(3.9062, device='cuda:0', dtype=torch.float64), 'masked_msa': tensor(3.0938, device='cuda:0'), 'supervised_chi': tensor(0.7941, device='cuda:0', dtype=torch.float64), 'violation': tensor(3.6495, device='cuda:0'), 'tm': tensor(4.1562, device='cuda:0', dtype=torch.float64), 'chain_center_of_mass': tensor([1.3754], device='cuda:0', dtype=torch.float64), 'unscaled_loss': tensor([10.5212], device='cuda:0', dtype=torch.float64), 'loss': tensor([84.1698], device='cuda:0', dtype=torch.float64)}
Epoch 0:   1%|▉                                 | 1/199 [02:55<9:38:06, 175.18s/it, loss=84.2, v_num=]

---