TorchEBM

https://soran-ghaderi.github.io/torchebm/latest/
Admin
TorchEBM is a high-performance PyTorch library that makes Energy-Based Models accessible and...
Tokens:449,671
Snippets:2,958
Trust Score:9.7
Update:5 months ago
Show doc for...
Tokens:
Raw
### Clone and Run TorchEBM Examples

Source: https://soran-ghaderi.github.io/torchebm/latest/examples

Instructions on how to clone the TorchEBM repository, set up the environment, list available examples, and run a specific example using the main script.

```bash
# Clone the repository
git clone https://github.com/soran-ghaderi/torchebm.git
cd torchebm

# Set up your environment
pip install -e .

# List all available examples
python examples/main.py --list

# Run a specific example
python examples/main.py samplers/langevin/visualization_trajectory
```

--------------------------------

### Clone and Run TorchEBM Examples

Source: https://soran-ghaderi.github.io/torchebm/latest/examples?q=

Instructions for cloning the TorchEBM repository, setting up the environment, listing available examples, and running a specific example using the main.py script.

```bash
git clone https://github.com/soran-ghaderi/torchebm.git
cd torchebm

pip install -e .

python examples/main.py --list

python examples/main.py samplers/langevin/visualization_trajectory
```

--------------------------------

### Install Documentation Dependencies

Source: https://soran-ghaderi.github.io/torchebm/latest/developer_guide/getting_started

Install additional dependencies required for working with the project's documentation. This enables serving a live preview of the documentation locally.

```bash
pip install -e ".[docs]"
mkdocs serve                     # live preview at http://127.0.0.1:8000
```

--------------------------------

### Adaptive Integration Setup and Execution

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_integrator

Handles the setup and execution of adaptive integration. Raises a ValueError if error_weights or order are not defined, which are required for adaptive=True. It calculates the initial step size and integrates the state adaptively.

```python
if self.error_weights is None or self.order is None:
            raise ValueError(
                f"{type(self).__name__} does not define error_weights/order "
                f"and cannot be used with adaptive=True."
            )

        x = state["x"]
        drift_fn = self._resolve_drift(drift)

        if not torch.is_tensor(step_size):
            step_size = torch.tensor(
                step_size, device=x.device, dtype=x.dtype
            )

        if t is not None:
            if t.ndim != 1 or t.numel() < 2:
                raise ValueError("t must be a 1D tensor with length >= 2")
            t_start = t[0].item()
            t_end = t[-1].item()
        else:
            t_start = 0.0
            t_end = float(n_steps) * step_size.item()

        h = min(step_size.item(), t_end - t_start, self.max_step_size)
        x = self._adaptive_integrate(x, drift_fn, t_start, t_end, h)
        return {"x": x}
```

--------------------------------

### Bosh3Integrator Example Usage

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/integrators/bosh3?q=

Example demonstrating how to initialize and use the Bosh3Integrator for solving an ordinary differential equation.

```python
from torchebm.integrators import Bosh3Integrator
import torch

integrator = Bosh3Integrator(atol=1e-4, rtol=1e-2)
state = {"x": torch.randn(100, 2)}
drift = lambda x, t: -x
result = integrator.integrate(
    state,
    step_size=0.1,
    n_steps=10,
    drift=drift,
)
```

--------------------------------

### Noise Scale Warmup Example

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core

Shows an example of using WarmupScheduler to control noise scale during a process.

```APIDOC
## Noise Scale Warmup

### Description
This example illustrates how WarmupScheduler can be applied to control a noise scale, starting with a warmup phase before transitioning to a linear decay.

### Usage
```python
linear_scheduler = LinearScheduler(
    start_value=1.0, end_value=0.01, n_steps=500
)
noise_scheduler = WarmupScheduler(
    main_scheduler=linear_scheduler,
    warmup_steps=25,
    warmup_init_factor=0.05
)
```
```

--------------------------------

### Install TorchEBM and Dependencies

Source: https://soran-ghaderi.github.io/torchebm/latest/developer_guide/getting_started

Clone the repository, set up a virtual environment, and install the project with development dependencies. This command installs TorchEBM in editable mode along with development tools.

```bash
git clone https://github.com/<your-user>/torchebm.git
cd torchebm
python -m venv .venv
source .venv/bin/activate        # Windows: .venv\Scripts\activate
pip install -e ".[dev]"
```

--------------------------------

### Direct Import Example

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/integrators?q=

Demonstrates how to directly import an integrator, such as LeapfrogIntegrator, from the torchebm.integrators submodule.

```APIDOC
## Direct Import

### Description
Allows for direct import of specific integrator classes, bypassing the lazy-loading mechanism for immediate use.

### Usage
```python
from torchebm.integrators import LeapfrogIntegrator
```
```

--------------------------------

### Example Usage of BaseRungeKuttaIntegrator

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_integrator

An example demonstrating how to implement and use a concrete subclass of BaseRungeKuttaIntegrator.

```APIDOC
## Example

```python
from torchebm.core import BaseRungeKuttaIntegrator
import torch

class MidpointIntegrator(BaseRungeKuttaIntegrator):
    @property
    def tableau_a(self):
        return ((), (0.5,))

    @property
    def tableau_b(self):
        return (0.0, 1.0)

    @property
    def tableau_c(self):
        return (0.0, 0.5)

integrator = MidpointIntegrator()
state = {"x": torch.randn(100, 2)}
drift = lambda x, t: -x
result = integrator.step(state, step_size=0.01, drift=drift)
```
```

--------------------------------

### LinearScheduler Warmup Strategy Example

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_scheduler

Shows how LinearScheduler can be used for learning rate warmup.

```APIDOC
## Example: Warmup Strategy
```python
warmup_scheduler = LinearScheduler(
    start_value=0.0, end_value=0.1, n_steps=100
)
# Use for learning rate warmup
for epoch in range(100):
    lr = warmup_scheduler.step()
    # Set learning rate in optimizer
```
```

--------------------------------

### MCMC Sampling with Warmup Example

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core

Illustrates using WarmupScheduler for step size control in MCMC sampling.

```APIDOC
## MCMC Sampling with Warmup

### Description
This example demonstrates how to use WarmupScheduler to manage the step size in Langevin Dynamics for MCMC sampling, starting with a warmup phase.

### Usage
```python
decay_scheduler = ExponentialDecayScheduler(
    start_value=0.05, decay_rate=0.999, min_value=0.001
)
step_scheduler = WarmupScheduler(
    main_scheduler=decay_scheduler,
    warmup_steps=50,
    warmup_init_factor=0.1
)

sampler = LangevinDynamics(
    energy_function=energy_fn,
    step_size=step_scheduler
)
```
```

--------------------------------

### Quick Start: Profile a Training Step

Source: https://soran-ghaderi.github.io/torchebm/latest/developer_guide/profiling

Profile a full training step by specifying an importable zero-argument callable or factory. Use `--trace` and `--top` for detailed analysis.

```bash
1
2
3

```

```bash
python benchmarks/profiler.py run \
    --callable examples.eqm.train_2d:profile_one_step \
    --trace --top 20

```

--------------------------------

### get_start_points

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_loss?q=

Gets the starting points for the MCMC sampler. For standard CD, this is the input data. For PCD, it's samples from the replay buffer.

```APIDOC
## get_start_points(x)

### Description
Gets the starting points for the MCMC sampler. For standard CD, this is the input data. For PCD, it's samples from the replay buffer.

### Method Signature
`get_start_points(self, x: torch.Tensor) -> torch.Tensor`

### Parameters
#### Path Parameters
- **x** (Tensor) - Required - The input data batch.

### Returns
- **Tensor** - torch.Tensor: The tensor of starting points for the sampler.
```

--------------------------------

### Get Starting Points for MCMC Sampler

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core

Determines the starting points for the MCMC sampler. For standard CD, it uses the input data. For PCD, it samples from the replay buffer, potentially with added noise or sampling with replacement if the buffer is smaller than the batch size.

```python
x = x.to(device=self.device, dtype=self.dtype)

batch_size = x.shape[0]
data_shape_no_batch = x.shape[1:]

if self.persistent:
    if not self.buffer_initialized:
        self.initialize_buffer(data_shape_no_batch)
        if not self.buffer_initialized:
            raise RuntimeError("Buffer initialization failed.")

    if self.buffer_size < batch_size:
        warnings.warn(
            f"Buffer size ({self.buffer_size}) is smaller than batch size ({batch_size}). Sampling with replacement.",
            UserWarning,
        )
        indices = torch.randint(
            0, self.buffer_size, (batch_size,), device=self.device
        )
    else:
        # stratified sampling for better buffer coverage
        stride = self.buffer_size // batch_size
        base_indices = torch.arange(0, batch_size, device=self.device) * stride
        offset = torch.randint(0, stride, (batch_size,), device=self.device)
        indices = (base_indices + offset) % self.buffer_size

    start_points = self.replay_buffer[indices]

    # add some noise for exploration
    if self.new_sample_ratio > 0.0:
        n_new = max(1, int(batch_size * self.new_sample_ratio))
        noise_indices = torch.randperm(batch_size, device=self.device)[:n_new]
        noise_scale = 0.01
        start_points[noise_indices] = (
            start_points[noise_indices]
            + torch.randn_like(
                start_points[noise_indices],
                device=self.device,
                dtype=self.dtype,
            )
            * noise_scale
        )
else:
    # standard CD-k uses data as starting points
    start_points = x.detach().clone()

return start_points
```

--------------------------------

### Get Starting Points for Sampler - BaseContrastiveDivergence

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_loss?q=

Retrieves starting points for the MCMC sampler. For standard CD, it uses input data. For PCD, it samples from the replay buffer, potentially adding noise for exploration. Handles buffer initialization and sampling strategies.

```python
def get_start_points(self, x: torch.Tensor) -> torch.Tensor:
        """
        Gets the starting points for the MCMC sampler.

        For standard CD, this is the input data. For PCD, it's samples from the replay buffer.

        Args:
            x (torch.Tensor): The input data batch.

        Returns:
            torch.Tensor: The tensor of starting points for the sampler.
        """
        x = x.to(device=self.device, dtype=self.dtype)

        batch_size = x.shape[0]
        data_shape_no_batch = x.shape[1:]

        if self.persistent:
            if not self.buffer_initialized:
                self.initialize_buffer(data_shape_no_batch)
                if not self.buffer_initialized:
                    raise RuntimeError("Buffer initialization failed.")

            if self.buffer_size < batch_size:
                warnings.warn(
                    f"Buffer size ({self.buffer_size}) is smaller than batch size ({batch_size}). Sampling with replacement.",
                    UserWarning,
                )
                indices = torch.randint(
                    0, self.buffer_size, (batch_size,), device=self.device
                )
            else:
                # stratified sampling for better buffer coverage
                stride = self.buffer_size // batch_size
                base_indices = torch.arange(0, batch_size, device=self.device) * stride
                offset = torch.randint(0, stride, (batch_size,), device=self.device)
                indices = (base_indices + offset) % self.buffer_size

            start_points = self.replay_buffer[indices]

            # add some noise for exploration
            if self.new_sample_ratio > 0.0:
                n_new = max(1, int(batch_size * self.new_sample_ratio))
                noise_indices = torch.randperm(batch_size, device=self.device)[:n_new]
                noise_scale = 0.01
                start_points[noise_indices] = (
                    start_points[noise_indices]
                    + torch.randn_like(
                        start_points[noise_indices],
                        device=self.device,
                        dtype=self.dtype,
                    )
                    * noise_scale
                )
        else:
            # standard CD-k uses data as starting points
            start_points = x.detach().clone()

        return start_points
```

--------------------------------

### Get Starting Points for Sampler

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_loss

Retrieves starting points for the MCMC sampler. For standard CD, it uses the input data. For PCD, it samples from the replay buffer, with options for stratified sampling and adding exploration noise. Handles cases where the buffer size is smaller than the batch size.

```python
    def get_start_points(self, x: torch.Tensor) -> torch.Tensor:
        """
        Gets the starting points for the MCMC sampler.

        For standard CD, this is the input data. For PCD, it's samples from the replay buffer.

        Args:
            x (torch.Tensor): The input data batch.

        Returns:
            torch.Tensor: The tensor of starting points for the sampler.
        """
        x = x.to(device=self.device, dtype=self.dtype)

        batch_size = x.shape[0]
        data_shape_no_batch = x.shape[1:]

        if self.persistent:
            if not self.buffer_initialized:
                self.initialize_buffer(data_shape_no_batch)
                if not self.buffer_initialized:
                    raise RuntimeError("Buffer initialization failed.")

            if self.buffer_size < batch_size:
                warnings.warn(
                    f"Buffer size ({self.buffer_size}) is smaller than batch size ({batch_size}). Sampling with replacement.",
                    UserWarning,
                )
                indices = torch.randint(
                    0, self.buffer_size, (batch_size,), device=self.device
                )
            else:
                # stratified sampling for better buffer coverage
                stride = self.buffer_size // batch_size
                base_indices = torch.arange(0, batch_size, device=self.device) * stride
                offset = torch.randint(0, stride, (batch_size,), device=self.device)
                indices = (base_indices + offset) % self.buffer_size

            start_points = self.replay_buffer[indices]

            # add some noise for exploration
            if self.new_sample_ratio > 0.0:
                n_new = max(1, int(batch_size * self.new_sample_ratio))
                noise_indices = torch.randperm(batch_size, device=self.device)[:n_new]
                noise_scale = 0.01
                start_points[noise_indices] = (
                    start_points[noise_indices]
                    + torch.randn_like(
                        start_points[noise_indices],
                        device=self.device,
                        dtype=self.dtype,
                    )
                    * noise_scale
                )
        else:
            # standard CD-k uses data as starting points
            start_points = x.detach().clone()

        return start_points
```

--------------------------------

### Quick Start: Profile a Component

Source: https://soran-ghaderi.github.io/torchebm/latest/developer_guide/profiling

Profile a specific component from the benchmark suite at a given scale. Use flags like `--trace`, `--top`, and `--memory` to customize the output.

```bash
1
2
3
4

```

```bash
python benchmarks/profiler.py run \
    --component LangevinDynamics \
    --scale medium \
    --trace --top 20 --memory

```

--------------------------------

### Set up Training Components

Source: https://soran-ghaderi.github.io/torchebm/latest/tutorials/getting_started?q=

Initializes the sampler (LangevinDynamics), loss function (ContrastiveDivergence), and optimizer (Adam) for training the EBM. Ensure necessary imports are present.

```python
from torchebm.losses import ContrastiveDivergence
from torchebm.samplers import LangevinDynamics
from torch.optim import Adam

sampler = LangevinDynamics(
    energy_function=neural_energy_fn,
    step_size=10.0,
    noise_scale=0.1,
    n_steps=60
)

cd_loss = ContrastiveDivergence(sampler=sampler)

optimizer = Adam(neural_energy_fn.parameters(), lr=1e-4)
```

--------------------------------

### Train Generative Model with Equilibrium Matching

Source: https://soran-ghaderi.github.io/torchebm/latest

This example demonstrates training a generative model using Equilibrium Matching with a custom VelocityNet and sampling with FlowSampler. Ensure you have PyTorch installed.

```python
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchebm.losses import EquilibriumMatchingLoss
from torchebm.samplers import FlowSampler
from torchebm.core import BaseModel
from torchebm.datasets import EightGaussiansDataset

dataset = EightGaussiansDataset(n_samples=8192)
loader = DataLoader(dataset, batch_size=256, shuffle=True)

# Any nn.Module with forward(x, t) works
class VelocityNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 256), nn.SiLU(),
                                 nn.Linear(256, 256), nn.SiLU(), nn.Linear(256, 2))
    def forward(self, x, t, **kwargs):
        return self.net(x)

model = VelocityNet()

loss_fn = EquilibriumMatchingLoss(
    model=model, interpolant="linear", energy_type="dot",
)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)

for epoch in range(50):
    for x in loader:
        loss = loss_fn(x)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# Sample via ODE flow
flow = FlowSampler(model=model, interpolant="linear", negate_velocity=True)
flow_samples = flow.sample_ode(torch.randn(1000, 2), num_steps=100)

```

--------------------------------

### __init__

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core?q=

Initializes the warmup scheduler, setting up the main scheduler and defining the warmup period.

```APIDOC
## __init__(main_scheduler, warmup_steps, warmup_init_factor=0.01)

### Description
Initialize the warmup scheduler.

### Parameters
#### Path Parameters
- **main_scheduler** (BaseScheduler) - Required - The scheduler to use after warmup.
- **warmup_steps** (int) - Required - Number of warmup steps.
- **warmup_init_factor** (float) - Optional - Factor to determine initial warmup value. The initial value will be main_scheduler.start_value * warmup_init_factor. Defaults to 0.01.
```

--------------------------------

### Implement Custom Sampler with BaseSampler

Source: https://soran-ghaderi.github.io/torchebm/latest/tutorials/samplers?q=

Subclass BaseSampler and implement the sample() method to create a custom sampler. This example shows a basic structure including custom_step and diagnostic setup.

```python
from torchebm.core import BaseSampler, BaseModel
import torch
from typing import Optional, Union, Tuple, List, Dict

class MyCustomSampler(BaseSampler):
    def __init__(
        self, 
        model: BaseModel,
        my_parameter: float = 0.1,
        dtype: torch.dtype = torch.float32,
        device: Optional[Union[str, torch.device]] = None,
    ):
        super().__init__(model=model, dtype=dtype, device=device)
        self.my_parameter = my_parameter

        self.register_scheduler("my_parameter", ConstantScheduler(my_parameter))

    def custom_step(self, x: torch.Tensor) -> torch.Tensor:
        param_value = self.get_scheduled_value("my_parameter")

        gradient = self.model.gradient(x)

        noise = torch.randn_like(x)
        new_x = x - param_value * gradient + noise * 0.01

        return new_x

    @torch.no_grad()
    def sample(
        self,
        x: Optional[torch.Tensor] = None,
        dim: int = 10,
        n_steps: int = 100,
        n_samples: int = 1,
        thin: int = 1,
        return_trajectory: bool = False,
        return_diagnostics: bool = False,
        *args,
        **kwargs,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[dict]]]:
        self.reset_schedulers()

        if x is None:
            x = torch.randn(n_samples, dim, dtype=self.dtype, device=self.device)
        else:
            x = x.to(self.device)

        if return_trajectory:
            trajectory = torch.empty(
                (n_samples, n_steps, dim), dtype=self.dtype, device=self.device
            )

        if return_diagnostics:
            diagnostics = self._setup_diagnostics(dim, n_steps, n_samples=n_samples)

        for i in range(n_steps):
            self.step_schedulers()

            x = self.custom_step(x)

            if return_trajectory:
                trajectory[:, i, :] = x

            if return_diagnostics:
                pass

        if return_trajectory:
            if return_diagnostics:
                return trajectory, diagnostics
            return trajectory
        if return_diagnostics:
            return x, diagnostics
        return x

    def _setup_diagnostics(self, dim: int, n_steps: int, n_samples: int = None) -> torch.Tensor:
        """Optional method to setup diagnostic storage"""
        if n_samples is not None:
            return torch.empty(
                (n_steps, 3, n_samples, dim), device=self.device, dtype=self.dtype
            )
        else:
            return torch.empty((n_steps, 3, dim), device=self.device, dtype=self.dtype)
```

--------------------------------

### sample(x=None, dim=10, n_steps=50, n_samples=1, thin=1, return_trajectory=False, return_diagnostics=False, *, mode='ode', shape=None, ode_method='dopri5', atol=1e-06, rtol=0.001, reverse=False, sde_method='euler', diffusion_form='SBDM', diffusion_norm=1.0, last_step='Mean', last_step_size=0.04, **model_kwargs)

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/samplers/flow?q=

Unified sampling entrypoint for flow/diffusion models. This method provides a convenient way to sample from the model, abstracting away the details of ODE or SDE solvers. For advanced control, it is recommended to use `sample_ode` or `sample_sde` directly.

```APIDOC
## sample

### Description
Unified sampling entrypoint for flow/diffusion models. This method exists for API compatibility with `BaseSampler`. For full control, prefer calling `sample_ode` or `sample_sde` directly.

### Parameters
* **x** (Optional[torch.Tensor]) - Initial tensor to start sampling from. If None, a random tensor is generated.
* **dim** (int) - The dimension of the latent space. Defaults to 10.
* **n_steps** (int) - The number of steps for the sampler. Defaults to 50.
* **n_samples** (int) - The number of samples to generate. Defaults to 1.
* **thin** (int) - Thinning factor for samples. Not supported for FlowSampler.
* **return_trajectory** (bool) - Whether to return the trajectory of samples. Not supported for FlowSampler.
* **return_diagnostics** (bool) - Whether to return diagnostic information. Not supported for FlowSampler.
* **mode** (Literal["ode", "sde"]) - The sampling mode, either 'ode' or 'sde'. Defaults to 'ode'.
* **shape** (Optional[Tuple[int, ...]]) - The shape of the initial tensor if `x` is None. If not provided, it defaults to `(n_samples, dim)`.
* **ode_method** (str) - The ODE solver method to use. Defaults to 'dopri5'.
* **atol** (float) - Absolute tolerance for the ODE solver. Defaults to 1e-6.
* **rtol** (float) - Relative tolerance for the ODE solver. Defaults to 1e-3.
* **reverse** (bool) - Whether to sample in reverse time. Defaults to False.
* **sde_method** (str) - The SDE solver method to use. Defaults to 'euler'.
* **diffusion_form** (str) - The form of the diffusion process for SDE sampling. Defaults to 'SBDM'.
* **diffusion_norm** (float) - Normalization factor for the diffusion process. Defaults to 1.0.
* **last_step** (Optional[str]) - Specifies how to determine the last step in SDE sampling. Defaults to 'Mean'.
* **last_step_size** (float) - The size of the last step in SDE sampling. Defaults to 0.04.
* ****model_kwargs** - Additional keyword arguments to pass to the model.

### Returns
* torch.Tensor - The generated samples.

### Raises
* ValueError - If `thin`, `return_trajectory`, or `return_diagnostics` are used, or if an unknown `mode` is specified.

### Example
```python
# Assuming 'sampler' is an instance of FlowSampler
# Sample using ODE solver
samples_ode = sampler.sample(n_samples=5, mode='ode', n_steps=100)

# Sample using SDE solver
samples_sde = sampler.sample(n_samples=5, mode='sde', n_steps=100, sde_method='milstein')
```
```

--------------------------------

### LinearScheduler Example: Warmup Strategy

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_scheduler

Illustrates using LinearScheduler for a learning rate warmup strategy, where the parameter increases linearly from a start value to an end value over a set number of steps.

```python
warmup_scheduler = LinearScheduler(
    start_value=0.0, end_value=0.1, n_steps=100
)
# Use for learning rate warmup
for epoch in range(100):
    lr = warmup_scheduler.step()
    # Set learning rate in optimizer
```

--------------------------------

### Initialize LangevinDynamics Sampler

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/samplers/langevin_dynamics

Instantiate the LangevinDynamics sampler with an energy model and sampling parameters. Use this to set up the sampler for generating samples.

```python
from torchebm.samplers import LangevinDynamics
from torchebm.core import DoubleWellEnergy
import torch

energy = DoubleWellEnergy()
sampler = LangevinDynamics(energy, step_size=0.01, noise_scale=1.0)
samples = sampler.sample(n_samples=100, dim=2, n_steps=500)
```

--------------------------------

### LinearScheduler Example: Linear Decay

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_scheduler

Demonstrates how LinearScheduler creates a linear decay from a start value to an end value over a specified number of steps. The output is clamped to the end value after n_steps.

```python
scheduler = LinearScheduler(start_value=1.0, end_value=0.0, n_steps=5)
for i in range(7):  # Go beyond n_steps to see clamping
    value = scheduler.step()
    print(f"Step {i+1}: {value:.2f}")
```

--------------------------------

### CosineScheduler Learning Rate Annealing

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_scheduler

Provides an example of using CosineScheduler for learning rate annealing. This setup defines a scheduler that will smoothly decrease the learning rate over a specified number of steps.

```python
lr_scheduler = CosineScheduler(
    start_value=0.01, end_value=0.0001, n_steps=1000
)
```

--------------------------------

### Get Starting Points for MCMC Sampler

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_loss?q=

This method retrieves the initial data points for the MCMC sampler. For standard CD, it uses the input data directly. For PCD, it samples from a replay buffer, potentially adding noise for exploration. Ensure the replay buffer is initialized if `persistent` is True.

```python
def get_start_points(self, x: torch.Tensor) -> torch.Tensor:
    """
    Gets the starting points for the MCMC sampler.

    For standard CD, this is the input data. For PCD, it's samples from the replay buffer.

    Args:
        x (torch.Tensor): The input data batch.

    Returns:
        torch.Tensor: The tensor of starting points for the sampler.
    """
    x = x.to(device=self.device, dtype=self.dtype)

    batch_size = x.shape[0]
    data_shape_no_batch = x.shape[1:]

    if self.persistent:
        if not self.buffer_initialized:
            self.initialize_buffer(data_shape_no_batch)
            if not self.buffer_initialized:
                raise RuntimeError("Buffer initialization failed.")

        if self.buffer_size < batch_size:
            warnings.warn(
                f"Buffer size ({self.buffer_size}) is smaller than batch size ({batch_size}). Sampling with replacement.",
                UserWarning,
            )
            indices = torch.randint(
                0, self.buffer_size, (batch_size,), device=self.device
            )
        else:
            # stratified sampling for better buffer coverage
            stride = self.buffer_size // batch_size
            base_indices = torch.arange(0, batch_size, device=self.device) * stride
            offset = torch.randint(0, stride, (batch_size,), device=self.device)
            indices = (base_indices + offset) % self.buffer_size

        start_points = self.replay_buffer[indices]

        # add some noise for exploration
        if self.new_sample_ratio > 0.0:
            n_new = max(1, int(batch_size * self.new_sample_ratio))
            noise_indices = torch.randperm(batch_size, device=self.device)[:n_new]
            noise_scale = 0.01
            start_points[noise_indices] = (
                start_points[noise_indices]
                + torch.randn_like(
                    start_points[noise_indices],
                    device=self.device,
                    dtype=self.dtype,
                )
                * noise_scale
            )
    else:
        # standard CD-k uses data as starting points
        start_points = x.detach().clone()

    return start_points
```

--------------------------------

### Quick Start: Profile a Training Step

Source: https://soran-ghaderi.github.io/torchebm/latest/developer_guide/profiling?q=

Profile a full training step by specifying an importable zero-argument callable or factory. This is useful for analyzing the performance of training loops.

```bash
python benchmarks/profiler.py run \
    --callable examples.eqm.train_2d:profile_one_step \
    --trace --top 20

```

--------------------------------

### Complete Training Example with Contrastive Divergence Loss

Source: https://soran-ghaderi.github.io/torchebm/latest/tutorials/loss_functions

Demonstrates a full training loop for an EBM using Contrastive Divergence loss. Includes model definition, data loading, sampler setup, loss function instantiation, optimization, and training execution. Gradient clipping is applied for stability.

```python
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader

from torchebm.core import BaseModel
from torchebm.samplers import LangevinDynamics
from torchebm.losses import ContrastiveDivergence
from torchebm.datasets import TwoMoonsDataset

class MLPModel(BaseModel):
    def __init__(self, input_dim, hidden_dim=64):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SELU(),
            nn.Linear(hidden_dim, 1),
            nn.Tanh(),
        )

    def forward(self, x):
        return self.network(x).squeeze(-1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

INPUT_DIM = 2
HIDDEN_DIM = 16
BATCH_SIZE = 256
EPOCHS = 100
LEARNING_RATE = 1e-3
CD_K = 10
USE_PCD = True

dataset = TwoMoonsDataset(n_samples=3000, noise=0.05, seed=42, device=device)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)

model = MLPModel(INPUT_DIM, HIDDEN_DIM).to(device)
sampler = LangevinDynamics(
    model=model,
    step_size=0.1,
    device=device,
)
loss_fn = ContrastiveDivergence(
    model=model,
    sampler=sampler,
    k_steps=CD_K,
    persistent=USE_PCD,
    buffer_size=BATCH_SIZE
).to(device)

optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

losses = []
print("Starting training...")
for epoch in range(EPOCHS):
    model.train()
    epoch_loss = 0.0

    for i, data_batch in enumerate(dataloader):
        optimizer.zero_grad()

        loss, negative_samples = loss_fn(data_batch)

        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()

        epoch_loss += loss.item()

    avg_epoch_loss = epoch_loss / len(dataloader)
    losses.append(avg_epoch_loss)
    print(f"Epoch [{epoch+1}/{EPOCHS}], Average Loss: {avg_epoch_loss:.4f}")

plt.figure(figsize=(10, 6))
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('docs/assets/images/loss_functions/cd_training_loss.png')
plt.show()
```

--------------------------------

### Install TorchEBM with pip

Source: https://soran-ghaderi.github.io/torchebm/latest/faq

Install TorchEBM using pip. Ensure you have a compatible PyTorch version and CUDA toolkit installed if GPU support is desired.

```bash
pip install torchebm
```

--------------------------------

### EquilibriumMatchingLoss Initialization and Usage

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/losses/equilibrium_matching?q=

Demonstrates initializing the EquilibriumMatchingLoss for both implicit and explicit EqM with different energy types. Shows basic usage with a dummy model and data.

```python
from torchebm.losses import EquilibriumMatchingLoss
import torch.nn as nn
import torch

# Implicit EqM with velocity prediction (default)
model = MyTimeConditionedModel()
loss_fn = EquilibriumMatchingLoss(
    model=model,
    prediction="velocity",
    energy_type="none",
)

# Explicit EqM-E with dot product (for OOD detection)
loss_fn_explicit = EquilibriumMatchingLoss(
    model=model,
    prediction="velocity",
    energy_type="dot",
)

x = torch.randn(32, 2)
loss = loss_fn(x)
```

--------------------------------

### __init__

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_scheduler

Initializes the warmup scheduler with a main scheduler, warmup steps, and an optional warmup factor.

```APIDOC
## __init__

### Description
Initialize the warmup scheduler.

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Method Signature
`__init__(self, main_scheduler: BaseScheduler, warmup_steps: int, warmup_init_factor: float = 0.01)`

### Parameters
- **main_scheduler** (BaseScheduler) - Required - The scheduler to use after warmup.
- **warmup_steps** (int) - Required - Number of warmup steps.
- **warmup_init_factor** (float) - Optional - Factor to determine initial warmup value. The initial value will be main_scheduler.start_value * warmup_init_factor. Defaults to 0.01.
```

--------------------------------

### Adaptive Integration Setup and Execution

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core?q=

Sets up and executes adaptive integration. Raises a ValueError if error_weights or order are not defined, as these are required for adaptive integration.

```python
# adaptive path
        if self.error_weights is None or self.order is None:
            raise ValueError(
                f"{type(self).__name__} does not define error_weights/order "
                f"and cannot be used with adaptive=True."
            )

        x = state["x"]
        drift_fn = self._resolve_drift(drift)

        if not torch.is_tensor(step_size):
            step_size = torch.tensor(
                step_size, device=x.device, dtype=x.dtype
            )

        if t is not None:
            if t.ndim != 1 or t.numel() < 2:
                raise ValueError("t must be a 1D tensor with length >= 2")
            t_start = t[0].item()
            t_end = t[-1].item()
        else:
            t_start = 0.0
            t_end = float(n_steps) * step_size.item()

        h = min(step_size.item(), t_end - t_start, self.max_step_size)
        x = self._adaptive_integrate(x, drift_fn, t_start, t_end, h)
        return {"x": x}
```

--------------------------------

### Learning Rate Warmup + Cosine Annealing Example

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core

Demonstrates how to use WarmupScheduler for learning rate warmup combined with cosine annealing.

```APIDOC
## Learning Rate Warmup + Cosine Annealing

### Description
This example shows how to set up a learning rate schedule that first warms up linearly and then follows a cosine annealing pattern.

### Usage
```python
main_scheduler = CosineScheduler(
    start_value=0.1, end_value=0.001, n_steps=1000
)
warmup_scheduler = WarmupScheduler(
    main_scheduler=main_scheduler,
    warmup_steps=100,
    warmup_init_factor=0.01
)

# First 100 steps: linear warmup from 0.001 to 0.1
# Next 1000 steps: cosine annealing from 0.1 to 0.001
for i in range(10):
    value = warmup_scheduler.step()
    print(f"Warmup step {i+1}: {value:.6f}")
```
```

--------------------------------

### Bosh3Integrator Initialization and Usage

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/integrators/bosh3

Demonstrates how to initialize the Bosh3Integrator with custom tolerances and use it to integrate a system.

```APIDOC
## Bosh3Integrator

### Description
Initializes the Bogacki-Shampine 3(2) explicit Runge-Kutta integrator.

### Parameters
#### Constructor Parameters
- **atol** (`float`) - Absolute tolerance for adaptive stepping. Default: `1e-06`
- **rtol** (`float`) - Relative tolerance for adaptive stepping. Default: `0.001`
- **max_steps** (`int`) - Maximum number of accepted steps before raising. Default: `10000`
- **safety** (`float`) - Safety factor for step-size adjustment (< 1). Default: `0.9`
- **min_factor** (`float`) - Minimum step-size shrink factor. Default: `0.2`
- **max_factor** (`float`) - Maximum step-size growth factor. Default: `10.0`
- **max_step_size** (`float`) - Maximum absolute step size during adaptive integration. Default: `float('inf')`
- **norm** (`Optional[Callable[[Tensor], Tensor]]`) - Callable `norm(tensor) -> scalar` for local error measurement. Default: `None`
- **device** (`Optional[device]`) - Device for computations. Default: `None`
- **dtype** (`Optional[dtype]`) - Data type for computations. Default: `None`

### Method
`integrate(state, step_size, n_steps, drift, **kwargs)`

### Parameters
#### `integrate` Method Parameters
- **state** (`dict`) - The initial state of the system.
- **step_size** (`float`) - The initial step size for integration.
- **n_steps** (`int`) - The number of steps to take.
- **drift** (`Callable[[Tensor, Tensor], Tensor]`) - The drift function `f(x, t)`.
- **adaptive** (`bool`) - Whether to use adaptive step sizing. Default: `True` (if `error_weights` is defined).

### Request Example
```python
from torchebm.integrators import Bosh3Integrator
import torch

integrator = Bosh3Integrator(atol=1e-4, rtol=1e-2)
state = {"x": torch.randn(100, 2)}
drift = lambda x, t: -x
result = integrator.integrate(
    state, step_size=0.1, n_steps=10, drift=drift,
)
```
```

--------------------------------

### MidpointIntegrator Example

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core?q=

Example of a concrete implementation of BaseRungeKuttaIntegrator using the Midpoint method.

```APIDOC
## MidpointIntegrator

### Description
A concrete implementation of `BaseRungeKuttaIntegrator` using the Midpoint method.

### Usage
```python
from torchebm.core import BaseRungeKuttaIntegrator
import torch

class MidpointIntegrator(BaseRungeKuttaIntegrator):
    @property
    def tableau_a(self):
        return ((), (0.5,))

    @property
    def tableau_b(self):
        return (0.0, 1.0)

    @property
    def tableau_c(self):
        return (0.0, 0.5)

integrator = MidpointIntegrator()
state = {"x": torch.randn(100, 2)}
drift = lambda x, t: -x
result = integrator.step(state, step_size=0.01, drift=drift)
```

### Parameters

- `device` (Optional[device]) - Optional - Device for computations.
- `dtype` (Optional[dtype]) - Optional - Data type for computations.
- `atol` (float) - Optional - Absolute tolerance for adaptive stepping. Default: `1e-06`
- `rtol` (float) - Optional - Relative tolerance for adaptive stepping. Default: `0.001`
- `max_steps` (int) - Optional - Maximum number of steps (accepted + rejected) before raising `RuntimeError`. Default: `10000`
- `safety` (float) - Optional - Safety factor for step-size adjustment (< 1). Default: `0.9`
- `min_factor` (float) - Optional - Minimum step-size shrink factor. Default: `0.2`
- `max_factor` (float) - Optional - Maximum step-size growth factor. Default: `10.0`
- `max_step_size` (float) - Optional - Maximum absolute step size allowed during adaptive integration. Defaults to `inf` (no limit).
- `norm` (Optional[Callable[[Tensor], Tensor]]) - Optional - Callable `norm(tensor) -> scalar` used to measure the local error. Defaults to the RMS norm mean(x2).

### Abstract Properties to Implement

- `tableau_a`: Butcher tableau coefficients for the 'a' matrix.
- `tableau_b`: Butcher tableau coefficients for the 'b' vector.
- `tableau_c`: Butcher tableau coefficients for the 'c' vector.

### Adaptive Step-Size Control

Adaptive step-size control is available automatically for subclasses that define `error_weights` and `order`. When `adaptive=True` is passed to `integrate` (or left as `None` for auto-detection), the integrator uses an embedded error pair to control the step size.
```

--------------------------------

### Quick Start: Profile a Component

Source: https://soran-ghaderi.github.io/torchebm/latest/developer_guide/profiling?q=

Profile a specific component at a scale defined by the benchmark suite. This command includes options for tracing, displaying top operations, and memory profiling.

```bash
python benchmarks/profiler.py run \
    --component LangevinDynamics \
    --scale medium \
    --trace --top 20 --memory

```

--------------------------------

### Use Langevin Dynamics Sampler

Source: https://soran-ghaderi.github.io/torchebm/latest/tutorials

This example demonstrates the usage of the LangevinDynamics sampler. It requires an initialized energy_function and specifies parameters like step_size, n_steps, and n_samples.

```python
from torchebm.samplers.langevin_dynamics import LangevinDynamics

sampler = LangevinDynamics(
    energy_function=energy_fn,
    step_size=0.01
)

samples = sampler.sample(
    dim=2, n_steps=100, n_samples=1000
)
```

--------------------------------

### Get and Set Dtype Property

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/device_mixin?q=

Gets the dtype, inferring from parameters or buffers if not set. Sets the dtype.

```python
@property
def dtype(self) -> torch.dtype:
        if self._dtype is not None:
            return self._dtype
        params = getattr(self, "parameters", None)
        if callable(params):
            try:
                p_dtype = next(self.parameters()).dtype
                self._dtype = p_dtype
                return p_dtype
            except StopIteration:
                pass
        bufs = getattr(self, "buffers", None)
        if callable(bufs):
            try:
                b_dtype = next(self.buffers()).dtype
                self._dtype = b_dtype
                return b_dtype
            except StopIteration:
                pass
        return torch.float32

    @dtype.setter
def dtype(self, value: torch.dtype):
        self._dtype = value
```

--------------------------------

### Create and Sample from a 2D Gaussian Energy Model

Source: https://soran-ghaderi.github.io/torchebm/latest/tutorials

This example demonstrates how to define a simple 2D Gaussian energy model and use Langevin Dynamics to generate samples from it. Ensure you have the necessary imports for torch, GaussianModel, and LangevinDynamics.

```python
import torch
from torchebm.core import GaussianModel
from torchebm.samplers.langevin_dynamics import LangevinDynamics

# Create an energy model (2D Gaussian)
energy_fn = GaussianModel(
    mean=torch.zeros(2),
    cov=torch.eye(2)
)

# Create a sampler
sampler = LangevinDynamics(
    energy_function=energy_fn,
    step_size=0.01
)

# Generate samples
samples = sampler.sample(
    dim=2, n_steps=100, n_samples=1000
)

# Print sample statistics
print(f"Sample mean: {samples.mean(0)}")
print(f"Sample std: {samples.std(0)}")
```

--------------------------------

### Initialize Base Scheduler

Source: https://soran-ghaderi.github.io/torchebm/latest/api/torchebm/core/base_scheduler

Initializes the scheduler with a starting parameter value. Ensures the start value is a finite float or integer.

```python
def __init__(self, start_value: float):
    r"""
    Initialize the base scheduler.

    Args:
        start_value (float): Initial parameter value. Must be a finite number.

    Raises:
        TypeError: If start_value is not a float or int.
    """
    if not isinstance(start_value, (float, int)):
        raise TypeError(
            f"{type(self).__name__} received an invalid start_value of type "
            f"{type(start_value).__name__}. Expected float or int."
        )

    self.start_value = float(start_value)
    self.current_value = self.start_value
    self.step_count = 0
```
Tokens:
Raw