加速余弦相似度计算：60万分子与单蛋白相似性计算优化求助

问题背景 #

需计算60万个分子与单个蛋白的余弦相似度：

• 分子通过rdkit.Chem.MolFromSmiles(smiles)生成
• 蛋白从文件读取
  当前用4核CPU通过pool.map()并行计算，速度极慢——仅处理1万个分子就耗时超2小时；尝试GPU加速但无GPU利用率，现寻求可行的提速方案。

原代码如下：

from sklearn.metrics.pairwise import cosine_similarity
import rdkit
from rdkit import Chem

# Function to compute chemical similarity
def compute_similarity(mol1, mol2):
    if mol1 is not None and mol2 is not None:
        fp1 = Chem.RDKFingerprint(mol1)
        fp2 = Chem.RDKFingerprint(mol2)
        similarity = cosine_similarity([fp1], [fp2])[0][0]
        return similarity
    else:
        return None

def compute_similarity_wrapper(args):
    return compute_similarity(*args)

# The calculation for similarity takes too long. Let's parallelize it
import multiprocessing
from multiprocessing import Pool

molecule_mols_subset = molecule_mols[:10_000]

# Get the number of available processors
num_processors = multiprocessing.cpu_count()
print(f" {num_processors} CPUs ---- available for parallel processing")

# Use multiprocessing.Pool to parallelize the computation
# Use the "map" function to preserve the order after parallel computing.
with multiprocessing.Pool(processes=num_processors) as pool:
    sims = pool.map(compute_similarity_wrapper, [(mol, alb_mol) for mol in molecule_mols_subset])

sims[:10]

优化方案 #

1. 预计算蛋白指纹，消除重复计算 #

蛋白是固定对象，无需在每个分子计算时重复生成其RDK指纹，提前生成一次即可复用，直接砍掉大量冗余计算。

2. 替换相似度计算函数为RDKit原生实现 #

sklearn的cosine_similarity是通用计算函数，针对分子二进制指纹的效率远低于RDKit内置的DataStructs.TanimotoSimilarity。对于二进制指纹，余弦相似度与Tanimoto相似度结果完全等价，直接替换可大幅提升计算速度。

3. 优化多进程数据传递，减少开销 #

原代码通过pool.map()重复传递alb_mol对象，进程间数据拷贝会消耗大量资源。用functools.partial将蛋白指纹绑定到计算函数，避免重复传递冗余数据。

4. 批量处理降低函数调用开销 #

将分子列表拆分为批量，批量生成指纹并计算相似度，减少单分子处理的函数调用次数，进一步提升效率。

优化后代码示例 #

基础优化版（单分子处理） #

import multiprocessing
from functools import partial
from rdkit import Chem
from rdkit.DataStructs import TanimotoSimilarity

# 预计算蛋白的RDK指纹
protein_fp = Chem.RDKFingerprint(alb_mol) if alb_mol is not None else None

def compute_similarity(mol, protein_fp):
    if mol is None or protein_fp is None:
        return None
    mol_fp = Chem.RDKFingerprint(mol)
    return TanimotoSimilarity(mol_fp, protein_fp)

molecule_mols_subset = molecule_mols[:10_000]
num_processors = multiprocessing.cpu_count()

# 使用partial绑定蛋白指纹，避免重复传递
with multiprocessing.Pool(processes=num_processors) as pool:
    sims = pool.map(partial(compute_similarity, protein_fp=protein_fp), molecule_mols_subset)

print(sims[:10])

进阶批量处理版 #

import multiprocessing
from functools import partial
from rdkit import Chem
from rdkit.DataStructs import BulkTanimotoSimilarity

# 预计算蛋白指纹
protein_fp = Chem.RDKFingerprint(alb_mol) if alb_mol is not None else None

def process_batch(mol_batch, protein_fp):
    if protein_fp is None:
        return [None]*len(mol_batch)
    # 批量生成分子指纹
    fps = [Chem.RDKFingerprint(mol) if mol is not None else None for mol in mol_batch]
    # 批量计算相似度
    sims = []
    for fp in fps:
        if fp is None:
            sims.append(None)
        else:
            sims.append(BulkTanimotoSimilarity(protein_fp, [fp])[0])
    return sims

# 拆分批量，每1000个分子为一批
batch_size = 1000
molecule_batches = [molecule_mols[i:i+batch_size] for i in range(0, len(molecule_mols_subset), batch_size)]

num_processors = multiprocessing.cpu_count()

with multiprocessing.Pool(processes=num_processors) as pool:
    results = pool.map(partial(process_batch, protein_fp=protein_fp), molecule_batches)

# 合并结果
sims = []
for batch_res in results:
    sims.extend(batch_res)

print(sims[:10])

可选：使用Dask处理超大规模数据 #

若60万分子全部加载内存有压力，Dask可实现分块并行计算，支持内存外处理：

import dask.bag as db
from rdkit import Chem
from rdkit.DataStructs import TanimotoSimilarity

protein_fp = Chem.RDKFingerprint(alb_mol) if alb_mol is not None else None

def compute_sim(mol):
    if mol is None or protein_fp is None:
        return None
    mol_fp = Chem.RDKFingerprint(mol)
    return TanimotoSimilarity(mol_fp, protein_fp)

# 将分子列表转为Dask Bag，自动分块并行
mol_bag = db.from_sequence(molecule_mols, npartitions=multiprocessing.cpu_count()*2)
sims = mol_bag.map(compute_sim).compute()

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