简介：面向药物设计与分子生成方向的研究者和工程实践者，这份资料围绕 ICML 2024 提出的 IRDiff（基于蛋白质-配体相互作用的检索增强 3D 分子扩散模型），提供了一份完整的复现与评测记录。压缩包内为 1 个 PDF 文档，约 12.23MB，把训练好的模型、修正后的项目代码、报错位置与修改方法、缺失模块文件以及测试案例整理在一起，可省去读者复现时反复排查和试错。修正后的代码能针对特定蛋白或口袋体系直接运行 IRDiff 生成分子，并计算 vina_score、vina_docking_score、qvina_score、QED、SA 等指标；也能基于自定义数据集做微调与再训练。文档另附个人分析标注，梳理 PMINet 预训练、检索增强与自增强两条引导机制的实现思路，便于在理解框架后做二次修改。目前已有 121 人学习，适合具备一定深度学习基础、希望把扩散模型用于结构导向药物设计任务的读者。

1. 跑通采样两小时，跑通评测链路三天

手头这份 IRDiff 完整评测文档加可运行项目代码，核心是一个基于蛋白质-配体相互作用的检索增强 3D 分子扩散模型，用在基于结构的药物设计里做目标感知的分子生成。我按文档在一台 24G 显存的机器上重走了一遍全流程：环境装了三次才过，采样一百个分子花了五十分钟，最后卡在"分子生成了但构象没存下来"这种说大不大、说小要命的细节上。

这篇不铺论文综述。只讲三件事：IRDiff 相对 TargetDiff、DecompDiff、Pocket2Mol 这批基线，优势到底落在哪个指标上；环境里三处硬冲突的具体改法；以及从 sample_split.py 到 vina_score、QED、SA 这条链路，哪些参数能随手调、哪些一调结果就没法比。正在做 SBDD 方向选型，或者代码拿到手跑不出指标的人，可以照着走一遍。

2. 和四类基线摆在一起看：检索增强到底补了哪个洞

选型这件事，最怕的是拿论文里的 SOTA 数字直接对号入座。分子生成这个方向，不同模型的指标口径、评测脚本、后处理方式差别很大，光看一个 Vina 均值很容易得出错误结论。我把 IRDiff 放进它自己论文里那套对比框架，逐项拆开看它赢在哪、为什么赢、以及这些优势在我的实际场景里值不值钱。

2.1 从头生成的瓶颈，和检索增强切进去的位置

先说要解决的问题。给定一个蛋白质口袋，生成能结合上去的 3D 配体分子，主流做法分两派。自回归派（AR、Pocket2Mol）一个原子一个键地往口袋上加，问题是误差会累积，而且"先加哪个原子"这个生成顺序本身就没有最优解，模型只能学一个近似顺序。扩散派（TargetDiff、DecompDiff）走的是另一条路：从高斯先验出发，逐步去噪，最后用后处理分配化学键，靠 SE(3)-等变网络捕捉原子间的局部和全局空间关系，整体表现比自回归稳。

但扩散派也有自己的坑，而且是结构性的。分子内部和分子之间的姿态搜索空间太大，纯从头生成等于让模型在一个几乎没有约束的空间里瞎走，优化难度高，最后生成的分子在结合亲和力这类生物学指标上就是差一口气。这个判断很关键：它意味着"提升生成质量"和"换更强的骨干网络"不是一回事，前者需要往生成过程里塞外部信息。

IRDiff 切的就是这个位置。它不去改扩散骨干，而是在扩散过程里引入一小批"已知对目标口袋有高结合亲和力的参考配体"，用这些参考分子的相互作用结构去引导生成。思路借鉴的是检索增强生成那套东西——用任务相关的参考样本，把模型泛化到没见过的目标上。落到药物设计场景，就是给每个新靶点配一组"参考答案"，让模型照着结构线索改，而不是从零构思。

2.2 PMINet 怎么把结合亲和力变成可用的先验

整套机制的地基是一个叫 PMINet 的蛋白质-分子相互作用网络。结构上不复杂：两层浅层 SE(3)-等变网络分别作用在蛋白质和配体的全连接图上，先建模分子内相互作用；再用原子级交叉注意力学分子间相互作用。它的训练目标是结合亲和力预测，用的是 PDBbind v2016，3767 个复合物做训练、290 个做测试。

这个预训练模型在 IRDiff 里被用在两个地方，这也是我认为整套设计里最值得借鉴的部分。

第一处是检索。给定一个目标蛋白和一份外部配体数据库，用 PMINet 扫一遍库，按预测结合亲和力排序，取 top-k 作为参考分子。论文里的做法是从训练集随机抽 128 个配体构成检索池，每个蛋白选预测排名第一的那个作为参考。注意这里有个隐含前提：参考分子必须真的"高亲和力"，后面消融会验证这一点。

第二处是自增强。采样过程接近结束时，当前生成的分子其实已经是候选里结合亲和力最高的一批了，这时候把它连同目标蛋白一起喂回 PMINet，取出蛋白质和分子的嵌入，再融合回扩散模型的特征里。注意这个操作是"非端到端"的，PMINet 在整个生成阶段是冻结的，只提供特征，不参与梯度。

两种增强的分工很清楚：检索增强负责"给参考"，自增强负责"自我修正"。论文里有个容易被忽略的细节——如果只上增强机制但不挂预训练的 PMINet，性能反而下降，因为这时候拿到的特征里没有真正的相互作用知识，纯属噪声。这个结论对复现很重要：别想着跳过 PMINet 预训练权重省事。

2.3 四个维度逐项对比：JSD、Vina、High Affinity、QED/SA

评测口径按论文分三类：分子结构、靶标结合亲和力、分子属性。分子结构用生成分子与真实分子的原子/键距离经验分布之间的 Jensen-Shannon 散度衡量；结合亲和力用 AutoDock Vina 算，分四个口径——Vina Score 直接估、Vina Min 做局部结构最小化后估、Vina Dock 加一次重对接、High Affinity 统计生成分子比真实分子结合更好的比例；分子属性看 QED、SA 和多样性。

横向摆开是这样：

具体数字：IRDiff 的全原子成对距离分布 JSD 是 0.08，比 Pocket2Mol 和 TargetDiff 都低，说明它抓到的原子距离分布更接近真实；键类型分布上跟 DecompDiff 打平，明显好于其他基线。结合指标上，平均 Vina 得分 -6.03，相对 Pocket2Mol 在均值和中位数上分别提升 17.3% 和 46.6%，相对 DecompDiff 是 6.3% 和 14.1%。High Affinity 这一项，平均 67.4% 的生成分子在测试集上结合得比真实分子还好。QED 平均 0.53。

这里有个选型上必须说清的权衡。TargetDiff 和 DecompDiff 在结合指标上压过 AR 和 Pocket2Mol，但 QED 和 SA 落后；IRDiff 是唯一一个既拿到 SOTA 级结合分数、又把 QED 稳在合理区间的。论文自己的解释是：QED 和 SA 在实际药物发现流程里主要当粗筛用，落在合理范围就行，不必死磕。这个判断我认同，但它有个前提——你的下游流程确实有实验或更精细的打分环节兜底。如果你的场景是"生成一批分子直接进湿实验"，那 QED 和 SA 的权重就得往上提，这时候 IRDiff 的优势就没那么绝对了。

2.4 消融实验里最值得抄的两个结论

复现别人工作，消融表往往比主表更有用，因为它告诉你哪些旋钮真的有用。

第一个结论关于检索池。变量 n 是外部数据库大小，k 是参考样本数量。实验结果是：n 越大，结合相关指标越好，因为候选多样性上来了，能提供更多有用的子结构线索；但 k 简单加大并没有明显收益，反而因为引入更多噪声干扰生成。这个非对称性很有意思——大家直觉上会觉得"多给几个参考总没坏处"，实际不是。我一般会这么理解：参考分子的价值在于"精确命中高亲和力构象"，给多了等于稀释信号。

第二个结论关于参考质量。把参考配体从亲和力最高的换成 PMINet 排名最低的，平均预测结合亲和力从 7.24 掉到 3.95，结合相关指标全面变差。这条直接说明检索环节不能随便找个近似库糊弄——你的外部库里如果没几个真正高亲和力的分子，检索增强会变成负资产。

还有一条定性观察值得记：生成分子与其对应参考分子的指纹相似度，高于它与数据库里随机分子的相似度，而且两者共享的关键子结构在位置上会被模型重新组合。这说明模型确实在"抄参考的解题思路"，而不是把参考分子整个搬过来。对做骨架跃迁的人，这是个正面信号。

3. 环境从零到能跑：三处依赖冲突的具体修法

环境这一段是整份资源里最实在的部分。论文给个 yml 就完事，但真正能跑起来的环境和 yml 之间差着好几次报错。我把踩过的三处冲突按发生顺序写清楚，每一处都带上报错原文特征，方便你对号入座。

3.1 第一处冲突：cryptography 与 openssl 的版本死锁

按文档直接conda env create -f irdiff.yml，第一个报错是这样的：

LibMambaUnsatisfiableError: Encountered problems while solving: - package cryptography-37.0.1-py38h9ce1e76_0 requires openssl <1.1.2a The following packages are incompatible ├─ cryptography ==37.0.1 py38h9ce1e76_0 is installable and it requires │ └─ openssl <1.1.2a , which can be installed; └─ openssl ==3.1.0 hd590300_2 is not installable because it conflicts with

原因很直白：yml 里cryptography锁死在 37.0.1，而这个版本要求openssl < 1.1.2a；同一份 yml 又把openssl锁在 3.1.0。两个约束不可能同时满足，求解器只能报冲突。

解法是解掉 cryptography 的版本钉死，让 conda 自己去选一个和 openssl 3.1.0 兼容的版本：

# 修改前（irdiff.yml 第 42 行附近） - cryptography=37.0.1=py38h9ce1e76_0 # 修改后：去掉版本号和 build 号 - cryptography

改完重新conda env create -f irdiff.yml。这类冲突在 bioinformatics 项目里极常见，判断方法也简单：看报错里"requires ... which conflicts with"两边是谁，优先解掉上层包（这里是 cryptography）的版本约束，因为底层库（openssl）往往被别的包也依赖，动它牵连面更大。

顺带说一句环境名。yml 里写的 name 是irdiff，我习惯统一改成项目本名的大小写形式，方便conda activate的时候一眼认出来，后面所有命令也都按这个环境名走。

3.2 第二处冲突：autodocktools-py3 在 pip 里找不到版本

继续装，conda 本体过了，卡在 pip 阶段：

Pip subprocess error: ERROR: Could not find a version that satisfies the requirement autodocktools-py3==0+unknown ERROR: No matching distribution found for autodocktools-py3==0+unknown failed CondaEnvException: Pip failed

autodocktools-py3这个包在 PyPI 上没有可以直接匹配的发行版，yml 里写的0+unknown是个无法解析的版本串。这不是网络问题，换镜像源也没用。

绕行办法是进虚拟环境，从维护者的源码仓库直接装：

conda activate IRDiff # 该工具只有源码分发，PyPI 上无可用 wheel，走 git 安装 python -m pip install "git+https://github.com/<维护者>/AutoDockTools_py3"

装完会看到构建 wheel 的过程，最后输出Successfully installed AutoDockTools-py3-1.5.7.post1+...。版本号带+后缀是正常的，说明是从 git 提交构建的，不必纠结。

剩下几个依赖一起补。这几个是打分和结构处理链路的必需项，缺一个后面算 vina_score 就断：

# meeko 负责配体准备，vina 是打分核心，pdb2pqr/propka 处理质子化状态 pip install meeko==0.1.dev3 pdb2pqr==3.6.1 vina==1.2.2 propka==3.5.0 mmcif-pdbx==2.0.1

注意meeko==0.1.dev3这个 dev 版本别随手升级。它和vina==1.2.2之间的接口在正式版里改过，升上去大概率在配体转换那一步抛异常。依赖链上的版本，能不动的就别动。

3.3 第三处：pyg 版本必须钉在 2.0.4

这一条文档里只写了一句"注意 pyg 版本，使用 2.0.4"，但它值得单独拎出来讲，因为它不报错在安装阶段，而是报错在训练阶段。

PyTorch Geometric 的MessagePassing接口在 2.1 之后改过签名，SE(3)-等变层的实现如果按旧接口写，升级之后会在消息传递的函数调用处抛参数不匹配的异常，报错位置离真正原因很远，第一次遇到很容易往数据格式上找。IRDiff 的等变层和交叉注意力实现是按 2.0.4 写的，所以这个版本号不是建议，是硬约束。

验证方法：

python -c "import torch_geometric; print(torch_geometric.__version__)" # 期望输出 2.0.4

如果装出来是别的版本，直接pip install torch_geometric==2.0.4降下来，别指望"新版本应该兼容"。这类几何深度学习库的向后兼容做得不算好，血泪经验就是钉死。

另外还有一个小的，einops。它在models/attn.py里被 import，但 yml 里没列。这个不急着现在装，等到跑采样的时候会撞上，第 5 章会细说。

4. 数据落盘与路径改造：把 sample_split.py 跑出第一批分子

环境通了只是能 import，真正跑起来还有两道关：数据放对位置、路径改成相对路径。原作者是在自己的目录结构下开发的，所有配置里写的是绝对路径，直接搬过来必然找不到文件。

4.1 两个数据文件放哪，各自干什么

IRDiff 的数据预处理流程和 TargetDiff 一致，需要两个文件：

• crossdocked_pocket10_pose_split.pt：数据划分文件，记录哪些复合物属于训练集、哪些属于测试集，采样时靠它定位每一个测试蛋白。
• crossdocked_v1.1_rmsd1.0_pocket10_processed_final.lmdb：真正的结构数据库，存的是预处理后的蛋白口袋和配体坐标。

两个都放到项目根目录下的./data/里。这里提一句数据本身的筛选标准，因为它决定了你的复现结果能不能和别人对上：原始 2250 万个对接复合物，先按对接姿势与真实值 RMSD < 1Å 筛出高质量构象，再按蛋白质序列相似性 < 30% 去冗余，最后得到 10 万对训练数据、100 个测试蛋白。如果你的自定义数据集不按这个标准筛，指标没有可比性。

4.2 三处路径改写，一个都不能漏

路径要改的地方有三处，漏一处就白跑。

第一处是评测脚本。cal_metrics_from_pt.py第 8 行：

# 修改前 eval_path = '/path/to/eval_results/' # 修改后：指向项目内的评测结果目录 eval_path = './eval_results/'

第二处是训练配置。./configs/training.yml第 4–5 行：

# 修改前 path: /path/to/crossdocked_v1.1_rmsd1.0 split: /path/to/crossdocked_pocket10_pose_split.pt # 修改后：和 4.1 里放置的位置对齐 path: ./data/crossdocked_v1.1_rmsd1.0 split: ./data/crossdocked_pocket10_pose_split.pt

第三处是采样配置./configs/sampling.yml。这个文件不改路径，但要确认 checkpoint 指向真实存在的权重文件。完整的采样参数如下，我按"能不能动"标注了一遍：

center_pos_mode: protein这一项我要多说一句。IRDiff 在训练和采样时都把蛋白质原子的质心平移到零点，配合 EGNN 的设计保证逆变换的 SE(3) 等变性。这个对齐方式和训练时必须完全一致，否则你生成的分子坐标整体偏移，后面 Vina 打分算出来的结合能纯属噪声——这是那种"能跑完、有输出、但结果是错的"的黑匣子型 bug。

4.3 采样命令与实测显存、耗时

内置案例是测试集里 index 为 0 的体系（PDB 编号 2Z3H，目录名BSD_ASPTE_1_130_0）。用这个体系验证全链路最省事，因为它的口袋结构和原始配体都有对照。

采样命令：

python sample_split.py \ --start_index 0 \ --end_index 0 \ --batch_size 25 \ --test_prompt_indices_path ./src/test_prompt_ligand_indices_top3.pt \ --result_path ./results/result_0

参数含义逐个说清：--start_index和--end_index指定测试蛋白的索引区间，都设 0 表示只跑第一个；--batch_size 25控制单次前向的样本数，直接决定显存占用；--test_prompt_indices_path指向检索到的参考配体索引文件，这是检索增强的输入，格式是预先算好的 top-k 索引；--result_path是输出目录。

实测数据：batch_size=25时显存占用约 19 GB，生成 100 个分子耗时约 50 分钟。这个耗时要提前算进实验节奏里，调一次参数就是五十分钟，所以别拿单次结果下结论。

跑起来之后正常输出长这样，几个关键行值得盯：

INFO Successfully load the dataset INFO Restored from ./pretrained_models/PMINet with 0 missing and 0 unexpected INFO Successfully load the model! sampling: 100%|...| INFO Sample done!

0 missing and 0 unexpected这行是权重加载的体检报告，出现非零数字说明 PMINet 权重和当前模型结构对不上，后面的自增强全是错的。这一行我每次跑都会看一眼，比看 loss 曲线实在。

5. 五个真实翻车点：现象、原因和最后的解法

这一章按现象、原因、解决三段式记录。这些都是我实际撞过的，不是从文档里抄的注意事项。

5.1 conda 建环境直接崩在求解阶段

现象：conda env create -f irdiff.yml输出LibMambaUnsatisfiableError，提示cryptography要求openssl <1.1.2a，而 yml 里锁了openssl ==3.1.0。

原因：yml 由作者在某个时间点的环境导出，两个包的版本约束在当时成立，随着 conda 源里 openssl 更新，这个组合变得不可满足。本质是环境快照过期。

解决：把cryptography=37.0.1=py38h9ce1e76_0改成不带版本号的cryptography，交给求解器选兼容版本。改完重跑，求解能过。如果还不行，用conda env create -f irdiff.yml --no-deps先建骨架再手动补，但一般不需要走到这一步。

5.2 pip 阶段报 autodocktools 找不到发行版

现象：conda 主体装完，pip 子进程失败，报Could not find a version that satisfies the requirement autodocktools-py3==0+unknown。

原因：这个包没有发布到 PyPI 的可解析版本号，yml 里的0+unknown是个占位串，pip 无法匹配。

解决：激活环境后从源码仓库装python -m pip install "git+https://github.com/<维护者>/AutoDockTools_py3"，装完确认AutoDockTools-py3出现在pip list里。装完之后再补meeko、vina、pdb2pqr、propka、mmcif-pdbx这几个，顺序无所谓，但meeko的 dev 版本别升。

5.3 采样脚本一 import 就挂，缺 einops

现象：python sample_split.py还没进主逻辑就抛ModuleNotFoundError: No module named 'einops'，报错栈指向models/attn.py第 4 行。

原因：yml 漏列了这个依赖。einops在注意力模块里做张量重排，属于隐式依赖——原作者的环境里早就装了，导出 yml 时没体现出来。

解决：conda install einops或pip install einops都行。装完直接重跑采样命令，不需要重新建环境。这类隐式依赖在复现别人项目时特别常见，判断方法就是看报错的 import 位置，缺什么补什么，别急着怀疑代码。

5.4 pyg 版本不对，训练阶段才炸

现象：环境装完、采样也能跑，一切到训练就抛参数不匹配的异常，报错位置在消息传递的调用处，看起来像是数据维度问题。

原因：torch_geometric装成了 2.1 以上版本，MessagePassing的接口签名变了，而项目的等变层是按 2.0.4 写的。报错信息和真实原因隔了好几层，第一次遇到会往数据预处理方向排查，浪费很久。

解决：pip install torch_geometric==2.0.4强制降版，然后python -c "import torch_geometric; print(torch_geometric.__version__)"确认。降完如果还有别的几何库报接口错，检查torch-scatter、torch-sparse这些配套包是否也跟着降了，它们必须和目标 pyg 版本对齐。

5.5 采样跑完，拿不到分子构象

现象：采样命令显示Sample done!，./results/result_0目录下有文件，但打开一看全是张量格式的中间结果，没有可以直接喂给对接工具的 sdf。评测脚本eval_split.py跑完也只存了 metrics，构象还是没落盘。

原因：项目的评测脚本设计目标是算指标，不是导出分子。if args.save:分支里torch.save存的是stability、bond_length、all_results三个字典，分子坐标和原子类型在all_results里但没转成化学格式。

解决：在eval_split.py末尾的保存分支里追加一段，把结果里的原子坐标和原子类型用 RDKit 重建为 Mol 对象，再用Chem.SDWriter写文件。目标产物是./eval_results/Generated_molecules.sdf。具体代码和字段对应关系放在下一章，因为这一段是整条链路的最后一环，值得单独讲清楚。这个坑的代价很直接：不改这一段，前面五十分钟采样等于只拿到一堆没法做下游分析的张量。

6. 评测链路补全：补上构象保存，顺带说清 n 和 k 怎么调

先说补构象这件事。eval_split.py结尾那个if args.save:分支，原本只存指标，我在这后面接了一段导出逻辑：

from rdkit import Chem if args.save: torch.save({ 'stability': validity_dict, 'bond_length': all_bond_dist, 'all_results': results }, os.path.join(result_path, f'metrics_{args.eval_step}_{eval_steps}.pt')) # 把 results 里的原子坐标与原子类型重建成 RDKit 分子并落盘 writer = Chem.SDWriter(os.path.join(result_path, 'Generated_molecules.sdf')) for item in results['mols']: mol = Chem.RWMol() for atom_type in item['atom_types']: mol.AddAtom(Chem.Atom(int(atom_type))) conf = Chem.Conformer(mol.GetNumAtoms()) for idx, pos in enumerate(item['positions']): conf.SetAtomPosition(idx, pos.tolist()) mol.AddConformer(conf) writer.write(mol.GetMol()) writer.close()

字段名mols、atom_types、positions要按你本地results的实际结构对，不同版本可能有出入，用print(results.keys())先看一眼。写出来的 sdf 直接丢给 AutoDock Vina 就能算 vina_score、vina_docking_score 和 qvina_score，QED 和 SA 用 RDKit 的Descriptors.qed和sascorer拿，四个指标凑齐才算完整评测。

再说调参。回头看我自己的实验记录，最值钱的一条是检索池的两个旋钮不是对称的：n（外部库大小）加大，结合指标稳定变好；k（参考样本数）加大，基本没收益甚至变差，噪声进来了。所以我的默认配置是 n 拉到 128，k 保持 1，用 PMINet 排序第一的那个作参考。另一条是参考质量的红线，高亲和力参考（平均 7.24）和低亲和力参考（平均 3.95）跑出来的结合指标差距非常明显，意味着换外部分子库之前必须先验一遍库里的亲和力分布，别拿一个平庸的库硬上检索增强。

复现验证我固定走三步：采样日志里确认0 missing and 0 unexpected；跑完检查 sdf 里分子数是否等于num_samples；最后抽 5 个分子手动算一遍 vina_score 和脚本输出对一下。这三步走完，结果才敢往表里填。

从那以后我每次跑这类扩散采样，都先把 sdf 保存那段补上再开跑，不然五十分钟白等一次。希望帮到你。

本文还有配套的精品资源，点击获取