E(3)等变扩散模型分子生成入门指南一次跑通训练、评估与条件生成【免费下载链接】e3_diffusion_for_molecules项目地址: https://gitcode.com/gh_mirrors/e3/e3_diffusion_for_molecules做分子生成的你一定遇到过这样的尴尬用 SMILES 或者 2D 图模型生成的结果放进三维空间里一看原子坐标是乱的、键长是离谱的甚至根本拼不出一个合理构象。问题不在于模型不够强而在于它从设计之初就没把分子的三维几何当作第一公民来对待。今天要介绍的E(3) 等变扩散模型EDMEquivariant Diffusion for Molecule Generation in 3D正是为了解决这个痛点而生的它把扩散模型和等变图神经网络结合直接在三维空间里联合生成原子的坐标与类型是目前分子生成方向绕不开的基线模型。从盲人摸象到三维直出EDM 到底解决了什么传统分子生成大多走两条路一是字符串SMILES生成二是二维图生成。前者把分子压扁成一串符号后者只保留拓扑连接两者都丢失了真正的三维构象信息。而现实是——分子的一切性质从热力学稳定性到光谱响应都取决于它立体的空间排布。EDM 的思路可以这样理解把分子生成想象成从噪声里显影一张三维照片。训练阶段模型不断学习把干净的分子一步步加噪直到变成一锅高斯噪声的过程生成阶段模型反向操作从纯噪声出发一步步去噪最终显影出一个带坐标、带原子类型的完整分子。而支撑整个过程的神经网络是 E(3) 等变图神经网络EGNN它保证无论你把分子旋转、平移还是镜像翻转模型给出的分布都不变。这就像一位雕塑家不管他把雕塑坯料转到哪个角度都能刻出同一个作品。上图直观展示了这种设计上半部分是抽象的离散点集下半部分是真实分子中间的 R 表示对坐标施加旋转等 E(3) 变换而模型输出的概率分布保持不变p(x, h) p(Rx, h)。换句话说模型学到的不是一个死板的坐标系而是分子内在的几何规律——这正是它在构象生成上表现稳健的根本原因。技术定位上EDM 属于连续扩散模型它用一条时间轴上的噪声调度如多项式调度 polynomial_2、cosine把连续坐标和离散原子类型统一建模配合 EGNN 主干代码在 egnn/egnn.py 的 E_GCL 层实现了对三维欧氏群 E(3) 的等变性。三步完成环境搭建从零到能跑动手之前先把环境铺好。项目核心依赖只有几个torch、numpy、scipy、imageio、tqdm、wandb都在 requirements.txt 里。好消息是没有 rdkit 也能运行只是拿不到有效性Validity等化学指标——所以强烈建议装一个。第一步创建 rdkit 环境可选但推荐conda create -c conda-forge -n my-rdkit-env rdkit conda activate my-rdkit-env第二步克隆并安装依赖git clone https://gitcode.com/gh_mirrors/e3/e3_diffusion_for_molecules cd e3_diffusion_for_molecules pip install -r requirements.txt第三步验证 GPU 可用性python -c import torch; print(torch.cuda.is_available())新手常踩的坑项目源码里 RDKit 的 import 必须放在最前面main_qm9.py 里特意注释了 Rdkit import should be first如果你自定义脚本记得保持同样的导入顺序否则可能出现诡异的符号解析错误。跑通第一次训练QM9 上的标准流程环境就绪后用 QM9 数据集约 13 万个小分子做第一次完整训练。第一次运行会自动下载并预处理 GDB9 数据来自 figshare 的公开文件耐心等它完成即可之后数据会缓存到本地。训练命令直接采用官方推荐的参数组合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 1000是去噪步数直接影响采样质量与耗时--diffusion_loss_type l2用最简单的均方误差损失训练噪声预测网络稳定且省显存--ema_decay 0.9999启用指数滑动平均最终保存的 generative_model_ema.npy 通常比原始权重更平滑、更可靠--normalize_factors [1,4,10]分别对坐标、离散特征和整型特征做归一化新手不建议改动。训练过程中模型每--test_epochs轮做一次验证并在验证集 NLL 最优时把权重存到outputs/edm_qm9/。下面的图说明了训练阶段的损失设计模型被要求从带噪样本 z_t 中预测出被加进去的噪声 ε通过最小化预测误差来学会去噪。新手常踩的坑项目默认会初始化 wandb 并尝试在线同步。如果你没有 wandb 账号或只想本地跑务必在命令末尾加上--no_wandb否则程序可能卡在登录等待上。训练期间或训练结束后随时可以采样看看模型学会了什么。先生成一批分子并保存为 xyz 文件再做稳定性与化学合理性分析python eval_analyze.py --model_path outputs/edm_qm9 --n_samples 10000 python eval_sample.py --model_path outputs/edm_qm9 --n_samples 10000eval_analyze.py会输出原子稳定性stability、价键有效性Validity、独特性Uniqueness和新颖性Novelty等关键指标——这些指标依赖 rdkit没装的话命令行会明确提示你 Install rdkit to obtain Validity, Uniqueness, Novelty。eval_sample.py则把采样到的分子渲染成球棍模型还会生成一条从噪声到分子的去噪链动图直观展示生成过程。新手常踩的坑EGNN 使用全连接的消息传递显存消耗随原子数增长很快。如果显存不够先调小--batch_size和--nf而不是硬扛 OOM。三个实战场景从药物筛选到性质导向生成跑通基线只是开始EDM 的价值在于可以嵌入真实研发流程。这里给出三个可以立刻上手的场景。场景一性质条件生成让模型按需造分子EDA 在药物发现中的典型诉求是生成具有特定理化性质的分子。EDM 原生支持条件生成把某个性质如极化率 alpha、HOMO/LUMO 能级、能隙 gap、偶极矩 mu、热容 Cv作为条件输入模型。先用 QM9 后半段数据训练一个条件版模型python main_qm9.py --exp_name exp_cond_alpha --model egnn_dynamics \ --lr 1e-4 --nf 192 --n_layers 9 --diffusion_steps 1000 \ --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然后对性质值做扫描采样sweep观察不同 alpha 取值下分子结构如何渐变python eval_conditional_qm9.py --generators_path outputs/exp_cond_alpha \ --property alpha --n_sweeps 10 --task qualitative场景二训练属性分类器量化评估生成质量生成得像不像还要用数字说话。项目在qm9/property_prediction/下提供了 EGNN 属性分类器训练好后可以直接在生成样本上打分验证条件生成的准确性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 cd ../.. 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有意思的是--model_name numnodes可以训练一个只看原子数的基线分类器用来对照验证如果只看节点数都能预测属性说明模型根本没学到几何信息——这是检验生成器是否作弊的巧妙手段。场景三挑战更大体系——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注意这里把n_layers从 9 降到 4正是为了缓解全连接消息传递带来的显存压力。数据准备细节见项目内data/geom的说明文档。生态联动让 EDM 融入你的技术栈EDM 不是孤岛把它放进下面的生态里威力会成倍放大。RDKit化学信息学的瑞士军刀。EDM 的指标计算Validity/Uniqueness/Novelty和分子可视化都挂在 rdkit 上它让生成和化学合理性判断无缝衔接。想快速上手可以先跑一遍qm9/rdkit_functions.py里的BasicMolecularMetrics。PyTorch torchvision模型骨架、数据加载、EMA、DataParallel 多卡训练全部建立在 PyTorch 之上代码风格清晰egnn/egnn_new.py与egnn/egnn.py还能对照学习 EGNN 的两种实现差异。WandB训练脚本内置了实验跟踪自动记录损失曲线与采样稳定性离线模式或关闭都可通过命令行参数一键切换configs/qm9_config.yaml里也给出了完整的开关项。imageio / scipy负责去噪链动图与统计直方图等可视化产出生成过程一览无余。下一步三条路线任你选到这里你已经具备了从训练到评估再到条件生成的完整闭环能力。接下来可以这样继续深入读源码从main_qm9.py的入口开始顺着qm9/models.py到equivariant_diffusion/en_diffusion.py理解扩散损失与 NLL 计算再对照configs/qm9_config.yaml的每一项配置含义。看效果仓库generated_samples/里提供了预生成的样本samples_edm.zip 与 gschnet 对比样本无需训练就能先感受 EDM 的输出质量。做对比用eval_analyze.py的稳定性指标把 EDM 和仓库中 gschnet 的生成结果放一起比较你会更直观地理解等变扩散的优势所在。分子生成这个领域还在快速演进但 EDM 作为三维扩散生成的开山之作值得你花一个下午把它彻底跑通——它是你理解后续所有 3D 分子生成工作最好的起点。【免费下载链接】e3_diffusion_for_molecules项目地址: https://gitcode.com/gh_mirrors/e3/e3_diffusion_for_molecules创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考