AI力场二次开发教程（02）：环境搭建——conda/mamba 安装全套 AI 力场栈并验证 GPU

AI力场二次开发教程（2）：环境搭建——用 conda/mamba 安装全套 AI 力场栈并验证 GPU

版本声明：本文使用mamba（conda-forge 通道）构建，目标版本为espaloma=0.3.2、openff-toolkit=0.19.0、OpenMM、OpenMMForceFields、PyTorch(支持 CUDA)；语言/环境为 Linux + Python 3.10/3.11。本文目标是让你在一台新机器上，用一组命令装出可复现、可跑通esp.get_model("latest")的完整环境。所有版本号及安装命令基于 conda-forge 真实包约束，个别依赖以官方文档为准。

一句话结论：用mamba create -n espaloma -c conda-forge "espaloma=0.3.2" openff-toolkit openmm openmmforcefields "pytorch=cuda12*" pytorch-cuda即可构建核心栈，之后用python -c "import torch; print(torch.cuda.is_available())"验证 GPU，并跑通get_model("/path/espaloma-0.3.2.pt")完成推理。

〇、本篇要解决的认知问题

1. 为什么推荐mamba而不是pip直装这一整套 AI 力场栈？conda-forge 解决了哪些依赖矛盾？
2. 正确的安装命令是什么？espaloma与openff-toolkit/OpenMM之间有哪些必须锁定的版本约束？
3. CPU 与 CUDA 两种 PyTorch 如何选择、如何验证torch.cuda.is_available()为真？
4. 从本地.pt权重加载模型应该做什么处理（如model.eval()）？esp.get_model("latest")与本地权重有什么区别？
5. 装了却不生效的常规场景（kernel 缓存、缺openff-forcefields、版本错配）如何定位？

一、机制解析

1.1 为什么是 conda-forge + mamba

Espaloma 生态包含 Python 包与二进制/编译依赖（OpenMM 的 CUDA 插件、libgpuarray、RDKit 等）。纯pip无法原生解析 CUDA 运行时与 C++ 扩展的二进制冲突。conda-forge提供了一套统一元数据 + 二进制分发的依赖图，mamba用更快求解器（libsolv）处理冲突，通常比conda原版 solve 快数倍且更少报“conflict”。

因此标准姿势是：

mamba create-nespaloma\-cconda-forge\"espaloma=0.3.2"\"openff-toolkit=0.19.0"\"openmmforcefields"\"openmm"\"pytorch=cuda12*""pytorch-cuda=12.1"

说明：pytorch=cuda12*与pytorch-cuda来自 conda-forge 的 pytorch 元包约定，具体 CUDA 版本号以你本机驱动支持的为准；若机型无 NVIDIA GPU，去掉这两行改为pytorch-cpu即可。

1.2espaloma=0.3.2的依赖画像

Espaloma 0.3.2 的依赖包括但不限于：torch、dgl、rdkit、openff-toolkit（>=某下限）、openff-units、typing-extensions等。它通过esp.get_model("id")或esp.get_model("/path/espaloma-0.3.2.pt")加载权重。已知.pt权重兼容规则为：espaloma-0.3.2.pt兼容 0.3.1 / 0.3.2 / 0.4.0（以官方发布说明为准）。DGL（Deep Graph Library）负责异构图算子，与 PyTorch 版本需匹配。

1.3 本地权重的eval()原则

用esp.get_model(path)加载本地.pt后，推理前应确保模型处于eval（评估）模式，即调用model.eval()。理由是：训练阶段需要的 dropout / batch-norm 统计行为在推理时不适用，eval()保证批归一化使用训练均值/方差。Espaloma 官方示例即在使用前对espaloma_charge做个eval()。若直接推理，虽然多数网络仍能给出结果，但数值可能不以与训练一致。

1.4 验证路径：从“能 import”到“能跑”

“环境搭好了”不等于“能跑通”。我们按三层递进验证：

1. import层：torch、espaloma、openff.toolkit、openmm都能 import。
2. GPU 层：torch.cuda.is_available()返回 True，且torch.cuda.get_device_name(0)打印 GPU 型号。
3. 端到端层：esp.get_model(...)成功，对咖啡因分子图前向一次推理并产出参数。

二、完整代码与逐行剖析

2.1 创建与激活环境（脚本，复制即跑）

# 在 linux shell 下执行；本段为一次性环境搭建mamba create-nespaloma-cconda-forge\"python=3.10"\"espaloma=0.3.2"\"openff-toolkit=0.19.0"\"openmm"\"openmmforcefields"\"pytorch=cuda12*""pytorch-cuda=12.1"-y# 激活环境mamba activate espaloma# 从中量依赖里确认核心版本python-c"import espaloma, openff.toolkit, openmm; \ print('espaloma', espaloma.__version__); \ print('openff-toolkit', openff.toolkit.__version__); \ print('openmm', openmm.__version__)"

逐行剖析：

• -c conda-forge：指定唯一通道，避免 defaults 与 forge 混合导致 URL 冲突。
• 显式python=3.10：锁死解释器版本，Espaloma 0.3.2 的二进制依赖多在 3.8–3.11 区间，3.12+ 常因dgl未预编译而失败。
• pytorch=cuda12*：取 CUDA 12.x 系 PyTorch，与pytorch-cuda配套；若无 GPU 改pytorch-cpu。

若mamba activate不生效，先执行source $(mamba info --base)/etc/profile.d/conda.sh。

2.2 验证 GPU（关键一跑）

importtorch# 验证 CUDA 是否可用print("CUDA 可用：",torch.cuda.is_available())iftorch.cuda.is_available():print("GPU 名称：",torch.cuda.get_device_name(0))x=torch.rand(4,4,device="cuda")# 实际执行一次张量搬运print("GPU 张量：",x.dtype,x.device)else:print("未检测到 GPU，本次会退化为 CPU 推理")

逐行剖析：

• torch.cuda.is_available()返回布尔值：True 代表 driver 与 CUDA 运行时都可用。
• 进一步把张量搬到cuda设备强制发生一次实际分配，避免“假阳性”——部分环境 import 不报错但设备不可用。
• 输出x.device应为cuda:0，否则说明拿到的是 CPU 张量。

2.3 跑通 Espaloma 最短示例（从本地 .pt 推理）

importespalomaasespfromopenff.toolkit.topologyimportMolecule# 1) 构造分子图molecule=Molecule.from_smiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")# 咖啡因molecule_graph=esp.Graph(molecule)# 2) 加载模型：优先本地权重（注意 eval()）try:model=esp.get_model("espaloma-0.3.2.pt")# 本地 .pt 权重路径exceptException:model=esp.get_model("latest")# 回退到自动拉取model.eval()# 关键：推理前进入评估模式# 3) 前向推理model(molecule_graph.heterograph)print("推理完成，原子数：",molecule_graph.heterograph["n1"].shape[0])

逐行剖析：

• esp.get_model("espaloma-0.3.2.pt")：若当前目录/缓存有该权重文件则加载本地权重；esp.get_model("latest")则从网络拉取官方最新权重。
• 权重兼容：espaloma-0.3.2.pt兼容 0.3.1 / 0.3.2 / 0.4.0，因此不一定要把包版本与权重严格同号。
• model.eval()至关重要，见 §1.3；若想偏向可重复性，可在推理后model.eval()前若有 torch 的torch.set_grad_enabled(False)包裹推理。

说明：esp.Graph与 heterograph 的内部结构（n1/n2/n3/f1/f2/fi1 等字段）将在第 03 篇展开。

三、常见报错与排查

上面的dgl版本约束以官方发布说明为准；若你的情况与表不符，先打印完整Solve failed文本再对照。

四、动手练习

1. 从零重建：删除当前espaloma环境，用mamba create一分钟内重建，并记录 solve 耗时对比原conda create。
2. 双模式验证：在 NVIDIA 机器上分别用pytorch=cuda12*与pytorch-cpu建两个环境，运行 §2.2 脚本，把两份输出贴到对比表。
3. 权重兼容验证：把espaloma-0.3.2.pt分别配合 0.3.1 与 0.4.0 环境加载，各跑一次咖啡因推理，记录是否一致、有无警告。
4. 离线兜底：在无外网沙箱中，仅用本地.pt权重完成 §2.3 的推理，写一句结论说明“本地权重 vs 网络拉取”的使用时机。

五、小结与下一篇预告

本套环境已就绪：mamba create一次构建espaloma、openff-toolkit、OpenMM、OpenMMForceFields与 CUDA 版 PyTorch；用torch.cuda.is_available()验证 GPU；能用esp.get_model+model.eval()对咖啡因跑通图推理。你已经把“看得懂”变成了“能跑起来”。

下一篇进入第一个真正的“分子”落地：第 03 篇《第一个AI力场分子——SMILES→OpenMM System 全流程》，我们会完整复现锚点 A 的Molecule.from_smiles → esp.Graph → get_model → 前向 → openmm_system_from_graph，并逐段剖析异构图字段与能量项个数。

本篇认知问题回显（FAQ）

Q：为什么用 conda-forge 加 mamba 构建 espaloma 力场环境而不直接用 pip？

A：Espaloma 依赖 OpenMM 的 CUDA 插件与 DGL/PyTorch 二进制，pip 无法统一解析 CUDA 运行时与 C++ 扩展冲突；conda-forge 提供统一依赖图，mamba 用 libsolv 快速求解并减少 conflict。

Q：安装 espaloma 全套 AI 力场栈时哪些版本必须保持约束一致？

A：需配套指定 python=3.10、espaloma=0.3.2、openff-toolkit=0.19.0，并用 pytorch-cuda 锁定 CUDA 运行时；DGL 与 PyTorch 版本必须匹配，否则推理前就报错。

Q：怎样设置 CPU 与 CUDA 两种 PyTorch，并用 torch.cuda.is_available 验证生效？

A：有 GPU 时安装 pytorch-cuda 与 pytorch=cuda12*，无 GPU 换 pytorch-cpu；再用torch.cuda.is_available()返回 True，并把随机张量搬到 cuda:0 做一次实际分配校验。

Q：加载本地 espaloma .pt 权重时应如何处理模型推理模式？

A：用esp.get_model("本地.pt")加载后调用model.eval()进入评估模式；否则训练阶段的 dropout/批归一化行为会污染推理数值，与训练口径不一致。

Q：装了包却调用失败（kernel 缓存、openff-forcefields 缺失、版本错配）如何定位？

A：先source $(mamba info --base)/etc/profile.d/conda.sh修正激活路径，再装 openff-forcefields 补齐 offxml，最后核对 python=3.10 与 pytorch-cuda 是否仍匹配，用版本打印命令逐层确认。