lightning/notebooks/07-cifar10-baseline.ipynb

395 lines
15 KiB
Plaintext

{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"accelerator": "GPU",
"colab": {
"name": "07-cifar10-baseline.ipynb",
"provenance": [],
"collapsed_sections": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "qMDj0BYNECU8"
},
"source": [
"<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>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ECu0zDh8UXU8"
},
"source": [
"# PyTorch Lightning CIFAR10 ~94% Baseline Tutorial ⚡\n",
"\n",
"Train a Resnet to 94% accuracy on Cifar10!\n",
"\n",
"Main takeaways:\n",
"1. Experiment with different Learning Rate schedules and frequencies in the configure_optimizers method in pl.LightningModule\n",
"2. Use an existing Resnet architecture with modifications directly with Lightning\n",
"\n",
"---\n",
"\n",
" - Give us a ⭐ [on Github](https://www.github.com/PytorchLightning/pytorch-lightning/)\n",
" - Check out [the documentation](https://pytorch-lightning.readthedocs.io/en/latest/)\n",
" - Join us [on Slack](https://join.slack.com/t/pytorch-lightning/shared_invite/zt-f6bl2l0l-JYMK3tbAgAmGRrlNr00f1A)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HYpMlx7apuHq"
},
"source": [
"### Setup\n",
"Lightning is easy to install. Simply `pip install pytorch-lightning`.\n",
"Also check out [bolts](https://github.com/PyTorchLightning/pytorch-lightning-bolts/) for pre-existing data modules and models."
]
},
{
"cell_type": "code",
"metadata": {
"id": "ziAQCrE-TYWG"
},
"source": [
"! pip install pytorch-lightning lightning-bolts -qU"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "L-W_Gq2FORoU"
},
"source": [
"# Run this if you intend to use TPUs\n",
"# !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n",
"# !python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "wjov-2N_TgeS"
},
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"from torch.optim.lr_scheduler import OneCycleLR\n",
"from torch.optim.swa_utils import AveragedModel, update_bn\n",
"import torchvision\n",
"\n",
"import pytorch_lightning as pl\n",
"from pytorch_lightning.callbacks import LearningRateMonitor\n",
"from pytorch_lightning.metrics.functional import accuracy\n",
"from pl_bolts.datamodules import CIFAR10DataModule\n",
"from pl_bolts.transforms.dataset_normalizations import cifar10_normalization"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "54JMU1N-0y0g"
},
"source": [
"pl.seed_everything(7);"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "FA90qwFcqIXR"
},
"source": [
"### CIFAR10 Data Module\n",
"\n",
"Import the existing data module from `bolts` and modify the train and test transforms."
]
},
{
"cell_type": "code",
"metadata": {
"id": "S9e-W8CSa8nH"
},
"source": [
"batch_size = 32\n",
"\n",
"train_transforms = torchvision.transforms.Compose([\n",
" torchvision.transforms.RandomCrop(32, padding=4),\n",
" torchvision.transforms.RandomHorizontalFlip(),\n",
" torchvision.transforms.ToTensor(),\n",
" cifar10_normalization(),\n",
"])\n",
"\n",
"test_transforms = torchvision.transforms.Compose([\n",
" torchvision.transforms.ToTensor(),\n",
" cifar10_normalization(),\n",
"])\n",
"\n",
"cifar10_dm = CIFAR10DataModule(\n",
" batch_size=batch_size,\n",
" train_transforms=train_transforms,\n",
" test_transforms=test_transforms,\n",
" val_transforms=test_transforms,\n",
")"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "SfCsutp3qUMc"
},
"source": [
"### Resnet\n",
"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)."
]
},
{
"cell_type": "code",
"metadata": {
"id": "GNSeJgwvhHp-"
},
"source": [
"def create_model():\n",
" model = torchvision.models.resnet18(pretrained=False, num_classes=10)\n",
" model.conv1 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
" model.maxpool = nn.Identity()\n",
" return model"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "HUCj5TKsqty1"
},
"source": [
"### Lightning Module\n",
"Check out the [`configure_optimizers`](https://pytorch-lightning.readthedocs.io/en/stable/common/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"
]
},
{
"cell_type": "code",
"metadata": {
"id": "03OMrBa5iGtT"
},
"source": [
"class LitResnet(pl.LightningModule):\n",
" def __init__(self, lr=0.05):\n",
" super().__init__()\n",
"\n",
" self.save_hyperparameters()\n",
" self.model = create_model()\n",
"\n",
" def forward(self, x):\n",
" out = self.model(x)\n",
" return F.log_softmax(out, dim=1)\n",
"\n",
" def training_step(self, batch, batch_idx):\n",
" x, y = batch\n",
" logits = F.log_softmax(self.model(x), dim=1)\n",
" loss = F.nll_loss(logits, y)\n",
" self.log('train_loss', loss)\n",
" return loss\n",
"\n",
" def evaluate(self, batch, stage=None):\n",
" x, y = batch\n",
" logits = self(x)\n",
" loss = F.nll_loss(logits, y)\n",
" preds = torch.argmax(logits, dim=1)\n",
" acc = accuracy(preds, y)\n",
"\n",
" if stage:\n",
" self.log(f'{stage}_loss', loss, prog_bar=True)\n",
" self.log(f'{stage}_acc', acc, prog_bar=True)\n",
"\n",
" def validation_step(self, batch, batch_idx):\n",
" self.evaluate(batch, 'val')\n",
"\n",
" def test_step(self, batch, batch_idx):\n",
" self.evaluate(batch, 'test')\n",
"\n",
" def configure_optimizers(self):\n",
" optimizer = torch.optim.SGD(self.parameters(), lr=self.hparams.lr, momentum=0.9, weight_decay=5e-4)\n",
" steps_per_epoch = 45000 // batch_size\n",
" scheduler_dict = {\n",
" 'scheduler': OneCycleLR(optimizer, 0.1, epochs=self.trainer.max_epochs, steps_per_epoch=steps_per_epoch),\n",
" 'interval': 'step',\n",
" }\n",
" return {'optimizer': optimizer, 'lr_scheduler': scheduler_dict}"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "3FFPgpAFi9KU"
},
"source": [
"model = LitResnet(lr=0.05)\n",
"model.datamodule = cifar10_dm\n",
"\n",
"trainer = pl.Trainer(\n",
" progress_bar_refresh_rate=20,\n",
" max_epochs=40,\n",
" gpus=1,\n",
" logger=pl.loggers.TensorBoardLogger('lightning_logs/', name='resnet'),\n",
" callbacks=[LearningRateMonitor(logging_interval='step')],\n",
")\n",
"\n",
"trainer.fit(model, cifar10_dm)\n",
"trainer.test(model, datamodule=cifar10_dm);"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "lWL_WpeVIXWQ"
},
"source": [
"### Bonus: Use [Stochastic Weight Averaging](https://arxiv.org/abs/1803.05407) to get a boost on performance\n",
"\n",
"Use SWA from torch.optim to get a quick performance boost. Also shows a couple of cool features from Lightning:\n",
"- Use `training_epoch_end` to run code after the end of every epoch\n",
"- Use a pretrained model directly with this wrapper for SWA"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bsSwqKv0t9uY"
},
"source": [
"class SWAResnet(LitResnet):\n",
" def __init__(self, trained_model, lr=0.01):\n",
" super().__init__()\n",
"\n",
" self.save_hyperparameters('lr')\n",
" self.model = trained_model\n",
" self.swa_model = AveragedModel(self.model)\n",
"\n",
" def forward(self, x):\n",
" out = self.swa_model(x)\n",
" return F.log_softmax(out, dim=1)\n",
"\n",
" def training_epoch_end(self, training_step_outputs):\n",
" self.swa_model.update_parameters(self.model)\n",
"\n",
" def validation_step(self, batch, batch_idx, stage=None):\n",
" x, y = batch\n",
" logits = F.log_softmax(self.model(x), dim=1)\n",
" loss = F.nll_loss(logits, y)\n",
" preds = torch.argmax(logits, dim=1)\n",
" acc = accuracy(preds, y)\n",
"\n",
" self.log(f'val_loss', loss, prog_bar=True)\n",
" self.log(f'val_acc', acc, prog_bar=True)\n",
"\n",
" def configure_optimizers(self):\n",
" optimizer = torch.optim.SGD(self.model.parameters(), lr=self.hparams.lr, momentum=0.9, weight_decay=5e-4)\n",
" return optimizer\n",
"\n",
" def on_train_end(self):\n",
" update_bn(self.datamodule.train_dataloader(), self.swa_model, device=self.device)"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "cA6ZG7C74rjL"
},
"source": [
"swa_model = SWAResnet(model.model, lr=0.01)\n",
"swa_model.datamodule = cifar10_dm\n",
"\n",
"swa_trainer = pl.Trainer(\n",
" progress_bar_refresh_rate=20,\n",
" max_epochs=20,\n",
" gpus=1,\n",
" logger=pl.loggers.TensorBoardLogger('lightning_logs/', name='swa_resnet'),\n",
")\n",
"\n",
"swa_trainer.fit(swa_model, cifar10_dm)\n",
"swa_trainer.test(swa_model, datamodule=cifar10_dm);"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "RRHMfGiDpZ2M"
},
"source": [
"# Start tensorboard.\n",
"%reload_ext tensorboard\n",
"%tensorboard --logdir lightning_logs/"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "RltpFGS-s0M1"
},
"source": [
"<code style=\"color:#792ee5;\">\n",
" <h1> <strong> Congratulations - Time to Join the Community! </strong> </h1>\n",
"</code>\n",
"\n",
"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",
"\n",
"### Star [Lightning](https://github.com/PyTorchLightning/pytorch-lightning) on GitHub\n",
"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",
"\n",
"* Please, star [Lightning](https://github.com/PyTorchLightning/pytorch-lightning)\n",
"\n",
"### Join our [Slack](https://join.slack.com/t/pytorch-lightning/shared_invite/zt-f6bl2l0l-JYMK3tbAgAmGRrlNr00f1A)!\n",
"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",
"\n",
"### Interested by SOTA AI models ! Check out [Bolt](https://github.com/PyTorchLightning/pytorch-lightning-bolts)\n",
"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",
"\n",
"* Please, star [Bolt](https://github.com/PyTorchLightning/pytorch-lightning-bolts)\n",
"\n",
"### Contributions !\n",
"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",
"\n",
"* [Lightning good first issue](https://github.com/PyTorchLightning/pytorch-lightning/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)\n",
"* [Bolt good first issue](https://github.com/PyTorchLightning/pytorch-lightning-bolts/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)\n",
"* You can also contribute your own notebooks with useful examples !\n",
"\n",
"### Great thanks from the entire Pytorch Lightning Team for your interest !\n",
"\n",
"<img src=\"https://github.com/PyTorchLightning/pytorch-lightning/blob/master/docs/source/_static/images/logo.png?raw=true\" width=\"800\" height=\"200\" />"
]
}
]
}