Project Structure¶
Codebase Organization
Understanding the TorchEBM project structure helps you navigate the codebase and contribute effectively. This guide provides an overview of the repository organization.
Repository Overview¶
The TorchEBM repository is organized as follows:
Main Package Structure¶
torchebm.core
¶
Contains the core functionality including base classes and essential components:
base.py
- Base classes for the entire libraryenergy_function.py
- Base energy function class and interfaceanalytical_functions.py
- Analytical energy function implementationsdistributions.py
- Probability distribution implementations
torchebm.samplers
¶
Implementations of various sampling algorithms:
base.py
- Base sampler classlangevin_dynamics.py
- Langevin dynamics implementationhmc.py
- Hamiltonian Monte Carlometropolis_hastings.py
- Metropolis-Hastings- [Other sampler implementations]
torchebm.losses
¶
BaseLoss functions for training energy-based models:
base.py
- Base loss classcontrastive_divergence.py
- Contrastive divergence implementationsscore_matching.py
- Score matching methods- [Other loss implementations]
torchebm.models
¶
Neural network model implementations:
mlp.py
- Multi-layer perceptron energy modelscnn.py
- Convolutional neural network energy modelsebm.py
- Generic energy-based model implementations- [Other model architectures]
torchebm.cuda
¶
CUDA-optimized implementations for performance-critical operations:
kernels/
- CUDA kernel implementationsbindings.cpp
- PyTorch C++ bindingsops.py
- Python interfaces to CUDA operations
torchebm.utils
¶
Utility functions and helpers:
device.py
- Device management utilitiesvisualization.py
- Visualization toolsdata.py
- Data loading and processing utilitieslogging.py
- Logging utilities
Tests Structure¶
The tests directory mirrors the package structure:
Documentation Structure¶
The documentation is built with MkDocs and organized as follows:
Examples Structure¶
Example applications to demonstrate TorchEBM usage:
Dependencies and Requirements¶
TorchEBM has the following dependencies:
-
PyTorch
Primary framework for tensor operations and automatic differentiation.
-
NumPy
Used for numerical operations when PyTorch isn't needed.
-
Matplotlib
For visualization capabilities.
-
tqdm
For progress bars during long operations.
Entry Points¶
The main entry points to the library are:
torchebm.core
- Import core functionalitytorchebm.samplers
- Import samplerstorchebm.losses
- Import loss functionstorchebm.models
- Import neural network models
Example of typical import patterns:
Package Management¶
TorchEBM uses setuptools for package management with setup.py:
Version Control Structure¶
- We follow Conventional Commits for our commit messages
- Feature branches should be named
feature/feature-name
- Bugfix branches should be named
bugfix/bug-name
- Release branches should be named
release/vX.Y.Z
Finding Your Way Around
When contributing to TorchEBM, start by exploring the relevant directory for your feature. For example, if you're adding a new sampler, look at the existing implementations in torchebm/samplers/
to understand the pattern.