ML Research Template

A Copier template for machine learning research projects using PyTorch Lightning + Hydra + W&B

日本語版 README

Features

🔥 Modern ML Stack

• PyTorch Lightning 2.x: High-level training framework
• Hydra 1.3+: Powerful configuration management
• W&B/TensorBoard/MLflow: Experiment tracking
• PyTorch Geometric: Graph Neural Networks (optional)

⚡ Fast Development Environment

• uv: 10-100x faster Python package manager than pip
• pixi: Conda alternative with isolated CUDA environment (no host interference)
• Unified pyproject.toml: Single configuration file for both package managers
• ruff: Fast linter & formatter
• ty: Fast type checker (by Astral)
• pytest: Testing framework

🎯 9 PyTorch + CUDA Presets

PyTorch	CUDA Versions
2.9.0	12.6, 13.0
2.8.0	12.6, 12.8
2.7.1	11.8, 12.6
2.6.0	12.4
2.5.1	12.1
2.4.1	11.8

Quick Start

Prerequisites

# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Or install pixi
curl -fsSL https://pixi.sh/install.sh | bash

Generate Project

# From GitHub (recommended)
uvx copier copy --trust gh:nishide-dev/ml-research-template my-project

# From local clone
git clone https://github.com/nishide-dev/ml-research-template.git
uvx copier copy --trust ./ml-research-template my-project

# Navigate to your project
cd my-project

# Start training
uv run python src/my_project/train.py  # if using uv
pixi run train                           # if using pixi

Tip: Auto-fill author information from Git config:

uvx copier copy --trust gh:nishide-dev/ml-research-template my-project \
  --data author_name="$(git config user.name)" \
  --data author_email="$(git config user.email)"

Configuration Options

Interactive prompts will ask for:

Project Basics:

• Project name, package name, description
• Author name, email
• Python version (3.10, 3.11, 3.12, 3.13)

Package Manager:

• uv: Fast, pip-compatible, PyTorch with bundled CUDA
• pixi: Conda-based, isolated CUDA environment (.pixi/envs/)

PyTorch/CUDA Configuration:

• 9 presets (PyTorch 2.4-2.9, CUDA 11.8-13.0)
• Custom configuration available
• Optional torchvision / torchaudio

ML Frameworks:

• PyTorch Lightning (2.2, 2.3, 2.4)
• Hydra (1.2, 1.3)
• PyTorch Geometric (for GNNs)

Experiment Tracking:

• TensorBoard (simple)
• Weights & Biases (feature-rich)
• MLflow (on-premise)
• Both (TensorBoard + W&B)
• None

Development Tools:

• ruff (linter/formatter)
• ty (type checker)
• pytest (testing)

Generated Project Structure

my-project/
├── src/
│   └── my_project/
│       ├── __init__.py
│       ├── train.py              # Main training script (Hydra)
│       ├── models/
│       │   ├── __init__.py
│       │   └── base_model.py     # LightningModule
│       ├── data/
│       │   ├── __init__.py
│       │   └── datamodule.py     # LightningDataModule
│       └── utils/
│           └── __init__.py
├── tests/
│   ├── __init__.py
│   └── test_my_project.py
├── configs/                      # Hydra configuration
│   ├── config.yaml               # Main config
│   ├── model/
│   │   └── default.yaml
│   ├── data/
│   │   └── default.yaml
│   ├── trainer/
│   │   └── default.yaml
│   ├── logger/
│   │   ├── tensorboard.yaml
│   │   ├── wandb.yaml
│   │   └── mlflow.yaml
│   └── experiment/
│       └── baseline.yaml
├── pyproject.toml                # Unified config (uv + pixi)
├── ruff.toml                     # ruff config
├── .gitignore
├── .python-version
└── README.md

Usage Examples

Basic Training

# Train with default config
uv run python src/my_project/train.py

# Override parameters
uv run python src/my_project/train.py trainer.max_epochs=50 data.batch_size=64

# Use specific experiment config
uv run python src/my_project/train.py experiment=baseline

Hydra Configuration Management

# Override multiple parameters
uv run python src/my_project/train.py \
  trainer.max_epochs=100 \
  data.batch_size=128 \
  model.learning_rate=0.001 \
  logger.name=my_experiment

# Multi-run sweep (hyperparameter search)
uv run python src/my_project/train.py -m \
  data.batch_size=32,64,128 \
  model.learning_rate=0.0001,0.001,0.01

Experiment Tracking

# TensorBoard
tensorboard --logdir logs/

# Weights & Biases (login required)
wandb login
uv run python src/my_project/train.py

# MLflow
mlflow ui --backend-store-uri file:./mlruns

Code Quality Checks

# Format and lint
uv run ruff format .
uv run ruff check . --fix

# Type check
uv run ty check src/

# Run tests
uv run pytest tests/ -v

# Run tests with coverage
uv run pytest --cov=src --cov-report=html

Development

Testing the Template

# Test with defaults
uvx copier copy --trust --defaults . /tmp/test-project

# Test with custom config
uvx copier copy --trust \
  --data project_name="test-ml" \
  --data python_version="3.12" \
  --data package_manager="uv" \
  --data pytorch_cuda_preset="pytorch-2.8.0-cuda-12.6" \
  --data use_lightning=true \
  --data use_hydra=true \
  --data logger_choice="tensorboard" \
  . /tmp/test-project

# Verify the generated project
cd /tmp/test-project
uv sync
uv run pytest tests/
uv run ruff check .

Updating the Template

# Update existing project with latest template
uvx copier update /path/to/my-project

Key Features

Unified pyproject.toml

Both uv and pixi package managers use a single pyproject.toml file:

• UV users: Uses [tool.uv] sections with PyTorch wheel indexes (bundled CUDA)
• Pixi users: Uses [tool.pixi.*] sections with conda CUDA toolkit (isolated in .pixi/envs/)

CUDA Isolation (Pixi)

Pixi projects provide complete CUDA isolation from the host system:

[tool.pixi.activation]
env = { PYTHONNOUSERSITE = "1" }  # Isolate from ~/.local

[tool.pixi.dependencies]
cuda = { version = "12.6.*", channel = "nvidia" }  # Isolated CUDA

Benefits:

• No interference with host CUDA installations
• Reproducible GPU environments across machines
• Multiple CUDA versions on same machine (different projects)

Verification:

pixi run python -c "import torch; print(torch.__file__)"
# Output: /path/to/project/.pixi/envs/default/lib/python3.11/site-packages/torch/

Related Projects

• uv: Fast Python package manager
• pixi: Fast Conda alternative
• PyTorch Lightning: PyTorch training framework
• Hydra: Configuration management framework

Contributing

Contributions are welcome! See CONTRIBUTING.md for details.

License

MIT License - see LICENSE for details.

Support

• Issues: GitHub Issues
• Discussions: GitHub Discussions

---

Made with ❤️ for ML researchers