gflownet

用于图结构的GFlowNet相关训练和环境代码。

简介

GFlowNet [1], [2], [3]，全称为生成流网络，是一种新颖的生成建模框架，特别适用于离散的组合对象。这里特别实现了图生成。

GFN的核心思想是估计有向无环网络*中的流。该网络表示构建对象的所有可能方式，因此了解流可以为我们提供一个策略，我们可以按顺序构建对象。这种部分构建对象的序列称为"轨迹"。*也许令人困惑的是，GFN中的"网络"指的是状态空间，而不是神经网络架构。

这个库的主要重点（尽管它可以做其他事情）是构建图（例如原子图），这些图是逐节点构建的。为了进行策略预测，我们使用图神经网络。这个GNN输出每个节点的logits（例如，向这个原子添加一个原子，或在这两个原子之间添加一个键），以及每个图的logits（例如，停止/"完成构建此对象"）。

这个库支持多种GFN算法（以及一些基线），并支持在现有数据（离线）和自生成数据（在线）的混合上进行训练，后者是通过顺序查询模型获得轨迹。

安装

PIP

这个包可以作为PIP包安装，但由于它依赖于一些torch-geometric包的轮子，必须指定--find-links参数：

pip install -e . --find-links https://data.pyg.org/whl/torch-2.1.2+cu121.html

或者用于CPU：

pip install -e . --find-links https://data.pyg.org/whl/torch-2.1.2+cpu.html

要安装或依赖特定标签，例如这里的v0.0.10，使用以下方案：

pip install git+https://github.com/recursionpharma/gflownet.git@v0.0.10 --find-links ...

如果包依赖似乎不起作用，你可能需要安装requirements/中列出的确切冻结版本，即pip install -r requirements/main-3.10.txt。

入门

立即开始的好地方是sEH基于片段的MOO任务。文件seh_frag_moo.py可以直接运行（尽管你可能想更改main()中的默认配置）。

对于更温和的库介绍，请参阅入门指南。对于更深入了解该库，请参阅实现说明。

仓库概览

• algo，包含GFlowNet算法实现（轨迹平衡，SubTB，流匹配），以及一些基线。这些实现了如何从模型中采样轨迹和从轨迹计算损失。
• data，包含数据集定义、数据加载和数据采样工具。
• envs，包含环境类；基础环境与正在制作的图类型无关，上下文类指定从图到对象（例如分子）和torch geometric Data的映射。
• examples，包含GFlowNet的简单示例实现。
• models，包含模型定义。
• tasks，包含训练代码。
  • qm9，基于QM9的HOMO-LUMO间隙数据作为奖励的温度条件分子采样器。
  • seh_frag，复现Bengio等人2021年的工作，基于片段的分子设计，针对sEH蛋白。
  • seh_frag_moo，与上述相同，但进行多目标优化（包括QED、SA和分子重量目标）。
• utils，包含工具（多进程、指标、条件设置）。
• trainer.py，定义了训练GFlowNet模型的通用框架。
• online_trainer.py，定义了典型的在线GFN训练循环。

更多信息请参见实现说明。

开发与贡献

欢迎外部贡献。

要安装开发者依赖项

pip install -e '.[dev]' --find-links https://data.pyg.org/whl/torch-2.1.2+cu121.html

我们使用tox运行测试和代码检查，使用pre-commit在提交前运行检查。 为确保这些检查通过，只需运行tox -e style和tox run分别运行代码检查和测试。

更多信息，请参见贡献指南。