从 0 到 1 上手 E(3) 等变扩散模型生成三维分子的完整实操指南【免费下载链接】e3_diffusion_for_molecules项目地址: https://gitcode.com/gh_mirrors/e3/e3_diffusion_for_molecules如果你正在研究分子生成很可能已经踩过这样的坑用二维图生成分子忽略了原子的真实空间排布或者硬生生把三维坐标塞进普通神经网络结果模型一遇到旋转、平移就失忆。E(3) 等变扩散模型EDMEquivariant 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.pyQM9 数据集的主训练入口main_geom_drugs.pyGEOM-Drugs 数据集训练入口eval_analyze.py分析生成分子的质量稳定性、有效性等指标eval_sample.py采样并可视化生成的分子eval_conditional_qm9.py条件生成与条件评估equivariant_diffusion/扩散模型核心实现en_diffusion、distributions、utilsegnn/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_layersEGNN 网络的宽度和深度直接决定模型容量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两点提醒运行前需要先按data/geom/README.md的说明准备 GEOM 数据这一步比较耗时显存占用会明显变大。EGNN 采用全连接的消息传递机制分子越大计算越吃内存。官方实验用的 GPU 相当大如果你的卡放不下优先把batch_size调小其次考虑减小nf或n_layers。常见坑与调优建议结合社区反馈和项目源码整理几个高频问题显存溢出OOM优先降低batch_size这是最直接的方案再考虑调低nf。全连接消息传递是内存大户别硬撑。生成分子稳定性差检查normalize_factors是否与数据集匹配QM9 用[1,4,10]条件实验用[1,8,1]尝试增大ema_decay到 0.9999EMA 权重通常比原始权重更稳。采样太慢采样复杂度是 O(diffusion_steps)可以把diffusion_steps从 1000 降到 500 甚至 200 试试质量会有损失但速度翻倍。不想用 WandB训练命令加--no_wandb或者显式传--wandb_usr 你的用户名绑定账号。从头恢复训练用--resume outputs/实验名配合--start_epoch程序会从存档的args.pickle恢复超参数。损失类型怎么选默认l2简单稳定适合快速出结果vlb是基于变分下界的损失理论更严谨但训练更慢、更容易出数值问题。新手先用 l2。常见问题 FAQQ一定要用 conda 装 rdkit 吗A不强制。训练环节不依赖 RDKit但eval_analyze.py等分析脚本需要它做化学验证建议装。QEDM 和普通的 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考