wehub-resource-sync
69 lines
1.8 KiB
Markdown
69 lines
1.8 KiB
Markdown
# 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)
|
|
|
|
```bash
|
|
python hpo_mnist.py --n-trials 3 --max-epochs 3
|
|
```
|
|
|
|
### Full optimization (10 trials, 5 epochs each)
|
|
|
|
```bash
|
|
python hpo_mnist.py --n-trials 10 --max-epochs 5
|
|
```
|
|
|
|
### Using MLflow projects
|
|
|
|
```bash
|
|
mlflow run . -P n_trials=5 -P max_epochs=3
|
|
```
|
|
|
|
## Viewing results
|
|
|
|
After running, view the results in MLflow UI:
|
|
|
|
```bash
|
|
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 training
|
|
- `torchvision>=0.15.1`: MNIST dataset
|
|
- `optuna>=3.0.0`: Hyperparameter optimization framework
|
|
- `mlflow`: Experiment tracking
|
|
|
|
**No Lightning, no torchmetrics, no transformers** = no dependency conflicts! 🎉
|