Skip to content

nishide-dev/claude-code-ml-research

BranchesTags

Repository files navigation

ML Research Plugin for Claude Code

Comprehensive Claude Code plugin for machine learning research and experimentation. Provides project scaffolding, experiment management, training support, and debugging tools for PyTorch Lightning, Hydra, PyTorch Geometric, and Hugging Face Transformers.

Installation

Via Claude Code Marketplace

# Add marketplace
/plugin marketplace add nishide-dev/claude-code-ml-research

# Install plugin
/plugin install ml-research

Via GitHub URL

/plugin install gh:nishide-dev/claude-code-ml-research

Via Local Clone

git clone https://github.com/nishide-dev/claude-code-ml-research.git
/plugin install ./claude-code-ml-research

Features

Skills

This plugin provides comprehensive ML workflow skills following the Agent Skills open standard. All skills are directory-based with supporting files (templates, examples, scripts).

Workflow Skills (user-invocable commands):

  • /ml-train - Execute training runs with proper monitoring, checkpointing, and experiment tracking
  • /ml-config-manager - Generate and manage Hydra configuration files (model, data, trainer, logger, experiment configs)
  • /ml-debug - Debug common ML training issues (NaN loss, OOM, slow training, convergence problems)
  • /ml-experiment - Manage ML experiments, track results, and compare performance across configurations
  • /ml-validate - Comprehensive validation of ML project structure, configurations, and training readiness
  • /ml-profile - Profile ML training performance to identify bottlenecks and optimize GPU utilization
  • /ml-data-pipeline - Create and manage data loading, preprocessing, and augmentation pipelines
  • /ml-setup - Setup development environment with modern Python tooling (uv/pixi)
  • /ml-project-init - Initialize new ML research project using the ML Research template
  • /ml-lint - Run comprehensive code quality checks (ruff format/check, type checking)
  • /ml-format - Format code with ruff and verify results
  • /ml-model-export - Export trained PyTorch models to various formats (ONNX, TorchScript, TensorRT)

Knowledge Skills (background knowledge loaded automatically):

  • /ml-lightning-basics - PyTorch Lightning patterns and best practices
  • /ml-hydra-config - Hydra configuration management
  • /ml-pytorch-geometric - Graph Neural Networks with PyTorch Geometric
  • /ml-wandb-tracking - Experiment tracking with Weights & Biases
  • /ml-transformers - Hugging Face Transformers + Lightning integration
  • /ml-cli-tools - Building CLIs with Typer and Rich
  • /tool-pixi - Pixi package manager for ML projects
  • /tool-marimo - Marimo reactive Python notebooks

Agents

Specialized sub-agents with visual color coding for easy identification:

  • ml-architect 🔵 (Opus): Design ML system architectures and model designs
  • training-debugger 🔴 (Sonnet): Diagnose and fix training issues
  • config-generator 🟣 (Sonnet): Generate and validate Hydra configurations
  • pytorch-expert 🟠 (Sonnet): Write efficient PyTorch code and optimize performance
  • geometric-specialist 🔷 (Sonnet): Implement Graph Neural Networks with PyTorch Geometric
  • transformers-specialist 💜 (Sonnet): Hugging Face Transformers expert for LLM fine-tuning, PEFT (LoRA/QLoRA), tokenization, and NLP tasks

Rules

Automatically enforced best practices for ML development:

  • coding-standards: Deterministic operations, type hints, tensor shape documentation, PyTorch Lightning patterns
  • security-practices: API key protection, environment variables, file path sanitization, sensitive data handling
  • workflow-constraints: Config validation, checkpointing requirements, experiment tracking, reproducibility

Hooks

Event-driven automation for code quality:

  • Auto-format Python files with Ruff after Write/Edit
  • Validate YAML configs after writing Hydra configs
  • Remind dependencies installation after package config changes
  • Welcome message on session start with available commands

Supported Frameworks

  • PyTorch Lightning: High-level training framework
  • PyTorch: Core deep learning framework
  • PyTorch Geometric: Graph Neural Networks
  • Hugging Face Transformers: NLP and vision transformers
  • Hydra: Configuration management
  • Weights & Biases: Experiment tracking

Supported Tools

  • pixi: Conda-based package manager with PyPI integration (recommended for GPU/CUDA projects)
  • uv: Fast Python package manager with pip compatibility (recommended for CPU or simple projects)
  • ruff: Linting and formatting
  • ty: Type checking

Development

Local Development

# Clone repository
git clone https://github.com/nishide-dev/claude-code-ml-research.git
cd claude-code-ml-research

# Install ML-specific rules (recommended for all ML projects)
cp -r rules/ml/* ~/.claude/rules/

# Test plugin locally
claude --plugin-dir . code

Rules Installation: Rules are organized in rules/ml/ to coexist with rules from other plugins (e.g., everything-claude-code/rules/common/, everything-claude-code/rules/python/). Install them to your global Claude rules directory for consistent ML development practices across all projects.

Prerequisites

This plugin uses copier (via uvx) to generate projects from templates. Ensure you have uv installed:

# Check if uv is installed
uvx --version

# If not installed, install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

Quick Start

Choosing a Package Manager

The template supports two package managers with different strengths:

Pixi (Recommended for GPU/ML Research)

Best for:

  • Projects requiring CUDA/GPU support
  • Complex ML dependencies (PyTorch, PyTorch Geometric)
  • Reproducible environments across platforms
  • Teams using conda ecosystems

Architecture:

  • Python from conda (base environment)
  • ML packages from PyPI via [pypi-dependencies] (faster, latest versions)
  • PyTorch includes bundled CUDA (no separate CUDA installation needed)
  • Automatic dependency resolution and locking

Installation:

curl -fsSL https://pixi.sh/install.sh | sh
pixi --version  # Should show v0.63+

Commands:

pixi install          # Install all dependencies
pixi run train        # Run training
pixi run test         # Run tests
pixi run format       # Format code

UV (Recommended for Simple/CPU Projects)

Best for:

  • CPU-only projects
  • Simple dependencies
  • Fast iteration and prototyping
  • Projects without complex CUDA requirements

Architecture:

  • Pure Python packages from PyPI
  • Fast dependency resolution
  • pip-compatible workflows

Commands:

uv sync                                      # Install dependencies
uv run python src/my_project/train.py       # Run training
uv run pytest tests/                        # Run tests

1. Initialize New ML Project

/project-init

Creates a new ML research project using the ML Research Template (maintained separately for independent versioning and broader reusability).

Interactive Configuration:

The template will ask you to configure:

  • Project Basics: Name, package name, author info, Python version (3.10-3.13)
  • Package Manager: uv (fast, pip-compatible) or pixi (conda-based, automatic CUDA)
  • PyTorch + CUDA: Choose from 16 presets (PyTorch 2.4-2.9, CUDA 11.8-13.0)
  • ML Frameworks: PyTorch Lightning, Hydra, PyTorch Geometric
  • Experiment Tracking: TensorBoard, Weights & Biases, MLflow, or both
  • Template Type: Image classification, segmentation, GNN, text classification, minimal
  • Dataset: MNIST, CIFAR-10, CIFAR-100, Fashion-MNIST (for vision tasks)
  • Development Tools: ruff, ty, pytest, GitHub Actions, Nix + direnv

What it creates:

  • src/{{ package_name }}/ directory with your Python package
  • Complete Hydra configuration structure
  • PyTorch Lightning model and data module templates
  • Test suite with pytest
  • Development tools: ruff, ty, pytest configured
  • Git repository initialized with proper .gitignore
  • Virtual environment with all dependencies installed

Manual Usage:

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

# From local template clone (for development)
uvx copier copy --trust /path/to/ml-research-template ~/projects/my-ml-project

# Follow the interactive prompts
# cd to your new project
cd ~/projects/my-ml-project

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

2. Generate Configuration

/ml-config

Create new model, data, trainer, or experiment configurations with proper Hydra structure.

3. Train Model

/train

Execute training with:

  • Pre-training validation checks
  • Real-time monitoring
  • Automatic checkpointing
  • Experiment tracking (W&B)
  • Multi-GPU support

4. Debug Issues

/debug

Diagnose and fix:

  • NaN/Inf loss
  • Memory errors (OOM)
  • Slow training
  • Convergence problems
  • Overfitting/underfitting

Example Project Structure

Projects generated from the template have this structure:

my-ml-project/
├── src/
│   └── my_ml_project/           # Package name auto-generated from project name
│       ├── __init__.py
│       ├── train.py             # Main training script with Hydra
│       ├── models/
│       │   ├── __init__.py
│       │   └── base_model.py    # LightningModule template
│       ├── data/
│       │   ├── __init__.py
│       │   └── datamodule.py    # LightningDataModule template
│       ├── utils/
│       │   └── __init__.py
│       └── py.typed             # PEP 561 type hint marker
├── configs/                     # Hydra configurations
│   ├── config.yaml              # Main config with defaults
│   ├── model/
│   │   └── default.yaml         # Model architecture config
│   ├── data/
│   │   └── default.yaml         # Dataset config
│   ├── trainer/
│   │   └── default.yaml         # Lightning Trainer config
│   ├── logger/
│   │   ├── tensorboard.yaml     # TensorBoard logger
│   │   ├── wandb.yaml           # W&B logger
│   │   └── mlflow.yaml          # MLflow logger
│   └── experiment/
│       └── baseline.yaml        # Experiment configuration
├── tests/
│   ├── __init__.py
│   └── test_my_ml_project.py    # Model and data tests
├── .venv/                       # Virtual environment (if using uv)
├── .pixi/                       # Pixi environment (if using pixi)
├── pyproject.toml               # Python project config (if using uv)
├── pixi.toml                    # Pixi config (if using pixi)
├── ruff.toml                    # Ruff linting config
├── .gitignore                   # Git ignore rules
├── .python-version              # Python version
├── README.md                    # Project documentation
└── uv.lock or pixi.lock         # Lock file

Key Features:

  • src/{{ package_name }}/ structure for proper Python packaging
  • Complete Hydra configuration with composition support
  • Lightning callbacks (checkpointing, early stopping, progress bar)
  • Multiple logger options (TensorBoard/W&B/MLflow)
  • Comprehensive tests with pytest
  • Type hints throughout (validated by ty)
  • Ruff-compliant code (no print statements, proper logging)

Training Commands

Generated projects use the src/{{ package_name }}/train.py structure:

Using Pixi (Recommended for GPU)

# Basic training
pixi run train

# Quick test (1 batch per epoch)
pixi run train-debug

# CPU-only training
pixi run train-cpu

# GPU training (explicit)
pixi run train-gpu

# Override Hydra parameters
pixi run train trainer.max_epochs=50 model.lr=0.001

# With specific experiment
pixi run train experiment=baseline

# Multi-GPU training
pixi run train trainer.devices=4 trainer.strategy=ddp

# Development tasks
pixi run test           # Run tests
pixi run test-cov       # Run tests with coverage
pixi run lint           # Check code style
pixi run format         # Format code
pixi run typecheck      # Type check with ty

# TensorBoard
pixi run tensorboard    # Launch TensorBoard (if configured)

Using UV

# Basic training
uv run python src/my_ml_project/train.py

# With specific experiment
uv run python src/my_ml_project/train.py experiment=baseline

# Override parameters
uv run python src/my_ml_project/train.py model.lr=0.001 data.batch_size=128

# Multi-GPU training
uv run python src/my_ml_project/train.py trainer.devices=4 trainer.strategy=ddp

# Hyperparameter sweep
uv run python src/my_ml_project/train.py --multirun model.lr=0.001,0.01,0.1

# Resume from checkpoint
uv run python src/my_ml_project/train.py ckpt_path=checkpoints/best.ckpt

# Quick test (1 batch per epoch)
uv run python src/my_ml_project/train.py trainer.fast_dev_run=true

# Development tasks
uv run pytest tests/ -v              # Run tests
uv run ruff check src/ tests/        # Lint
uv run ruff format src/ tests/       # Format

Configuration Examples

Model Config

# configs/model/resnet50.yaml
_target_: src.models.resnet.ResNetModel

arch: resnet50
num_classes: 1000
pretrained: true

optimizer:
  _target_: torch.optim.AdamW
  lr: 0.001
  weight_decay: 0.0001

scheduler:
  _target_: torch.optim.lr_scheduler.CosineAnnealingLR
  T_max: ${trainer.max_epochs}

Experiment Config

# configs/experiment/baseline.yaml
# @package _global_

defaults:
  - override /model: resnet50
  - override /data: imagenet
  - override /trainer: gpu_ddp

name: "resnet50_imagenet_baseline"
seed: 42

model:
  lr: 0.0001  # Fine-tuning LR

data:
  batch_size: 256

trainer:
  max_epochs: 50
  devices: 4

PyTorch Geometric Support

# Initialize GNN project
/project-init
# Select: PyTorch Geometric, Graph ML task

# Train node classification
python src/train.py model=gnn data=cora

# Train graph classification with batching
python src/train.py model=gnn data=proteins data.batch_size=32

# Large graph with sampling
python src/train.py model=gnn data=ogbn-products data.use_sampling=true

Debugging Guide

NaN Loss

# Run debug command
/debug

# Common fixes:
# 1. Lower learning rate
python src/train.py model.optimizer.lr=0.0001

# 2. Enable gradient clipping
python src/train.py trainer.gradient_clip_val=1.0

# 3. Use full precision
python src/train.py trainer.precision=32

Out of Memory

# 1. Reduce batch size
python src/train.py data.batch_size=32

# 2. Enable gradient accumulation
python src/train.py trainer.accumulate_grad_batches=4

# 3. Mixed precision
python src/train.py trainer.precision=16-mixed

Experiment Tracking

Weights & Biases

# Login to W&B
wandb login

# Train with W&B logging (automatic if logger configured)
python src/train.py

# View experiments
wandb workspace

Compare Experiments

# List all experiments
python scripts/list_experiments.py

# Compare specific experiments
python scripts/compare_experiments.py exp_001 exp_002 exp_003

# Generate report
python scripts/generate_report.py --output report.md

Best Practices

General ML Practices

  1. Always use self.save_hyperparameters() in LightningModule
  2. Use DataModule for all data loading logic
  3. Enable mixed precision for faster training
  4. Use W&B for experiment tracking
  5. Set seeds for reproducibility
  6. Use callbacks for checkpointing and early stopping
  7. Test with fast_dev_run before long training
  8. Monitor GPU utilization (should be >80%)

Package Manager Best Practices

For Pixi Projects

  1. Use pixi.toml tasks instead of remembering long commands

    pixi run train-debug  # Instead of: pixi run python src/...

  2. Lock file is critical - Always commit pixi.lock

    git add pixi.lock  # Ensures reproducible environments

  3. PyPI packages in [pypi-dependencies] - ML packages install faster from PyPI

    • Template already optimized: PyTorch, Lightning from PyPI
    • CUDA is bundled with PyTorch (no separate installation)

  4. Platform-specific environments - Pixi handles this automatically

    platforms = ["linux-64", "osx-arm64"]  # Intel macOS not supported by PyTorch 2.5+

  5. Feature environments for optional dependencies

    pixi run --environment dev test  # Dev environment with extra tools

For UV Projects

  1. Use uv.lock for reproducibility - Commit lock files

  2. Leverage uv's speed - uv sync is much faster than pip

  3. Use extras for optional features

    uv pip install -e ".[dev,wandb]"

Plugin Development

Adding New Skills

# skills/my-skill/SKILL.md
---
name: my-skill
description: Comprehensive guide for... (use third-person voice)
disable-model-invocation: true  # Optional: prevent auto-invocation
---

# My Skill Title

Skill content providing comprehensive guidance...

## Available Templates

See [templates/](templates/) for...

## Examples

See [examples/](examples/) for...

Create supporting files:

mkdir -p skills/my-skill/{templates,examples,scripts}

Note: Skills are invoked as /my-skill (either manually by user or automatically by Claude).

Adding New Agents

# agents/my-agent.md
---
name: my-agent
description: Agent description (should include when to trigger)
tools: ["Read", "Write", "Edit"]
model: sonnet
color: blue
---

You are an expert in...

## Your Role

- ...

Note: Use keyword colors (blue, cyan, green, yellow, magenta, red), not hex codes.

Adding New Rules

# Create rule file in ml subdirectory
echo "# My Rule\n\nRule content..." > rules/ml/my-rule.md

# Reference in plugin.json
# Add "./rules/ml/my-rule.md" to the "rules" array

Rules Organization: Keep rules in rules/ml/ subdirectory for compatibility with other plugins. Users can install them with:

cp -r rules/ml/* ~/.claude/rules/

Pre-commit Hooks

This project uses pre-commit hooks to ensure code quality before commits.

Setup:

# Install dependencies (includes pre-commit)
uv sync --all-extras --dev

# Install pre-commit hooks
uv run pre-commit install

# Run hooks manually on all files
uv run pre-commit run --all-files

Hooks included:

  • ✅ Ruff: Auto-fix linting issues and format code
  • ✅ ty: Type checking
  • ✅ Markdownlint: Auto-fix markdown formatting issues
  • ✅ YAML validation: Check YAML syntax
  • ✅ Python AST check: Verify Python syntax
  • ✅ Trailing whitespace: Remove trailing spaces
  • ✅ Merge conflict detection: Prevent accidental commits with conflicts
  • ✅ Plugin validation: Validate plugin.json structure

Run tests manually:

# Run only fast hooks (recommended for commits)
uv run pre-commit run --all-files

# Run including slow tests
uv run pre-commit run --all-files --hook-stage manual

Skip hooks (not recommended):

git commit --no-verify -m "message"

Contributing

Contributions are welcome! Please:

  1. Fork the repository
  2. Create a feature branch
  3. Make your changes
  4. Add tests if applicable
  5. Submit a pull request

License

MIT License - see LICENSE file for details

Resources

Documentation

Package Managers

Recommended Plugins

This plugin focuses on ML-specific workflows. For Python development tools, we recommend installing the official Astral plugin:

Astral Plugin (UV, Ruff, Ty)

The official Astral plugin provides skills for:

  • uv: Python package and project manager
  • ruff: Linter and formatter
  • ty: Type checker

Installation:

# Add Astral marketplace
/plugin marketplace add astral-sh/claude-code-plugins

# Install Astral plugin
/plugin install astral@astral-sh

Team Installation:

Add to your project's .claude/settings.json:

{
  "extraKnownMarketplaces": {
    "astral-sh": {
      "source": {
        "source": "github",
        "repo": "astral-sh/claude-code-plugins"
      }
    }
  },
  "enabledPlugins": {
    "astral@astral-sh": true
  }
}

What it provides:

Skills for uv, ruff, and ty best practices. Claude will automatically reference these skills when working with Python development tasks.

For more information, see: https://github.com/astral-sh/claude-code-plugins

Everything Claude Code

For general software development workflows (TDD, verification loops, continuous learning), install the comprehensive plugin by Anthropic hackathon winner:

  • Plugin: everything-claude-code
  • Repository: https://github.com/affaan-m/everything-claude-code
  • Agents: 13 agents including architect, code-reviewer, planner, tdd-guide, security-reviewer
  • Commands: 30+ commands including /plan, /tdd, /code-review, /verify, /learn, and more
  • Skills: 29 skills including coding-standards, tdd-workflow, verification-loop, continuous-learning
  • Languages: TypeScript, Python, Go, Java, Django, Spring Boot patterns
  • License: MIT

Installation:

# Add marketplace
/plugin marketplace add affaan-m/everything-claude-code

# Install plugin
/plugin install everything-claude-code@everything-claude-code

# Install rules manually (required - plugins cannot distribute rules automatically)
git clone https://github.com/affaan-m/everything-claude-code.git

# Install common rules (recommended for all projects)
cp -r everything-claude-code/rules/common/* ~/.claude/rules/

# Install language-specific rules as needed
cp -r everything-claude-code/rules/python/* ~/.claude/rules/  # For Python projects
# cp -r everything-claude-code/rules/typescript/* ~/.claude/rules/  # For TypeScript
# cp -r everything-claude-code/rules/golang/* ~/.claude/rules/  # For Go

# Combine with ML Research Plugin rules
git clone https://github.com/nishide-dev/claude-code-ml-research.git
cp -r claude-code-ml-research/rules/ml/* ~/.claude/rules/  # ML-specific rules

Rules Organization: Both plugins organize rules in subdirectories (common/, python/, ml/) to allow seamless coexistence in ~/.claude/rules/.

Usage:

Use commands like /plan, /tdd, /code-review for development workflows and best practices. Combine with /train, /debug, /experiment from ML Research Plugin for comprehensive ML development.

Citation

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

@software{claude_code_ml_research,
  author = {Ryusei Nishide},
  title = {ML Research Plugin for Claude Code},
  year = {2026},
  url = {https://github.com/nishide-dev/claude-code-ml-research}
}

Support

Built with ❤️ for ML researchers using Claude Code

About

Claude Code Plugin for ML Research and Development.

Topics

pytorch pixi uv pytorch-lightning claude-code claude-code-plugin claude-code-skills

Resources

Readme

License

MIT license
Activity

Stars

0 stars

Watchers

1 watching

Forks

0 forks
Report repository

Releases 3

v0.1.2 Latest
Apr 6, 2026
+ 2 releases

Packages

 
 
 

Contributors

Languages