aiqm/torchani

TorchANI 2.0 – PyTorch 中的神經網絡原子間勢能

TorchANI 是一個基於 PyTorch 構建的開源 Python 函式庫，允許研究人員訓練、評估及使用 ANI 風格的神經網絡原子間勢能 (NN-IP)。這些是機器學習模型，可以直接根據原子的座標預測分子的勢能與作用力，從而實現快速、準確的分子動力學 (MD) 及混合 ML/MM 模擬。

函式庫功能

• 模型庫 – 提供即插即用的預訓練 ANI 模型以及定義新模型的簡潔 API。
• 訓練工具 – 將 NN-IP 擬合到量子化學參考數據的工具。
• 快速推理 – 透過 C++/CUDA 擴充功能計算 ANI 描述符並在 GPU 上執行網絡，實現近乎實時的 MD。
• 命令列介面 – ani 可執行檔提供快速工具（例如：能量/力評估、構建擴充功能）。
• 整合點 – 獨立的 TorchANI-Amber 套件允許您將勢能引入 Amber 的經典 MD 引擎，用於 ML/MM 模擬。

適用對象

• 需要用於模擬的高性能、基於 ML 的力場的計算化學家或材料科學家。
• 正在開發新的 NN-IP 架構並希望擁有 PyTorch 原生基準的研究人員。
• 希望在無需編寫底層 CUDA 程式碼的情況下，在 GPU 上執行 ML 驅動的分子動力學的任何人。

入門指南 (摘自 README)

1. 設置 Python 環境 (conda/mamba 或 virtualenv)。該函式庫在 PyPI 上發布，但作者建議即使在 conda 環境中也使用 pip 進行安裝，因為 conda 套件已不再維護。
2. 安裝相容的 PyTorch (例如，針對 CUDA 13.0，使用 pip install torch==2.13 --index-url https://download.pytorch.org/whl/cu130)。TorchANI 2 需要 PyTorch ≥ 2.0。
3. 安裝 TorchANI: pip install torchani。
4. 為描述符計算與快速推理構建選配的 C++/CUDA 擴充功能:

   ani build-extensions          # 為所有偵測到的 GPU SM 構建
   # 或限制於特定 SM，例如：
   ani build-extensions --sm 8.0 --sm 8.9
   

5. （選配）如果您需要構建文件、執行測試或修改原始碼，請克隆儲存庫並透過提供的 environment.yaml 建立開發環境。

平台說明

• GPU 支援：適用於帶有 CUDA 的 NVIDIA GPU；對於任何非平凡的工作負載，強烈建議使用。未測試 AMD/ROCm。
• macOS：無 CUDA；該函式庫在 Apple 的 MPS 後端上尚未經過測試。
• Conda 套件：面向開發者提供；可以使用 recipe 目錄在本地構建。

文件與幫助

• 完整的用戶指南與 API 參考：https://aiqm.github.io/torchani/
• 從 TorchANI 1.x 遷移到 2.0 的遷移指南。
• 歡迎透過 GitHub issues 提交錯誤報告與功能請求。

引用 TorchANI

如果您在發表的著作中使用該函式庫，請引用 README 中列出的兩篇論文（原始 TorchANI 論文與 TorchANI 2.0 擴充論文）。

---

摘要 (TL;DR): TorchANI 2.0 是一個基於 PyTorch 的工具包，用於構建與執行分子模擬的神經網絡勢能，提供 GPU 加速推理、訓練工具以及與 Amber 等經典 MD 套件的整合。這是一個面向計算化學界的真正的 AI/ML 軟體專案。