High-Level library for Efficient and Fast Architecture Evaluation and Comparison of Machine Learning models. It provides a unified interface for training, benchmarking, and hyperparameter optimization with features like distributed training, mixed precision, and real-time visualization.
- Unified Training Interface: Train single models with easy-to-use configuration options.
- Multi-Model Benchmarking: Compare multiple architectures sequentially or in parallel (thread/process-based).
- Distributed Training: Built-in support for DataParallel, DistributedDataParallel (DDP), and FSDP.
- Advanced Mixed Precision: AMP with float16, bfloat16, and experimental FP8 support.
- Gradient Checkpointing: Reduce memory footprint for large models.
- Rich Visualization: Real-time training windows, video recording of metrics, and publication-ready plots.
- Logging: Integration with Weights & Biases and TensorBoard.
- Hyperparameter Optimization: Grid search and random search out of the box.
- Extensible Plugin System: Custom hooks and callbacks for maximum flexibility.
- Data Handling: Supports PyTorch Datasets, synthetic data, torchvision datasets, Hugging Face datasets, and streaming.
- Transformer Support: Seamless compatibility with Hugging Face Transformers and custom transformer architectures.
- Production-Ready: Configurable timeouts, retry logic, checkpointing, and deterministic execution.
Install from the GitHub repository:
# Clone the repository
git clone --depth=1 https://github.com/lof310/arch_eval.git
cd arch_eval
# Install in development mode (recommended)
pip install -e .# Install normally
pip install .importtorch.nnasnnfromarch_evalimportTrainer, TrainingConfig, discover_plugins# Initialize plugins (must be called before using plugins)discover_plugins()
# Define a global configuration# Datasetn_samples, n_features, n_classes=5000, 128, 64# Modelinput_size, hidden=n_features, n_features*2# Trainingbatch_size, num_epochs=16, 4# Define a simple modelclassMLP(nn.Module):
def__init__(self, input_size=128, hidden=256, num_classes=64):
super().__init__()
self.net=nn.Sequential(
nn.Linear(input_size, hidden),
nn.GELU(),
nn.Linear(hidden, num_classes)
)
defforward(self, x):
returnself.net(x)
# Configure trainingconfig=TrainingConfig(
dataset="synthetic classification",
dataset_params={"n_samples": n_samples, "n_features": n_features, "n_classes": n_classes},
training_args={"num_epochs": num_epochs, "batch_size": batch_size},
task="classification",
realtime="auto", # Options: "auto", "gui", "terminal", "none"save_plot=["loss", "accuracy"]
)
model=MLP(input_size, hidden, n_classes)
trainer=Trainer(model, config)
history=trainer.train()fromarch_evalimportBenchmark, BenchmarkConfigmodels= [
{"name": "Small MLP", "model": MLP(hidden=256)},
{"name": "Large MLP", "model": MLP(hidden=512)}
]
config=BenchmarkConfig(
dataset="synthetic classification",
dataset_params={"n_samples": 10000, "n_features": 128, "n_classes": 64},
compare_metrics=["accuracy", "loss"],
parallel=True
)
benchmark=Benchmark(models, config)
results=benchmark.run()
print(results)fromarch_evalimportHyperparameterOptimizerdefmodel_fn():
returnMLP()
base_config=TrainingConfig(
dataset="synthetic classification",
dataset_params={"n_samples": 1000, "n_features": 128, "n_classes": 64},
training_args={"num_epochs": 3},
task="classification",
realtime="none"# disable live plots during search
)
param_grid= {
"learning_rate": [0.001, 0.01, 0.1],
"hidden": [10, 20, 50]
}
optimizer=HyperparameterOptimizer(
model_fn, base_config, param_grid,
search_type="grid", metric="val_accuracy", mode="max"
)
results=optimizer.run()Full documentation is available at https://lof310.github.io/arch_eval
The documentation includes:
- Quick Start Guide: Get up and running in minutes
- User Guide: Walkthrough of all features
- API Reference: Detailed documentation of all classes and functions
- Examples: Complete, runnable code examples
To build the documentation locally:
cd docs
pip install -r requirements-docs.txt
make html
# Open docs/build/html/index.html in your browser# you can usecd docs/build/html
python -m http.serverContributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.
To contribute:
- Fork the repository
- Create a feature branch (
git checkout -b feature/amazing-feature) - Commit your changes (
git commit -m 'Add some amazing feature') - Push to the branch (
git push origin feature/amazing-feature) - Open a Pull Request
Distributed under the Apache License 2.0. See LICENSE for more information.
If you use arch_eval in your research, please cite:
@software{arch_eval2026,
author = {Leinier Orama},
title = {arch_eval: High-level Library for Architecture Evaluation of ML Models},
year = {2026},
publisher = {GitHub},
url = {https://github.com/lof310/arch_eval}
}