Linux下运行分子描述符计算代码触发OOM的解决方法

原代码存在多处无意义的内存占用和不合理的全量加载逻辑，是触发OOM的直接原因：

• 代码中定义的列表l没有任何实际业务用途，循环处理过程中持续往里面追加所有蛋白、配体的RDKit Mol对象，6.2万组分子对象累计占用内存可达数十GB，是首要内存泄漏点
• 特征计算环节一次性将所有配体PDB文件加载为Mol对象，同时一次性生成全量6.2万*2048维的特征矩阵，内存峰值直接超过设备可用阈值
• 文件写入操作未使用上下文管理器，文件句柄无法正常释放；循环中生成的PDBFixer、Mol对象未及时回收，处理过程中内存持续线性上涨
• 结果存储环节同时在内存中持有Numpy特征数组、DeepChem数据集对象、Pandas DataFrame多份数据副本，进一步抬高内存占用

修改点包括移除无用缓存、分批计算特征、主动回收内存、自动释放文件句柄、支持断点续跑：

import os
import gc
import numpy as np
import pandas as pd
from rdkit import Chem
import deepchem as dc
from simtk.openmm.app import PDBFile
from pdbfixer import PDBFixer
from deepchem.utils.vina_utils import prepare_inputs

# 批次大小可根据自身内存调整，16G内存建议设为500-1000
BATCH_SIZE = 1000
DATA_PATH = 'data.csv'
OUTPUT_PATH = 'descr.csv'

if __name__ == '__main__':
    data = pd.read_csv(DATA_PATH)
    pdbids = data['pdb_id'].values
    ligand_smiles = data['smiles'].values

    # 逐样本处理PDB与配体结构，不缓存无用分子对象
    for pdbid, ligand in zip(pdbids, ligand_smiles):
        # 已处理过的样本直接跳过，支持断点续跑
        if os.path.exists(f'ligand_{pdbid}.pdb'):
            continue
        try:
            fixer = PDBFixer(url=f'https://files.rcsb.org/download/{pdbid}.pdb')
            # 上下文管理器自动关闭文件句柄，避免资源泄漏
            with open(f'{pdbid}.pdb', 'w') as f:
                PDBFile.writeFile(fixer.topology, fixer.positions, f)
            
            p, m = None, None
            try:
                p, m = prepare_inputs(
                    f'{pdbid}.pdb', 
                    ligand, 
                    replace_nonstandard_residues=False,
                    remove_heterogens=False, 
                    remove_water=False,
                    add_hydrogens=False
                )
            except:
                print(f'{pdbid} failed sanitization')
                if os.path.exists(f'{pdbid}.pdb'):
                    os.remove(f'{pdbid}.pdb')
                # 主动释放当前循环大对象
                del fixer
                gc.collect()
                continue

            if p and m:
                Chem.rdmolfiles.MolToPDBFile(p, f'{pdbid}.pdb')
                Chem.rdmolfiles.MolToPDBFile(m, f'ligand_{pdbid}.pdb')
            
            # 主动回收当前循环占用的内存
            del fixer, p, m
            gc.collect()

        except:
            print(f'{pdbid} failed PDB download')
            continue

    # 扫描所有有效生成的文件
    proteins = [f for f in os.listdir('.') if len(f) == 8 and f.endswith('.pdb')]
    ligands = [f for f in os.listdir('.') if f.startswith('ligand') and f.endswith('.pdb')]

    # 过滤处理失败的无效样本
    failures = set([f[:-4] for f in proteins]) - set([f[7:-4] for f in ligands])
    for pdbid in failures:
        proteins.remove(f'{pdbid}.pdb')
    valid_pdbids = [f[:-4] for f in proteins]
    valid_ligand_files = [f'ligand_{pid}.pdb' for pid in valid_pdbids]
    # 建立pdbid到标签的映射，避免全量dataframe长期驻留内存
    label_map = data.set_index('pdb_id')['labels'].to_dict()

    # 初始化指纹生成器
    fp_featurizer = dc.feat.CircularFingerprint(size=2048)
    # 预写入输出文件表头
    init_df = pd.DataFrame(columns=['pdb_id', 'label'] + [f'fp_{i}' for i in range(2048)])
    init_df.to_csv(OUTPUT_PATH, index=False)
    del init_df
    gc.collect()

    # 分批计算分子指纹，算完一批写入一批，不缓存全量特征
    for i in range(0, len(valid_ligand_files), BATCH_SIZE):
        batch_files = valid_ligand_files[i:i+BATCH_SIZE]
        batch_pdbids = valid_pdbids[i:i+BATCH_SIZE]
        # 仅加载当前批次的分子对象
        batch_mols = []
        valid_batch_pids = []
        for f, pid in zip(batch_files, batch_pdbids):
            mol = Chem.MolFromPDBFile(f)
            if mol is not None:
                batch_mols.append(mol)
                valid_batch_pids.append(pid)
        # 提取当前批次标签
        batch_labels = [label_map[pid] for pid in valid_batch_pids]
        # 计算当前批次指纹
        batch_features = fp_featurizer.featurize(batch_mols)
        # 拼接结果并追加写入CSV
        batch_df = pd.DataFrame(batch_features, columns=[f'fp_{i}' for i in range(2048)])
        batch_df.insert(0, 'label', batch_labels)
        batch_df.insert(0, 'pdb_id', valid_batch_pids)
        batch_df.to_csv(OUTPUT_PATH, mode='a', header=False, index=False)
        # 释放当前批次占用的所有内存
        del batch_files, batch_pdbids, batch_mols, valid_batch_pids, batch_labels, batch_features, batch_df
        gc.collect()
        print(f'Processed {min(i+BATCH_SIZE, len(valid_ligand_files))}/{len(valid_ligand_files)} samples')

• 若运行时仍有内存压力，可将BATCH_SIZE调小至200-500，内存占用会进一步降低，仅计算速度稍有下降
• 若不需要留存临时生成的PDB文件，可在每批特征计算完成后删除对应批次的蛋白、配体PDB文件，节省磁盘空间
• 若对分子结构标准度要求较高，可恢复prepare_inputs的默认参数（开启去水、替换非标准残基、补氢操作），提前过滤无效分子，减少后续计算量
• 运行时可在系统层面监控内存占用，若单样本处理阶段内存仍持续上涨，可适当增加gc.collect()的调用频率

内容的提问来源于stack exchange，提问作者LTE