Skip to content
TorchEBM Logo

PyTorch Toolkit for Generative Modeling

A high-performance PyTorch library that makes Energy-Based Models accessible and efficient for researchers and practitioners alike.

PyPI License GitHub Stars Ask DeepWiki Build Status Docs Downloads Python Versions

TorchEBM provides components for 🔬 sampling, 🧠 inference, and 📊 model training.

What is 🍓 TorchEBM?

TorchEBM is a PyTorch library for Energy-Based Models (EBMs), a powerful class of generative models. It provides a flexible framework to define, train, and generate samples using energy-based models.

Core Components

TorchEBM is structured around several key components:

  • Models

    Define energy functions using BaseModel, from analytical forms to custom neural networks.

    Details

  • Samplers

    Generate samples with MCMC samplers like Langevin Dynamics and Hamiltonian Monte Carlo.

    Details

  • Loss Functions

    Train models with loss functions like Contrastive Divergence and Score Matching.

    Details

  • Datasets

    Use synthetic dataset generators for testing and visualization.

    Details

  • Visualization

    Visualize energy landscapes, sampling, and training dynamics.

    Details

  • Accelerated Computing

    Accelerate sampling and training with CUDA implementations.

    Details

Quick Start

Install the library using pip:

1
pip install torchebm

Here's a minimal example of defining an energy function and a sampler:

  • Create and Sample from Energy Models

     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
    import torch
from torchebm.core import GaussianModel
from torchebm.samplers import LangevinDynamics

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GaussianModel(mean=torch.zeros(2), cov=torch.eye(2), device=device)

sampler = LangevinDynamics(model=model, step_size=0.01, device=device)

initial_points = torch.randn(500, 2, device=device)
samples = sampler.sample(x=initial_points, n_steps=100)

print(f"Output batch_shape: {samples.shape}") # (B, len) -> torch.Size([500, 2])

Latest Release

TorchEBM is currently in early development. Check our GitHub repository for the latest updates and features.

Community & Contribution

TorchEBM is an open-source project developed with the research community in mind.

  • Bug Reports & Feature Requests: Please use the GitHub Issues.
  • Contributing Code: We welcome contributions! Please see the Contributing Guidelines. Consider following the Commit Conventions.
  • Show Support: If you find TorchEBM helpful for your work, please consider starring the repository on GitHub! ⭐

Citation

Please consider citing the TorchEBM repository if it contributes to your research:

1
2
3
4
5
6
@misc{torchebm_library_2025,
  author       = {Ghaderi, Soran and Contributors},
  title        = {TorchEBM: A PyTorch Library for Training Energy-Based Models},
  year         = {2025},
  url          = {https://github.com/soran-ghaderi/torchebm},
}

License

TorchEBM is available under the MIT License. See the LICENSE file for details.