Skip to content

Getting Started

This guide will help you get started with TorchEBM by walking you through the installation process and demonstrating some basic usage examples.

Installation

TorchEBM can be installed directly from PyPI:

pip install torchebm

Prerequisites

  • Python 3.8 or newer
  • PyTorch 1.10.0 or newer
  • CUDA (optional, but recommended for performance)

Installation from Source

If you wish to install the development version:

git clone https://github.com/soran-ghaderi/torchebm.git
cd torchebm
pip install -e .

Quick Start

Here's a simple example to get you started with TorchEBM:

import torch
from torchebm.core import GaussianEnergy
from torchebm.samplers.langevin_dynamics import LangevinDynamics

# Set device for computation
device = "cuda" if torch.cuda.is_available() else "cpu"

# Define a 2D Gaussian energy function for visualization
energy_fn = GaussianEnergy(
    mean=torch.zeros(2, device=device),
    cov=torch.eye(2, device=device)
)

# Initialize Langevin dynamics sampler
sampler = LangevinDynamics(
    energy_function=energy_fn,
    step_size=0.01,
    device=device
).to(device)

# Generate 1000 samples
samples = sampler.sample_chain(
    dim=2,
    n_steps=100,
    n_samples=1000,
    return_trajectory=False
)

print(f"Generated {samples.shape[0]} samples of dimension {samples.shape[1]}")

Next Steps

Common Issues

CUDA Out of Memory

If you encounter CUDA out of memory errors, try: - Reducing the number of samples - Reducing the dimension of the problem - Switching to CPU if needed

Support

If you encounter any issues or have questions: - Check the FAQ - Open an issue on GitHub