integrations-and-supported-tools/pytorch-lightning/notebooks/Neptune_PyTorch_Lightning.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Neptune + PyTorch Lightning](https://neptune.ai/wp-content/uploads/2023/09/lightning.svg)\n",
    "\n",
    "# Neptune + PyTorch Lightning\n",
    "\n",
    "<a target=\"_blank\" href=\"https://lightning.ai/neptuneai/studios/neptune-lightning\">\n",
    "  <img src=\"https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg\" height=\"20\" alt=\"Open In Studio\"/>\n",
    "</a> <a target=\"_blank\" href=\"https://colab.research.google.com/github/neptune-ai/examples/blob/main/integrations-and-supported-tools/pytorch-lightning/notebooks/Neptune_PyTorch_Lightning.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open in Colab\"/>\n",
    "</a><a target=\"_blank\" href=\"https://github.com/neptune-ai/examples/blob/main/integrations-and-supported-tools/pytorch-lightning/notebooks/Neptune_PyTorch_Lightning.ipynb\">\n",
    "  <img alt=\"Open in GitHub\" src=\"https://img.shields.io/badge/Open_in_GitHub-blue?logo=github&labelColor=black\">\n",
    "</a><a target=\"_blank\" href=\"https://app.neptune.ai/o/showcase/org/pytorch-lightning/runs/details?viewId=9c2ac844-b98c-40f8-a776-4d5db0547545&detailsTab=dashboard&dashboardId=Overview-9c2ac3da-751f-4048-af0a-76707b4e036b&shortId=PTL-8&type=run\"> \n",
    "  <img alt=\"Explore in Neptune\" src=\"https://neptune.ai/wp-content/uploads/2024/01/neptune-badge.svg\">\n",
    "</a><a target=\"_blank\" href=\"https://docs.neptune.ai/integrations/lightning/\">\n",
    "  <img alt=\"View tutorial in docs\" src=\"https://neptune.ai/wp-content/uploads/2024/01/docs-badge-2.svg\">\n",
    "</a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "header",
     "comment"
    ]
   },
   "source": [
    "## Introduction\n",
    "\n",
    "This guide will show you how to:\n",
    "\n",
    "* Create a `NeptuneLogger()` instance\n",
    "* Use the logger to log metadata to Neptune"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before you start\n",
    "\n",
    "This notebook example lets you try out Neptune as an anonymous user, with zero setup.\n",
    "\n",
    "If you want to see the example logged to your own workspace instead:\n",
    "\n",
    "  1. Create a Neptune account. [Register &rarr;](https://neptune.ai/register)\n",
    "  1. Create a Neptune project that you will use for tracking metadata. For instructions, see [Creating a project](https://docs.neptune.ai/setup/creating_project) in the Neptune docs."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Neptune and dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -q -U neptune lightning torchvision \"scipy<1.12\"\n",
    "! pip install -q -U --user pydantic scikit-learn"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: If running on Google Colab, restart the kernel and continue execution from the next cell to avoid a `ContextualVersionConflict` error.\n",
    "\n",
    "This error is caused by Colab coming with `future==0.16.0` preinstalled, while `torchvision` updates this to a newer version."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "header"
    ]
   },
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "code"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pytorch_lightning as pl\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from lightning.pytorch.loggers.neptune import NeptuneLogger\n",
    "from neptune.types import File\n",
    "from sklearn.metrics import accuracy_score\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import MNIST"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "tags": [
     "header"
    ]
   },
   "source": [
    "## Define hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab_type": "code",
    "tags": [
     "code"
    ]
   },
   "outputs": [],
   "source": [
    "params = {\n",
    "    \"batch_size\": 64,\n",
    "    \"linear\": 32,\n",
    "    \"lr\": 0.001,\n",
    "    \"decay_factor\": 0.8,\n",
    "    \"max_epochs\": 3,\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "header"
    ]
   },
   "source": [
    "## Define Lightning model\n",
    "\n",
    "To log metrics to Neptune, define a Lightning model with both inbuilt logging (`self.log()`) and custom logging (`neptune_logger.experiment`).\n",
    "\n",
    "`neptune_logger.experiment` is a reference to the Neptune run object created by the `NeptuneLogger()` instance. You can use it to log outside the `prefix` passed to `NeptuneLogger()`. It accepts Neptune's logging methods such as `append()` and `upload()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "code"
    ]
   },
   "outputs": [],
   "source": [
    "# (neptune) define model with logging (self.log)\n",
    "class LitModel(pl.LightningModule):\n",
    "    def __init__(self, linear, learning_rate, decay_factor):\n",
    "        super().__init__()\n",
    "        self.training_step_outputs = []\n",
    "        self.validation_step_outputs = []\n",
    "        self.test_step_outputs = []\n",
    "        self.linear = linear\n",
    "        self.learning_rate = learning_rate\n",
    "        self.decay_factor = decay_factor\n",
    "        self.train_img_max = 10\n",
    "        self.train_img = 0\n",
    "        self.layer_1 = torch.nn.Linear(28 * 28, linear)\n",
    "        self.layer_2 = torch.nn.Linear(linear, 20)\n",
    "        self.layer_3 = torch.nn.Linear(20, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.layer_1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.layer_2(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.layer_3(x)\n",
    "        return x\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)\n",
    "        scheduler = LambdaLR(optimizer, lambda epoch: self.decay_factor**epoch)\n",
    "        return [optimizer], [scheduler]\n",
    "\n",
    "    def training_step(self, batch):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = F.cross_entropy(y_hat, y)\n",
    "        self.log(\"train/batch/loss\", loss, prog_bar=False)  # Log training batch loss to Neptune\n",
    "\n",
    "        y_true = y.cpu().detach().numpy()\n",
    "        y_pred = y_hat.argmax(axis=1).cpu().detach().numpy()\n",
    "        acc = accuracy_score(y_true, y_pred)\n",
    "        self.log(\"train/batch/acc\", acc)  # Log training batch accuracy to Neptune\n",
    "        self.training_step_outputs.append({\"loss\": loss, \"y_true\": y_true, \"y_pred\": y_pred})\n",
    "\n",
    "        return {\"loss\": loss, \"y_true\": y_true, \"y_pred\": y_pred}\n",
    "\n",
    "    def on_train_epoch_end(self):\n",
    "        loss = np.array([])\n",
    "        y_true = np.array([])\n",
    "        y_pred = np.array([])\n",
    "        for results_dict in self.training_step_outputs:\n",
    "            loss = np.append(loss, results_dict[\"loss\"].detach().numpy())\n",
    "            y_true = np.append(y_true, results_dict[\"y_true\"])\n",
    "            y_pred = np.append(y_pred, results_dict[\"y_pred\"])\n",
    "        acc = accuracy_score(y_true, y_pred)\n",
    "        self.log(\"train/epoch/loss\", loss.mean())  # Log training epoch accuracy to Neptune\n",
    "        self.log(\"train/epoch/acc\", acc)  # Log training epoch loss to Neptune\n",
    "        self.training_step_outputs.clear()  # free memory\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = F.cross_entropy(y_hat, y)\n",
    "\n",
    "        y_true = y.cpu().detach().numpy()\n",
    "        y_pred = y_hat.argmax(axis=1).cpu().detach().numpy()\n",
    "\n",
    "        self.validation_step_outputs.append({\"loss\": loss, \"y_true\": y_true, \"y_pred\": y_pred})\n",
    "\n",
    "        return {\"loss\": loss, \"y_true\": y_true, \"y_pred\": y_pred}\n",
    "\n",
    "    def on_validation_epoch_end(self):\n",
    "        loss = np.array([])\n",
    "        y_true = np.array([])\n",
    "        y_pred = np.array([])\n",
    "        for results_dict in self.validation_step_outputs:\n",
    "            loss = np.append(loss, results_dict[\"loss\"].detach().numpy())\n",
    "            y_true = np.append(y_true, results_dict[\"y_true\"])\n",
    "            y_pred = np.append(y_pred, results_dict[\"y_pred\"])\n",
    "        acc = accuracy_score(y_true, y_pred)\n",
    "        self.log(\"val/loss\", loss.mean())  # Log validation loss to Neptune\n",
    "        self.log(\"val/acc\", acc)  # Log validation accuracy to Neptune\n",
    "        self.validation_step_outputs.clear()  # free memory\n",
    "\n",
    "    def test_step(self, batch):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = F.cross_entropy(y_hat, y)\n",
    "\n",
    "        y_true = y.cpu().detach().numpy()\n",
    "        y_pred = y_hat.argmax(axis=1).cpu().detach().numpy()\n",
    "\n",
    "        # Log misclassified test images to Neptune\n",
    "        for j in np.where(np.not_equal(y_true, y_pred))[0]:\n",
    "            img = np.squeeze(x[j].cpu().detach().numpy())\n",
    "            img[img < 0] = 0\n",
    "            img = img / np.amax(img)\n",
    "            neptune_logger.experiment[f\"{neptune_logger._prefix}/test/misclassified_images\"].append(\n",
    "                File.as_image(img),  # You can log numpy arrays as images using File.as_image()\n",
    "                description=f\"y_pred={y_pred[j]}, y_true={y_true[j]}\",\n",
    "            )\n",
    "\n",
    "        self.test_step_outputs.append({\"loss\": loss, \"y_true\": y_true, \"y_pred\": y_pred})\n",
    "\n",
    "        return {\"loss\": loss, \"y_true\": y_true, \"y_pred\": y_pred}\n",
    "\n",
    "    def on_test_epoch_end(self):\n",
    "        loss = np.array([])\n",
    "        y_true = np.array([])\n",
    "        y_pred = np.array([])\n",
    "        for results_dict in self.test_step_outputs:\n",
    "            loss = np.append(loss, results_dict[\"loss\"].detach().numpy())\n",
    "            y_true = np.append(y_true, results_dict[\"y_true\"])\n",
    "            y_pred = np.append(y_pred, results_dict[\"y_pred\"])\n",
    "        acc = accuracy_score(y_true, y_pred)\n",
    "        self.log(\"test/loss\", loss.mean())  # Log test loss to Neptune\n",
    "        self.log(\"test/acc\", acc)  # Log test accuracy to Neptune\n",
    "        self.validation_step_outputs.clear()  # free memory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LitModel(\n",
    "    linear=params[\"linear\"],\n",
    "    learning_rate=params[\"lr\"],\n",
    "    decay_factor=params[\"decay_factor\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define DataModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNISTDataModule(pl.LightningDataModule):\n",
    "    def __init__(self, batch_size, normalization_vector):\n",
    "        super().__init__()\n",
    "        self.batch_size = batch_size\n",
    "        self.normalization_vector = normalization_vector\n",
    "        self.mnist_train = None\n",
    "        self.mnist_val = None\n",
    "        self.mnist_test = None\n",
    "\n",
    "    def prepare_data(self):\n",
    "        MNIST(os.getcwd(), train=True, download=True)\n",
    "        MNIST(os.getcwd(), train=False, download=True)\n",
    "\n",
    "    def setup(self, stage):\n",
    "        # transforms\n",
    "        transform = transforms.Compose(\n",
    "            [\n",
    "                transforms.ToTensor(),\n",
    "                transforms.Normalize(self.normalization_vector[0], self.normalization_vector[1]),\n",
    "            ]\n",
    "        )\n",
    "        if stage == \"fit\":\n",
    "            mnist_train = MNIST(os.getcwd(), train=True, transform=transform)\n",
    "            self.mnist_train, self.mnist_val = random_split(mnist_train, [55000, 5000])\n",
    "        if stage == \"test\":\n",
    "            self.mnist_test = MNIST(os.getcwd(), train=False, transform=transform)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.mnist_train, batch_size=self.batch_size, num_workers=0)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.mnist_val, batch_size=self.batch_size, num_workers=0)\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return DataLoader(self.mnist_test, batch_size=self.batch_size, num_workers=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize DataModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dm = MNISTDataModule(\n",
    "    normalization_vector=((0.1307,), (0.3081,)),\n",
    "    batch_size=params[\"batch_size\"],\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "header"
    ]
   },
   "source": [
    "## Create a NeptuneLogger() instance\n",
    "\n",
    "To create a new run for tracking the metadata, you tell Neptune who you are (`api_token`) and where to send the data (`project`).\n",
    "\n",
    "You can use the default code cell below to create an anonymous run in a public project.  \n",
    "**Note**: Public projects are cleaned regularly, so anonymous runs are only stored temporarily.\n",
    "\n",
    "### Log to your own project instead\n",
    "\n",
    "Replace the code below with the following:\n",
    "\n",
    "```python\n",
    "from getpass import getpass\n",
    "\n",
    "neptune_logger = NeptuneLogger(\n",
    "    project=\"workspace-name/project-name\",  # replace with your own (see instructions below)\n",
    "    api_token=getpass(\"Enter your Neptune API token: \"),\n",
    "    tags=[\"notebook\"],\n",
    "    log_model_checkpoints=False,  # Update to True to log model checkpoints\n",
    "    dependencies=\"infer\",  # You can also pass the path to your dependencies file or omit this parameter\n",
    ")\n",
    "```\n",
    "\n",
    "To find your API token and full project name:\n",
    "\n",
    "1. [Log in to Neptune](https://app.neptune.ai/).\n",
    "1. In the bottom-left corner, expand your user menu and select **Get your API token**.\n",
    "1. You can copy the project path from the project settings (select **Details & privacy**).\n",
    "\n",
    "For more help, see [Setting Neptune credentials](https://docs.neptune.ai/setup/setting_credentials) in the Neptune docs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "code"
    ]
   },
   "outputs": [],
   "source": [
    "from neptune import ANONYMOUS_API_TOKEN\n",
    "\n",
    "neptune_logger = NeptuneLogger(\n",
    "    api_key=ANONYMOUS_API_TOKEN,\n",
    "    project=\"common/pytorch-lightning\",\n",
    "    tags=[\"notebook\"],\n",
    "    log_model_checkpoints=False,  # Update to True to log model checkpoints\n",
    "    dependencies=\"infer\",  # You can also pass the path to your dependencies file or omit this parameter\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "header"
    ]
   },
   "source": [
    "## Initialize a trainer and pass neptune_logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "code"
    ]
   },
   "outputs": [],
   "source": [
    "trainer = pl.Trainer(\n",
    "    logger=neptune_logger,\n",
    "    max_epochs=params[\"max_epochs\"],\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "header"
    ]
   },
   "source": [
    "## Log hyperparameters to the run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "neptune_logger.log_hyperparams(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Log model summary to the run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neptune_logger.log_model_summary(model=model, max_depth=-1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "header"
    ]
   },
   "source": [
    "## Train and test the model and log metadata to the run live"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "code"
    ]
   },
   "outputs": [],
   "source": [
    "trainer.fit(model, datamodule=dm)\n",
    "trainer.test(model, datamodule=dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see the metadata being logged live, open the run in the Neptune app.\n",
    "\n",
    "You can ignore any \"X-coordinates (step) must be strictly increasing\" errors."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stop logging\n",
    "\n",
    "Once you are done logging, you should stop the run using the `stop()` method, which is a method of the run object. The run object itself is part of the `neptune_logger`.\n",
    "This is needed only while logging from a notebook environment. While logging through a script, Neptune automatically stops tracking once the script has completed execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "neptune_logger.experiment.stop()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "header"
    ]
   },
   "source": [
    "## Analyze logged metadata in the Neptune app"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "comment"
    ]
   },
   "source": [
    "To explore the metadata in Neptune, follow the link in the console output.\n",
    "\n",
    "It looks something like this: https://app.neptune.ai/showcase/pytorch-lightning/e/PTL-3"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Neptune_PyTorch_Lightning.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "py38",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "vscode": {
   "interpreter": {
    "hash": "a9715cf0b0024f6e1c62cb31a4f1f43970eb41991212681878768b4bfe53050a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}