395 lines
15 KiB
Plaintext
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 pytorch-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/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\" />"
|
|
]
|
|
}
|
|
]
|
|
}
|