企业官网建设流程全解析

从 0 到 1 上手 E(3) 等变扩散模型：生成三维分子的完整实操指南

【免费下载链接】e3_diffusion_for_molecules项目地址: https://gitcode.com/gh_mirrors/e3/e3_diffusion_for_molecules

如果你正在研究分子生成，很可能已经踩过这样的坑：用二维图生成分子，忽略了原子的真实空间排布；或者硬生生把三维坐标塞进普通神经网络，结果模型一遇到旋转、平移就"失忆"。E(3) 等变扩散模型（EDM，Equivariant Diffusion Model）正是为解决这类问题而生的三维分子生成方案。本文不堆理论，直接带你从环境搭建开始，一步步在 QM9 数据集上训练、评估、可视化你自己的 EDM 模型，并讲解条件生成、参数调优和常见报错的处理方法。

先聊聊痛点：为什么"画"一个分子这么难

分子不是平面图上的点线连接，而是一组带坐标的原子的三维排布。同一个分子，稍微旋转一下、平移一段距离，化学性质完全不变——但如果你用普通网络直接学坐标，模型会把"旋转后的分子"当成全新样本，数据利用率极低。

更麻烦的是，分子生成任务要同时处理两类信息：

• 连续信息：每个原子的三维坐标；
• 离散信息：原子的种类（C、N、O、F……）、电荷数等。

把这两类数据放在同一个生成框架里，还要保证几何变换下的行为一致，这就是传统方法迟迟做不好的原因。EDM 的设计目标，恰恰就是把"等变性"和"扩散生成"揉在一起。

EDM 的核心原理：一句话版本

EDM 用扩散模型逐步给分子的坐标和原子类型加噪、再去噪，而中间的"去噪网络"采用 E(3) 等变网络（EGNN），保证无论分子怎么旋转，模型学到的分布都保持不变。

用公式说就是：生成分布满足 p(x, h) = p(Rx, h)，其中 R 是任意旋转操作，x 是坐标，h 是原子特征。旋转前后，分布一致。这一点正是整个模型的灵魂。

有了这个特性，训练时不需要对每个分子做数据增强（旋转复制），模型天然理解"旋转不改本质"，生成的样本也更有化学合理性。

动手前先认清仓库结构

克隆代码后，先花两分钟把目录结构过一遍，后面找文件会快很多：

• main_qm9.py：QM9 数据集的主训练入口；
• main_geom_drugs.py：GEOM-Drugs 数据集训练入口；
• eval_analyze.py：分析生成分子的质量（稳定性、有效性等指标）；
• eval_sample.py：采样并可视化生成的分子；
• eval_conditional_qm9.py：条件生成与条件评估；
• equivariant_diffusion/：扩散模型核心实现（en_diffusion、distributions、utils）；
• egnn/：EGNN 等变网络实现；
• qm9/：数据加载、损失函数、rdkit 辅助函数、可视化工具；
• qm9/property_prediction/：用于条件生成的性质预测分类器。

训练产物默认统一输出到outputs/实验名/目录下，包括模型权重、优化器状态和参数存档。

第一步：搭环境（conda + rdkit + 依赖）

EDM 依赖 RDKit 做分子后处理分析，建议用 conda 单独建一个环境，避免污染其他项目：

git clone https://gitcode.com/gh_mirrors/e3/e3_diffusion_for_molecules cd e3_diffusion_for_molecules conda create -c conda-forge -n my-rdkit-env rdkit conda activate my-rdkit-env pip install -r requirements.txt

一个小提示：即使没装 RDKit，训练流程本身也能跑，只是分子分析环节会受限。所以除非你只跑纯生成，否则建议老老实实把 rdkit 装上。

第二步：在 QM9 上训练你的第一个模型

项目官方推荐的 QM9 训练命令如下，建议直接复制使用：

python main_qm9.py \ --n_epochs 3000 \ --exp_name edm_qm9 \ --n_stability_samples 1000 \ --diffusion_noise_schedule polynomial_2 \ --diffusion_noise_precision 1e-5 \ --diffusion_steps 1000 \ --diffusion_loss_type l2 \ --batch_size 64 \ --nf 256 \ --n_layers 9 \ --lr 1e-4 \ --normalize_factors [1, 4, 10] \ --test_epochs 20 \ --ema_decay 0.9999

几个关键参数先混个脸熟：

• diffusion_steps：扩散步数，影响采样质量，数值越大生成越精细，但采样耗时线性增长；
• nf和n_layers：EGNN 网络的宽度和深度，直接决定模型容量；
• normalize_factors：对坐标、类别特征、整数特征（电荷）分别做归一化的系数，[1, 4, 10]是官方在 QM9 上调好的经验值；
• ema_decay：指数滑动平均系数，0.9999 能让训练过程更稳，生成质量更好；
• n_stability_samples：每隔test_epochs轮采样多少个分子用于稳定性评估。

训练过程中，模型会周期性在验证集和测试集上计算负对数似然（NLL）作为损失参考，并采样一批分子做稳定性分析。项目默认接入 WandB 记录训练曲线，如果你不想用，加一个--no_wandb参数即可。

第三步：训练完怎么验收（分析 + 可视化）

训练结束不等于万事大吉，你需要客观评估生成质量。运行：

python eval_analyze.py --model_path outputs/edm_qm9 --n_samples 10000

脚本会采样 1 万个分子，借助 RDKit 计算一系列指标：分子有效性、原子稳定性、整体稳定性，以及键长、键角等几何分布是否合理。这些数字能帮你判断模型是"真的学会了化学"，还是"只会复读训练集"。

想看生成分子的长什么样，再用：

python eval_sample.py --model_path outputs/edm_qm9 --n_samples 10000

它会输出可视化结果，你可以直观检查生成的分子结构是否自然、有无不合理的长键或原子重叠。

进阶玩法：按目标性质做条件生成

无条件的随机生成只是入门，药物发现里更常见的是"我想要带特定性质的分子"。EDM 支持对以下性质做条件生成：alpha（极化率）、gap（能隙）、homo、lumo、mu（偶极矩）、Cv（热容）。

以 alpha 为例，训练条件模型：

python main_qm9.py \ --exp_name exp_cond_alpha \ --model egnn_dynamics \ --lr 1e-4 --nf 192 --n_layers 9 \ --save_model True --diffusion_steps 1000 \ --sin_embedding False --n_epochs 3000 \ --n_stability_samples 500 \ --diffusion_noise_schedule polynomial_2 \ --diffusion_noise_precision 1e-5 \ --dequantization deterministic \ --include_charges False \ --diffusion_loss_type l2 --batch_size 64 \ --normalize_factors [1,8,1] \ --conditioning alpha \ --dataset qm9_second_half

注意这里用了qm9_second_half，即只用训练集的后 5 万个样本，这是官方条件实验的默认设置。

然后扫描不同性质取值，观察生成分子的性质分布：

python eval_conditional_qm9.py \ --generators_path outputs/exp_cond_alpha \ --property alpha --n_sweeps 10 --task qualitative

如果想定量评估"生成分子是否真的具有目标性质"，还需要一个性质预测分类器来"打分"。先训练分类器：

cd qm9/property_prediction python main_qm9_prop.py \ --num_workers 2 --lr 5e-4 \ --property alpha --exp_name exp_class_alpha --model_name egnn

再用它评估 EDM 生成样本：

python eval_conditional_qm9.py \ --generators_path outputs/exp_cond_alpha \ --classifiers_path qm9/property_prediction/outputs/exp_class_alpha \ --property alpha --iterations 100 --batch_size 100 --task edm

官方还提供了预训练的 alpha 生成器和分类器权重，对应参数为--generators_path outputs/exp_35_conditional_nf192_9l_alpha和--classifiers_path qm9/property_prediction/outputs/exp_class_alpha_pretrained，可以拿来直接对比基线。

换数据集：跑 GEOM-Drugs 需要注意什么

QM9 是小分子玩具，想生成药物级别的分子就要换到 GEOM-Drugs 数据集。训练命令大同小异：

python main_geom_drugs.py \ --n_epochs 3000 --exp_name edm_geom_drugs \ --n_stability_samples 500 \ --diffusion_noise_schedule polynomial_2 \ --diffusion_steps 1000 \ --diffusion_noise_precision 1e-5 \ --diffusion_loss_type l2 \ --batch_size 64 --nf 256 --n_layers 4 \ --lr 1e-4 --normalize_factors [1,4,10] \ --test_epochs 1 --ema_decay 0.9999 \ --normalization_factor 1 \ --model egnn_dynamics \ --visualize_every_batch 10000

两点提醒：

1. 运行前需要先按data/geom/README.md的说明准备 GEOM 数据，这一步比较耗时；
2. 显存占用会明显变大。EGNN 采用全连接的消息传递机制，分子越大，计算越吃内存。官方实验用的 GPU 相当大，如果你的卡放不下，优先把batch_size调小，其次考虑减小nf或n_layers。

常见坑与调优建议

结合社区反馈和项目源码，整理几个高频问题：

• 显存溢出（OOM）：优先降低batch_size，这是最直接的方案；再考虑调低nf。全连接消息传递是内存大户，别硬撑。
• 生成分子稳定性差：检查normalize_factors是否与数据集匹配（QM9 用[1,4,10]，条件实验用[1,8,1]）；尝试增大ema_decay到 0.9999，EMA 权重通常比原始权重更稳。
• 采样太慢：采样复杂度是 O(diffusion_steps)，可以把diffusion_steps从 1000 降到 500 甚至 200 试试，质量会有损失但速度翻倍。
• 不想用 WandB：训练命令加--no_wandb，或者显式传--wandb_usr 你的用户名绑定账号。
• 从头恢复训练：用--resume outputs/实验名配合--start_epoch，程序会从存档的args.pickle恢复超参数。
• 损失类型怎么选：默认l2简单稳定，适合快速出结果；vlb是基于变分下界的损失，理论更严谨，但训练更慢、更容易出数值问题。新手先用 l2。

常见问题 FAQ

Q：一定要用 conda 装 rdkit 吗？A：不强制。训练环节不依赖 RDKit，但eval_analyze.py等分析脚本需要它做化学验证，建议装。

Q：EDM 和普通的 3D 生成模型（如基于 GAN 或 VAE）比，优势在哪？A：扩散模型在生成多样性和稳定性上通常更优，而等变结构让模型在训练数据不足时也能正确泛化到分子的任意旋转姿态。

Q：--conditioning能同时指定多个性质吗？A：脚本支持传入多个值（如--conditioning homo lumo），但需要对应的性质数据都存在，实际项目中建议从单个性质开始实验。

Q：怎么判断训练有没有收敛？A：看两个信号：验证/测试 NLL 是否趋于平稳；n_stability_samples采样的分子稳定性指标是否持续上升。EMA 权重下的指标才是最终参考。

延伸资源

想深入理解，可以按这个顺序阅读项目源码：

• 扩散过程与去噪实现：equivariant_diffusion/en_diffusion.py
• 等变网络 EGNN 细节：egnn/egnn.py 与 egnn/egnn_new.py
• 数据加载与原子编码：qm9/dataset.py
• 分子后处理与 rdkit 接口：qm9/rdkit_functions.py
• 条件生成全流程：eval_conditional_qm9.py
• 配置模板参考：configs/qm9_config.yaml

另外，generated_samples/目录下附带了一些官方采样结果（如samples_edm.zip和 gschnet 分析脚本），可以拿来做参照系，对比你自己训练出来的样本质量。

最后给一句总结：EDM 的价值不只是"能生成分子"，而是用等变性把化学直觉真正编码进了模型结构。按照本文的顺序跑通一遍，你就能亲手验证这套机制，并在这个基础上做属于自己的分子生成实验。祝你跑得顺利，生成出漂亮的分子。

【免费下载链接】e3_diffusion_for_molecules项目地址: https://gitcode.com/gh_mirrors/e3/e3_diffusion_for_molecules