ML Doctor (mldoct) 🩺

ML Doctor is an enterprise-grade AI developer tool designed to optimize PyTorch engineering workflows. It combines V2 Data-Flow AST analysis, C++ GPU VRAM profiling, and XGBoost-powered diagnostics to eliminate silent bugs, prevent OOM crashes, and save cloud compute costs.

🔥 Enterprise Features (v1.0.0)

1. V2 Data-Flow Linter (check):

   • Detects missing optimization steps (optimizer.zero_grad(), model.eval()).
   • Identifies Semantic Logic Bugs (e.g., Double Softmax from CrossEntropyLoss).
   • Flags CPU Bottlenecks (e.g., num_workers=0 in DataLoaders).
   • Features an Automated Fix Engine (--fix) to dynamically patch source files.

2. Smart GPU Watchdog (watch):

   • Spawns a background C++ profiling thread via nvidia-smi to catch "Staircase" VRAM leaks before an OOM crash occurs.
   • Intercepts NaN, Inf, and exploding loss spikes in real-time.
   • Cost-Saver (Early Stopping): Automatically terminates cloud instances if validation loss plateaus, saving expensive AWS/GCP bills.

3. Pre-Flight Auto-Tuner API:

   • Employs a Binary Search algorithm to stress-test your GPU prior to training, locking in the absolute maximum safe batch size for 100% hardware utilization.

4. Machine Learning Diagnoser (diagnose):

   • Uses an embedded scikit-learn HistGradientBoostingClassifier to analyze training CSV logs.
   • Diagnoses Vanishing/Exploding Gradients, Severe Overfitting, and Missing Gradients with 99% accuracy.
   • Exports professional Markdown post-mortem reports.

📦 Installation

pip install mldoct

💻 CLI Usage Guide

1. V2 Semantic Linting & Auto-Fixing

Scan a script for structural bugs, logic errors, and GPU bottlenecks.

mldoctor check train.py

# Automatically patch identified omissions in-place
mldoctor check train.py --fix

2. Live Watchdog & Memory Profiler

Launch training with live VRAM tracking and Smart Early Stopping active.

mldoctor watch train.py --epochs 100

3. ML Curve Diagnosis

Diagnose a specific metric log or let the tool auto-discover the latest run in a directory.

# Auto-discover CSV in logs/ and export a markdown report
mldoctor diagnose ./logs/ --export-report

🚀 API Usage: Pre-Flight Auto-Tuner

Maximize your GPU utilization by dynamically finding the perfect batch size before initializing your DataLoader. Add this to your PyTorch script:

from mldoctor.tuner import tune_batch_size
from torch.utils.data import DataLoader

# 1. Initialize your model
model = MyResNet().to('cuda')

# 2. Run the Pre-Flight Dry Run (Pass the shape of a single dataset item)
optimal_batch = tune_batch_size(model, dataset_sample_shape=(3, 224, 224))

# 3. Proceed with 100% hardware utilization
train_loader = DataLoader(dataset, batch_size=optimal_batch, num_workers=4)

👨‍💻 Author & Maintainer

• Author: Shravan Gupta
• Portfolio & Docs: shravangupta.dev
• License: MIT