gflownet  by alexhernandezgarcia

PyTorch library for training and extending Generative Flow Networks (GFlowNets)

created 2 years ago
265 stars

Top 97.2% on sourcepulse

GitHubView on GitHub
Project Summary

This library provides a PyTorch-based framework for Generative Flow Networks (GFNs), enabling probabilistic and generative modeling with applications in scientific discovery. It allows users to sample diverse objects with high rewards by decomposing generation into compositional states and learning transition probabilities, targeting researchers and engineers in fields like drug discovery, materials science, and reinforcement learning.

How It Works

GFNs model the generation of objects by learning to navigate a state space, transitioning from a source state to a final state. The library implements this by training neural network policies to approximate forward and backward state transitions. It utilizes a "proxy" function to define rewards for generated states, allowing the GFN agent to optimize sampling towards high-reward configurations through various loss functions like Flow Matching (FM) and Trajectory Balance (TB).

Quick Start & Requirements

  • Install: Clone the repo and run source install.sh.
  • Prerequisites: Python 3.10, CUDA 11.8 (for GPU), Ubuntu (recommended, OSX supported with manual dependency handling).
  • Dependencies: PyTorch 2.5.1, torch-scatter. Optional extras: dev, materials, molecules.
  • Setup: Installation script simplifies environment setup. See installation guide for detailed steps and CPU-only options.
  • Training: Use python train.py with Hydra configurations (e.g., +experiments=grid/corners).
  • Evaluation: Use python eval.py rundir=<path_to_run_directory>.
  • Resuming: Use python resume.py rundir=<path_to_run_directory>.

Highlighted Details

  • Implements core GFN components: environment, proxy, policy models, and agent.
  • Supports multiple GFN objectives: Flow Matching (FM), Trajectory Balance (TB), Detailed Balance (DB), Forward-Looking (FL).
  • Includes example environments like Scrabble and Grid, with extensibility for custom environments and proxies.
  • Integrates with Weights & Biases for logging.

Maintenance & Community

  • Active contributors include Alex Hernandez-Garcia, Nikita Saxena, Alexandra Volokhova, and others.
  • Used in several research articles, including ICML 2023 and NeurIPS 2023 workshops.
  • Project citation available in BibTeX and CFF formats.

Licensing & Compatibility

  • The repository does not explicitly state a license in the README.

Limitations & Caveats

  • Setup is primarily supported on Ubuntu, with OSX requiring manual dependency management.
  • The README notes that PyTorch defaults to CPU due to unobserved performance improvements on GPU, suggesting manual configuration (device=cuda) is needed to leverage GPUs.
Health Check
Last commit

2 weeks ago

Responsiveness

1 day

Pull Requests (30d)
1
Issues (30d)
0
Star History
19 stars in the last 90 days

Explore Similar Projects

Starred by Chip Huyen Chip Huyen(Author of AI Engineering, Designing Machine Learning Systems), Nathan Lambert Nathan Lambert(AI Researcher at AI2), and
1 more.

tianshou by thu-ml

0.1%
9k
PyTorch RL library for algorithm development and application
created 7 years ago
updated 1 day ago
Feedback? Help us improve.