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
分别运行代码检查和测试。
更多信息,请参见贡献指南。