A comprehensive PyTorch-based toolkit for managing machine learning experiments. Streamlines experiment tracking, metrics monitoring, and training workflows.

• Experiment Management

  tracker = RunTracker("experiments")
  run_id = tracker.create_run(description="Training run", hparams={"lr": 0.001})

  • Automatic experiment tracking and metadata storage
  • Checkpoint management with best model saving

• Training Pipeline

  trainer = BaseTrainer(
      model=model,
      optimizer=optimizer,
      train_dataloader=train_loader,
      test_dataloader=test_loader,
      tracker=tracker,
      hyperparams={"num_epochs": 10, "device": "cuda"}
  )
  metrics = trainer.train()

  • Flexible base trainer with customizable training loops
  • Progress tracking with configurable intervals
  • Gradient monitoring and clipping
  • Learning rate scheduling

• Time Series Processing

  splitter = TimeWindowSplitter(data, index_to_info, window_size=24)
  windowed_data, new_info = splitter.split()

  • Window-based sequence splitting
  • Dataset partitioning with stratification
  • Balanced splitting for imbalanced data

• Visualization

  plotter = MatplotlibMetricsPlotter(train_metrics, val_metrics)
  plotter.plot_task_metrics("classification", save="metrics.png")

  • Real-time training metrics visualization
  • Support for both Matplotlib and Plotly
  • Customizable plotting options

pip install ml-training-utils

For development:

git clone https://github.com/yourusername/ml-training-utils
cd ml-training-utils
pip install -e ".[dev]"

from ml_training_utils import RunTracker, BaseTrainer
import torch.nn as nn

# Initialize tracking
tracker = RunTracker("experiments")

# Define model and optimizer
model = YourModel()
optimizer = torch.optim.Adam(model.parameters())

# Create trainer
trainer = BaseTrainer(
    model=model,
    optimizer=optimizer,
    train_dataloader=train_loader,
    test_dataloader=test_loader,
    tracker=tracker,
    hyperparams={
        "num_epochs": 10,
        "device": "cuda",
        "max_norm": 1.0,
        "print_percentage": 0.1
    }
)

# Train with automatic tracking
metrics_train, metrics_test = trainer.train()

1. Fork the repository
2. Create your feature branch (git checkout -b feature/amazing-feature)
3. Commit your changes (git commit -m 'Add amazing feature')
4. Push to the branch (git push origin feature/amazing-feature)
5. Open a Pull Request

This project is licensed under the MIT License - see the LICENSE file for details.

If you use this package in your research, please cite:

@software{ml_training_utils,
  author = {Sagi Nathan},
  title = {ML Training Utils: A PyTorch Experiment Management Toolkit},
  year = {2024},
  publisher = {GitHub},
  url = {https://github.com/saginat/ml-training-utils}
}