Skip to content

Commit 820d5c7

Browse files
hemildesaitchatonrohitgr7
authored
Add a notebook example to reach a quick baseline of ~94% accuracy on CIFAR (#4818)
* Add a notebook example to reach a quick baseline of ~94% accuracy on CIFAR10 using Resnet in Lightning * Remove outputs * PR Feedback * some changes * some more changes Co-authored-by: chaton <[email protected]> Co-authored-by: rohitgr7 <[email protected]>
1 parent 4ebce38 commit 820d5c7

File tree

2 files changed

+402
-7
lines changed

2 files changed

+402
-7
lines changed
Lines changed: 394 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,394 @@
1+
{
2+
"nbformat": 4,
3+
"nbformat_minor": 0,
4+
"metadata": {
5+
"accelerator": "GPU",
6+
"colab": {
7+
"name": "06_cifar10_baseline.ipynb",
8+
"provenance": [],
9+
"collapsed_sections": []
10+
},
11+
"kernelspec": {
12+
"display_name": "Python 3",
13+
"name": "python3"
14+
}
15+
},
16+
"cells": [
17+
{
18+
"cell_type": "markdown",
19+
"metadata": {
20+
"id": "qMDj0BYNECU8"
21+
},
22+
"source": [
23+
"<a href=\"https://colab.research.google.com/github/PytorchLightning/pytorch-lightning/blob/master/notebooks/06-cifar10-pytorch-lightning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
24+
]
25+
},
26+
{
27+
"cell_type": "markdown",
28+
"metadata": {
29+
"id": "ECu0zDh8UXU8"
30+
},
31+
"source": [
32+
"# PyTorch Lightning CIFAR10 ~94% Baseline Tutorial ⚡\n",
33+
"\n",
34+
"Train a Resnet to 94% accuracy on Cifar10!\n",
35+
"\n",
36+
"Main takeaways:\n",
37+
"1. Experiment with different Learning Rate schedules and frequencies in the configure_optimizers method in pl.LightningModule\n",
38+
"2. Use an existing Resnet architecture with modifications directly with Lightning\n",
39+
"\n",
40+
"---\n",
41+
"\n",
42+
" - Give us a ⭐ [on Github](https://www.github.com/PytorchLightning/pytorch-lightning/)\n",
43+
" - Check out [the documentation](https://pytorch-lightning.readthedocs.io/en/latest/)\n",
44+
" - Join us [on Slack](https://join.slack.com/t/pytorch-lightning/shared_invite/zt-f6bl2l0l-JYMK3tbAgAmGRrlNr00f1A)"
45+
]
46+
},
47+
{
48+
"cell_type": "markdown",
49+
"metadata": {
50+
"id": "HYpMlx7apuHq"
51+
},
52+
"source": [
53+
"### Setup\n",
54+
"Lightning is easy to install. Simply `pip install pytorch-lightning`.\n",
55+
"Also check out [bolts](https://github.com/PyTorchLightning/pytorch-lightning-bolts/) for pre-existing data modules and models."
56+
]
57+
},
58+
{
59+
"cell_type": "code",
60+
"metadata": {
61+
"id": "ziAQCrE-TYWG"
62+
},
63+
"source": [
64+
"! pip install pytorch-lightning pytorch-lightning-bolts -qU"
65+
],
66+
"execution_count": null,
67+
"outputs": []
68+
},
69+
{
70+
"cell_type": "code",
71+
"metadata": {
72+
"id": "L-W_Gq2FORoU"
73+
},
74+
"source": [
75+
"# Run this if you intend to use TPUs\n",
76+
"# !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n",
77+
"# !python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev"
78+
],
79+
"execution_count": null,
80+
"outputs": []
81+
},
82+
{
83+
"cell_type": "code",
84+
"metadata": {
85+
"id": "wjov-2N_TgeS"
86+
},
87+
"source": [
88+
"import torch\n",
89+
"import torch.nn as nn\n",
90+
"import torch.nn.functional as F\n",
91+
"from torch.optim.lr_scheduler import OneCycleLR\n",
92+
"from torch.optim.swa_utils import AveragedModel, update_bn\n",
93+
"import torchvision\n",
94+
"\n",
95+
"import pytorch_lightning as pl\n",
96+
"from pytorch_lightning.callbacks import LearningRateMonitor\n",
97+
"from pytorch_lightning.metrics.functional import accuracy\n",
98+
"from pl_bolts.datamodules import CIFAR10DataModule\n",
99+
"from pl_bolts.transforms.dataset_normalizations import cifar10_normalization"
100+
],
101+
"execution_count": null,
102+
"outputs": []
103+
},
104+
{
105+
"cell_type": "code",
106+
"metadata": {
107+
"id": "54JMU1N-0y0g"
108+
},
109+
"source": [
110+
"pl.seed_everything(7);"
111+
],
112+
"execution_count": null,
113+
"outputs": []
114+
},
115+
{
116+
"cell_type": "markdown",
117+
"metadata": {
118+
"id": "FA90qwFcqIXR"
119+
},
120+
"source": [
121+
"### CIFAR10 Data Module\n",
122+
"\n",
123+
"Import the existing data module from `bolts` and modify the train and test transforms."
124+
]
125+
},
126+
{
127+
"cell_type": "code",
128+
"metadata": {
129+
"id": "S9e-W8CSa8nH"
130+
},
131+
"source": [
132+
"batch_size = 32\n",
133+
"\n",
134+
"train_transforms = torchvision.transforms.Compose([\n",
135+
" torchvision.transforms.RandomCrop(32, padding=4),\n",
136+
" torchvision.transforms.RandomHorizontalFlip(),\n",
137+
" torchvision.transforms.ToTensor(),\n",
138+
" cifar10_normalization(),\n",
139+
"])\n",
140+
"\n",
141+
"test_transforms = torchvision.transforms.Compose([\n",
142+
" torchvision.transforms.ToTensor(),\n",
143+
" cifar10_normalization(),\n",
144+
"])\n",
145+
"\n",
146+
"cifar10_dm = CIFAR10DataModule(\n",
147+
" batch_size=batch_size,\n",
148+
" train_transforms=train_transforms,\n",
149+
" test_transforms=test_transforms,\n",
150+
" val_transforms=test_transforms,\n",
151+
")"
152+
],
153+
"execution_count": null,
154+
"outputs": []
155+
},
156+
{
157+
"cell_type": "markdown",
158+
"metadata": {
159+
"id": "SfCsutp3qUMc"
160+
},
161+
"source": [
162+
"### Resnet\n",
163+
"Modify the pre-existing Resnet architecture from TorchVision. The pre-existing architecture is based on ImageNet images (224x224) as input. So we need to modify it for CIFAR10 images (32x32)."
164+
]
165+
},
166+
{
167+
"cell_type": "code",
168+
"metadata": {
169+
"id": "GNSeJgwvhHp-"
170+
},
171+
"source": [
172+
"def create_model():\n",
173+
" model = torchvision.models.resnet18(pretrained=False, num_classes=10)\n",
174+
" model.conv1 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
175+
" model.maxpool = nn.Identity()\n",
176+
" return model"
177+
],
178+
"execution_count": null,
179+
"outputs": []
180+
},
181+
{
182+
"cell_type": "markdown",
183+
"metadata": {
184+
"id": "HUCj5TKsqty1"
185+
},
186+
"source": [
187+
"### Lightning Module\n",
188+
"Check out the [`configure_optimizers`](https://pytorch-lightning.readthedocs.io/en/stable/lightning_module.html#configure-optimizers) method to use custom Learning Rate schedulers. The OneCycleLR with SGD will get you to around 92-93% accuracy in 20-30 epochs and 93-94% accuracy in 40-50 epochs. Feel free to experiment with different LR schedules from https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate"
189+
]
190+
},
191+
{
192+
"cell_type": "code",
193+
"metadata": {
194+
"id": "03OMrBa5iGtT"
195+
},
196+
"source": [
197+
"class LitResnet(pl.LightningModule):\n",
198+
" def __init__(self, lr=0.05):\n",
199+
" super().__init__()\n",
200+
"\n",
201+
" self.save_hyperparameters()\n",
202+
" self.model = create_model()\n",
203+
"\n",
204+
" def forward(self, x):\n",
205+
" out = self.model(x)\n",
206+
" return F.log_softmax(out, dim=1)\n",
207+
"\n",
208+
" def training_step(self, batch, batch_idx):\n",
209+
" x, y = batch\n",
210+
" logits = F.log_softmax(self.model(x), dim=1)\n",
211+
" loss = F.nll_loss(logits, y)\n",
212+
" self.log('train_loss', loss)\n",
213+
" return loss\n",
214+
"\n",
215+
" def evaluate(self, batch, stage=None):\n",
216+
" x, y = batch\n",
217+
" logits = self(x)\n",
218+
" loss = F.nll_loss(logits, y)\n",
219+
" preds = torch.argmax(logits, dim=1)\n",
220+
" acc = accuracy(preds, y)\n",
221+
"\n",
222+
" if stage:\n",
223+
" self.log(f'{stage}_loss', loss, prog_bar=True)\n",
224+
" self.log(f'{stage}_acc', acc, prog_bar=True)\n",
225+
"\n",
226+
" def validation_step(self, batch, batch_idx):\n",
227+
" self.evaluate(batch, 'val')\n",
228+
"\n",
229+
" def test_step(self, batch, batch_idx):\n",
230+
" self.evaluate(batch, 'test')\n",
231+
"\n",
232+
" def configure_optimizers(self):\n",
233+
" optimizer = torch.optim.SGD(self.parameters(), lr=self.hparams.lr, momentum=0.9, weight_decay=5e-4)\n",
234+
" steps_per_epoch = 45000 // batch_size\n",
235+
" scheduler_dict = {\n",
236+
" 'scheduler': OneCycleLR(optimizer, 0.1, epochs=self.trainer.max_epochs, steps_per_epoch=steps_per_epoch),\n",
237+
" 'interval': 'step',\n",
238+
" }\n",
239+
" return {'optimizer': optimizer, 'lr_scheduler': scheduler_dict}"
240+
],
241+
"execution_count": null,
242+
"outputs": []
243+
},
244+
{
245+
"cell_type": "code",
246+
"metadata": {
247+
"id": "3FFPgpAFi9KU"
248+
},
249+
"source": [
250+
"model = LitResnet(lr=0.05)\n",
251+
"model.datamodule = cifar10_dm\n",
252+
"\n",
253+
"trainer = pl.Trainer(\n",
254+
" progress_bar_refresh_rate=20,\n",
255+
" max_epochs=40,\n",
256+
" gpus=1,\n",
257+
" logger=pl.loggers.TensorBoardLogger('lightning_logs/', name='resnet'),\n",
258+
" callbacks=[LearningRateMonitor(logging_interval='step')],\n",
259+
")\n",
260+
"\n",
261+
"trainer.fit(model, cifar10_dm)\n",
262+
"trainer.test(model, datamodule=cifar10_dm);"
263+
],
264+
"execution_count": null,
265+
"outputs": []
266+
},
267+
{
268+
"cell_type": "markdown",
269+
"metadata": {
270+
"id": "lWL_WpeVIXWQ"
271+
},
272+
"source": [
273+
"### Bonus: Use [Stochastic Weight Averaging](https://arxiv.org/abs/1803.05407) to get a boost on performance\n",
274+
"\n",
275+
"Use SWA from torch.optim to get a quick performance boost. Also shows a couple of cool features from Lightning:\n",
276+
"- Use `training_epoch_end` to run code after the end of every epoch\n",
277+
"- Use a pretrained model directly with this wrapper for SWA"
278+
]
279+
},
280+
{
281+
"cell_type": "code",
282+
"metadata": {
283+
"id": "bsSwqKv0t9uY"
284+
},
285+
"source": [
286+
"class SWAResnet(LitResnet):\n",
287+
" def __init__(self, trained_model, lr=0.01):\n",
288+
" super().__init__()\n",
289+
"\n",
290+
" self.save_hyperparameters('lr')\n",
291+
" self.model = trained_model\n",
292+
" self.swa_model = AveragedModel(self.model)\n",
293+
"\n",
294+
" def forward(self, x):\n",
295+
" out = self.swa_model(x)\n",
296+
" return F.log_softmax(out, dim=1)\n",
297+
"\n",
298+
" def training_epoch_end(self, training_step_outputs):\n",
299+
" self.swa_model.update_parameters(self.model)\n",
300+
"\n",
301+
" def validation_step(self, batch, batch_idx, stage=None):\n",
302+
" x, y = batch\n",
303+
" logits = F.log_softmax(self.model(x), dim=1)\n",
304+
" loss = F.nll_loss(logits, y)\n",
305+
" preds = torch.argmax(logits, dim=1)\n",
306+
" acc = accuracy(preds, y)\n",
307+
"\n",
308+
" self.log(f'val_loss', loss, prog_bar=True)\n",
309+
" self.log(f'val_acc', acc, prog_bar=True)\n",
310+
"\n",
311+
" def configure_optimizers(self):\n",
312+
" optimizer = torch.optim.SGD(self.model.parameters(), lr=self.hparams.lr, momentum=0.9, weight_decay=5e-4)\n",
313+
" return optimizer\n",
314+
"\n",
315+
" def on_train_end(self):\n",
316+
" update_bn(self.datamodule.train_dataloader(), self.swa_model, device=self.device)"
317+
],
318+
"execution_count": null,
319+
"outputs": []
320+
},
321+
{
322+
"cell_type": "code",
323+
"metadata": {
324+
"id": "cA6ZG7C74rjL"
325+
},
326+
"source": [
327+
"swa_model = SWAResnet(model.model, lr=0.01)\n",
328+
"swa_model.datamodule = cifar10_dm\n",
329+
"\n",
330+
"swa_trainer = pl.Trainer(\n",
331+
" progress_bar_refresh_rate=20,\n",
332+
" max_epochs=20,\n",
333+
" gpus=1,\n",
334+
" logger=pl.loggers.TensorBoardLogger('lightning_logs/', name='swa_resnet'),\n",
335+
")\n",
336+
"\n",
337+
"swa_trainer.fit(swa_model, cifar10_dm)\n",
338+
"swa_trainer.test(swa_model, datamodule=cifar10_dm);"
339+
],
340+
"execution_count": null,
341+
"outputs": []
342+
},
343+
{
344+
"cell_type": "code",
345+
"metadata": {
346+
"id": "RRHMfGiDpZ2M"
347+
},
348+
"source": [
349+
"# Start tensorboard.\n",
350+
"%reload_ext tensorboard\n",
351+
"%tensorboard --logdir lightning_logs/"
352+
],
353+
"execution_count": null,
354+
"outputs": []
355+
},
356+
{
357+
"cell_type": "markdown",
358+
"metadata": {
359+
"id": "RltpFGS-s0M1"
360+
},
361+
"source": [
362+
"<code style=\"color:#792ee5;\">\n",
363+
" <h1> <strong> Congratulations - Time to Join the Community! </strong> </h1>\n",
364+
"</code>\n",
365+
"\n",
366+
"Congratulations on completing this notebook tutorial! If you enjoyed this and would like to join the Lightning movement, you can do so in the following ways!\n",
367+
"\n",
368+
"### Star [Lightning](https://github.com/PyTorchLightning/pytorch-lightning) on GitHub\n",
369+
"The easiest way to help our community is just by starring the GitHub repos! This helps raise awareness of the cool tools we're building.\n",
370+
"\n",
371+
"* Please, star [Lightning](https://github.com/PyTorchLightning/pytorch-lightning)\n",
372+
"\n",
373+
"### Join our [Slack](https://join.slack.com/t/pytorch-lightning/shared_invite/zt-f6bl2l0l-JYMK3tbAgAmGRrlNr00f1A)!\n",
374+
"The best way to keep up to date on the latest advancements is to join our community! Make sure to introduce yourself and share your interests in `#general` channel\n",
375+
"\n",
376+
"### Interested by SOTA AI models ! Check out [Bolt](https://github.com/PyTorchLightning/pytorch-lightning-bolts)\n",
377+
"Bolts has a collection of state-of-the-art models, all implemented in [Lightning](https://github.com/PyTorchLightning/pytorch-lightning) and can be easily integrated within your own projects.\n",
378+
"\n",
379+
"* Please, star [Bolt](https://github.com/PyTorchLightning/pytorch-lightning-bolts)\n",
380+
"\n",
381+
"### Contributions !\n",
382+
"The best way to contribute to our community is to become a code contributor! At any time you can go to [Lightning](https://github.com/PyTorchLightning/pytorch-lightning) or [Bolt](https://github.com/PyTorchLightning/pytorch-lightning-bolts) GitHub Issues page and filter for \"good first issue\". \n",
383+
"\n",
384+
"* [Lightning good first issue](https://github.com/PyTorchLightning/pytorch-lightning/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)\n",
385+
"* [Bolt good first issue](https://github.com/PyTorchLightning/pytorch-lightning-bolts/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)\n",
386+
"* You can also contribute your own notebooks with useful examples !\n",
387+
"\n",
388+
"### Great thanks from the entire Pytorch Lightning Team for your interest !\n",
389+
"\n",
390+
"<img src=\"https://github.com/PyTorchLightning/pytorch-lightning/blob/master/docs/source/_images/logos/lightning_logo-name.png?raw=true\" width=\"800\" height=\"200\" />"
391+
]
392+
}
393+
]
394+
}

0 commit comments

Comments
 (0)