Skip to content

Commit e90e925

Browse files
charlesmartin14charlesmartin14
committed
Add CIFAR-10 ViT optimizer baseline notebook
1 parent 72e80e7 commit e90e925

1 file changed

Lines changed: 203 additions & 0 deletions

File tree

‎baseline/notebooks/CIFAR10_ViT_Optimizer_Baselines.ipynb‎

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# CIFAR-10 Small ViT — SGD, AdamW, and Muon baselines\n",
8+
"\n",
9+
"This notebook trains the **same small Vision Transformer from scratch on CIFAR-10** with three optimizer baselines: SGD + Nesterov momentum, AdamW, and Muon + auxiliary AdamW.\n",
10+
"\n",
11+
"The architecture, data, augmentations, seeds, batch size, epoch budget, loss, evaluation protocol, checkpoint cadence, and WeightWatcher diagnostics are held fixed. Only optimizer-specific hyperparameters differ.\n",
12+
"\n",
13+
"Default model: 32×32 input, 4×4 patches, embedding dimension 192, 6 transformer blocks, 3 attention heads, MLP ratio 4. Default budget: **120 epochs × 3 seeds per optimizer**.\n",
14+
"\n",
15+
"Training uses random crop, horizontal flip, RandAugment, mixup α=0.2, label smoothing 0.1, gradient clipping at 1.0, five-epoch linear warm-up, and cosine decay. Muon is used only for eligible hidden 2-D transformer matrices; all other parameters use auxiliary AdamW.\n"
16+
]
17+
},
18+
{
19+
"cell_type": "code",
20+
"execution_count": null,
21+
"metadata": {},
22+
"outputs": [],
23+
"source": [
24+
"from pathlib import Path\n",
25+
"import os\n",
26+
"import sys\n",
27+
"import math\n",
28+
"import pandas as pd\n",
29+
"import matplotlib.pyplot as plt\n",
30+
"\n",
31+
"ROOT = None\n",
32+
"for path in [Path.cwd(), *Path.cwd().parents]:\n",
33+
" candidate = path / 'baseline'\n",
34+
" if (candidate / 'rg_baselines').is_dir():\n",
35+
" ROOT = candidate\n",
36+
" break\n",
37+
" if (path / 'rg_baselines').is_dir():\n",
38+
" ROOT = path\n",
39+
" break\n",
40+
"if ROOT is None:\n",
41+
" raise RuntimeError('Run this notebook from a clone of CalculatedContent/rg_optimizers.')\n",
42+
"ROOT = ROOT.resolve()\n",
43+
"if str(ROOT) not in sys.path:\n",
44+
" sys.path.insert(0, str(ROOT))\n",
45+
"\n",
46+
"from rg_baselines.vit_cifar10 import (\n",
47+
" DEFAULT_VIT_SEEDS,\n",
48+
" ViTBaselineConfig,\n",
49+
" choose_device,\n",
50+
" muon_parameter_names,\n",
51+
" run_vit_baseline,\n",
52+
" SmallViT,\n",
53+
" summarize_final,\n",
54+
")\n",
55+
"\n",
56+
"def resolve_dir(env_name, default):\n",
57+
" raw = os.environ.get(env_name)\n",
58+
" path = Path(raw).expanduser() if raw else default\n",
59+
" if not path.is_absolute():\n",
60+
" path = Path.cwd() / path\n",
61+
" return path.resolve()\n",
62+
"\n",
63+
"RUN_ROOT = resolve_dir('RG_BASELINE_RUN_ROOT', ROOT / 'runs')\n",
64+
"DATA_DIR = resolve_dir('RG_BASELINE_DATA_DIR', ROOT / 'data')\n",
65+
"EXPERIMENT_ROOT = RUN_ROOT / 'cifar10_vit'\n",
66+
"RUN_ROOT.mkdir(parents=True, exist_ok=True)\n",
67+
"DATA_DIR.mkdir(parents=True, exist_ok=True)\n",
68+
"EXPERIMENT_ROOT.mkdir(parents=True, exist_ok=True)\n",
69+
"DEVICE = choose_device()\n",
70+
"print('device:', DEVICE)\n",
71+
"print('experiment root:', EXPERIMENT_ROOT)\n"
72+
]
73+
},
74+
{
75+
"cell_type": "code",
76+
"execution_count": null,
77+
"metadata": {},
78+
"outputs": [],
79+
"source": [
80+
"CONFIG = ViTBaselineConfig(\n",
81+
" epochs=120,\n",
82+
" batch_size=128,\n",
83+
" warmup_epochs=5,\n",
84+
" patch_size=4,\n",
85+
" embed_dim=192,\n",
86+
" depth=6,\n",
87+
" num_heads=3,\n",
88+
" mlp_ratio=4.0,\n",
89+
" dropout=0.1,\n",
90+
" mixup_alpha=0.2,\n",
91+
" label_smoothing=0.1,\n",
92+
" grad_clip=1.0,\n",
93+
" ww_every=10,\n",
94+
" checkpoint_every=10,\n",
95+
" sgd_lr=0.10, sgd_momentum=0.9, sgd_weight_decay=5e-4,\n",
96+
" adamw_lr=5e-4, adamw_weight_decay=0.05,\n",
97+
" muon_lr=0.02, muon_momentum=0.95, muon_weight_decay=0.01,\n",
98+
" muon_ns_steps=5, muon_aux_lr=3e-4,\n",
99+
" muon_aux_beta1=0.9, muon_aux_beta2=0.95, muon_aux_weight_decay=0.01,\n",
100+
")\n",
101+
"SEEDS = DEFAULT_VIT_SEEDS\n",
102+
"model = SmallViT(CONFIG)\n",
103+
"print(f'parameters: {sum(p.numel() for p in model.parameters()):,}')\n",
104+
"print('Muon hidden matrices:', len(muon_parameter_names(model)))\n",
105+
"display(pd.DataFrame([CONFIG.__dict__]))\n",
106+
"del model\n"
107+
]
108+
},
109+
{
110+
"cell_type": "markdown",
111+
"metadata": {},
112+
"source": [
113+
"## Run the complete baseline suite\n",
114+
"\n",
115+
"This executes **9 trainings total**: 3 optimizers × 3 independent seeds. For a smoke test only, temporarily use `ViTBaselineConfig(epochs=2, warmup_epochs=1, ww_every=2, checkpoint_every=2)` and `SEEDS=(17,)`. The committed defaults above are the intended baseline settings.\n"
116+
]
117+
},
118+
{
119+
"cell_type": "code",
120+
"execution_count": null,
121+
"metadata": {},
122+
"outputs": [],
123+
"source": [
124+
"all_history = []\n",
125+
"all_spectral = []\n",
126+
"for optimizer_name in ('sgd_momentum', 'adamw', 'muon'):\n",
127+
" for seed in SEEDS:\n",
128+
" history, spectral = run_vit_baseline(\n",
129+
" optimizer_name, seed, data_dir=DATA_DIR, output_dir=EXPERIMENT_ROOT,\n",
130+
" config=CONFIG, device=DEVICE, progress=True,\n",
131+
" )\n",
132+
" history.insert(0, 'seed', seed)\n",
133+
" history.insert(0, 'optimizer', optimizer_name)\n",
134+
" spectral.insert(0, 'seed', seed)\n",
135+
" spectral.insert(0, 'optimizer', optimizer_name)\n",
136+
" all_history.append(history)\n",
137+
" all_spectral.append(spectral)\n",
138+
"\n",
139+
"performance = pd.concat(all_history, ignore_index=True)\n",
140+
"spectral = pd.concat(all_spectral, ignore_index=True)\n",
141+
"performance.to_csv(EXPERIMENT_ROOT / 'performance_all_runs.csv', index=False)\n",
142+
"spectral.to_csv(EXPERIMENT_ROOT / 'weightwatcher_all_runs.csv', index=False)\n",
143+
"final_summary = summarize_final(performance, CONFIG.epochs)\n",
144+
"final_summary.to_csv(EXPERIMENT_ROOT / 'final_summary_95ci.csv', index=False)\n",
145+
"display(final_summary)\n"
146+
]
147+
},
148+
{
149+
"cell_type": "code",
150+
"execution_count": null,
151+
"metadata": {},
152+
"outputs": [],
153+
"source": [
154+
"fig, ax = plt.subplots(figsize=(9, 5))\n",
155+
"for optimizer_name, group in performance.groupby('optimizer'):\n",
156+
" stats = group.groupby('epoch')['test_accuracy'].agg(['mean', 'std']).reset_index()\n",
157+
" ax.plot(stats['epoch'], 100 * stats['mean'], label=optimizer_name)\n",
158+
" ax.fill_between(stats['epoch'], 100*(stats['mean']-stats['std'].fillna(0)),\n",
159+
" 100*(stats['mean']+stats['std'].fillna(0)), alpha=0.15)\n",
160+
"ax.set_xlabel('Epoch')\n",
161+
"ax.set_ylabel('CIFAR-10 test accuracy (%)')\n",
162+
"ax.set_title('Small ViT optimizer baselines')\n",
163+
"ax.grid(alpha=0.25)\n",
164+
"ax.legend()\n",
165+
"fig.tight_layout()\n",
166+
"fig.savefig(EXPERIMENT_ROOT / 'test_accuracy_comparison.png', dpi=160)\n",
167+
"plt.show()\n",
168+
"\n",
169+
"fig, ax = plt.subplots(figsize=(9, 5))\n",
170+
"for optimizer_name, group in performance.groupby('optimizer'):\n",
171+
" stats = group.groupby('epoch')['test_loss'].agg(['mean', 'std']).reset_index()\n",
172+
" ax.plot(stats['epoch'], stats['mean'], label=optimizer_name)\n",
173+
" ax.fill_between(stats['epoch'], stats['mean']-stats['std'].fillna(0),\n",
174+
" stats['mean']+stats['std'].fillna(0), alpha=0.15)\n",
175+
"ax.set_xlabel('Epoch')\n",
176+
"ax.set_ylabel('CIFAR-10 test cross-entropy')\n",
177+
"ax.set_title('Small ViT optimizer baselines')\n",
178+
"ax.grid(alpha=0.25)\n",
179+
"ax.legend()\n",
180+
"fig.tight_layout()\n",
181+
"fig.savefig(EXPERIMENT_ROOT / 'test_loss_comparison.png', dpi=160)\n",
182+
"plt.show()\n"
183+
]
184+
},
185+
{
186+
"cell_type": "markdown",
187+
"metadata": {},
188+
"source": [
189+
"## Persisted outputs\n",
190+
"\n",
191+
"Each optimizer/seed run saves `history.csv`, `weightwatcher_by_epoch_layer.csv`, checkpoints every 10 epochs, `config.json`, and `final_state.pt`. The suite also saves aggregate performance and WeightWatcher CSVs, a final 95% Student-t confidence-interval table, and comparison plots.\n",
192+
"\n",
193+
"WeightWatcher is called with `analyze(ERG=True, randomize=True)` so the run records the package-provided spectral diagnostics including alpha, randomized correlation-trap counts when available, and ERG gap. No proxy trap count is substituted.\n"
194+
]
195+
}
196+
],
197+
"metadata": {
198+
"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
199+
"language_info": {"name": "python", "version": "3"}
200+
},
201+
"nbformat": 4,
202+
"nbformat_minor": 5
203+
}

0 commit comments

Comments
 (0)