wehub-resource-sync
1.8 KiB
1.8 KiB
PyTorch Hyperparameter Optimization Example
This example demonstrates hyperparameter optimization with MLflow tracking using pure PyTorch (no Lightning dependencies).
What it demonstrates
- MLflow nested runs: Parent run tracks the overall HPO experiment, child runs track individual trials
- Hyperparameter tuning: Uses Optuna to optimize learning rate, hidden layer size, dropout rate, and batch size
- Pure PyTorch: Simple, clean implementation without framework overhead
- Fast training: MNIST classification completes quickly for rapid iteration
Architecture
The model is a simple 2-layer neural network:
Input (784) → FC1 (hidden_size) → ReLU → Dropout → FC2 (10) → LogSoftmax
Hyperparameters optimized
lr: Learning rate (1e-4 to 1e-1, log scale)hidden_size: Hidden layer size (64 to 512, step 64)dropout_rate: Dropout probability (0.1 to 0.5)batch_size: Batch size (32, 64, or 128)
Running the example
Quick test (3 trials, 3 epochs each)
python hpo_mnist.py --n-trials 3 --max-epochs 3
Full optimization (10 trials, 5 epochs each)
python hpo_mnist.py --n-trials 10 --max-epochs 5
Using MLflow projects
mlflow run . -P n_trials=5 -P max_epochs=3
Viewing results
After running, view the results in MLflow UI:
mlflow server
Navigate to http://localhost:5000 to see:
- Parent run with overall HPO results
- Child runs for each trial with their hyperparameters and metrics
- Comparison view to analyze which hyperparameters work best
Dependencies
torch>=2.1: PyTorch for model trainingtorchvision>=0.15.1: MNIST datasetoptuna>=3.0.0: Hyperparameter optimization frameworkmlflow: Experiment tracking
No Lightning, no torchmetrics, no transformers = no dependency conflicts! 🎉