update PyTorch Lightning usage

Author: rasbt

ch13/ch13_part3_lightning.ipynb

@@ -166,7 +166,6 @@
     "        super().__init__()\n",
     "        \n",
     "        # new PL attributes:\n",
-    "        \n",
     "        if parse_version(torchmetrics_version) > parse_version(\"0.8\"):\n",
     "            self.train_acc = Accuracy(task=\"multiclass\", num_classes=10)\n",
     "            self.valid_acc = Accuracy(task=\"multiclass\", num_classes=10)\n",
@@ -201,23 +200,35 @@
     "        self.log(\"train_loss\", loss, prog_bar=True)\n",
     "        return loss\n",
     "\n",
-    "    def training_epoch_end(self, outs):\n",
-    "        self.log(\"train_acc\", self.train_acc.compute())\n",
-    "        self.train_acc.reset()\n",
-    "    \n",
-    "    def validation_step(self, batch, batch_idx):\n",
-    "        x, y = batch\n",
-    "        logits = self(x)\n",
-    "        loss = nn.functional.cross_entropy(logits, y)\n",
-    "        preds = torch.argmax(logits, dim=1)\n",
-    "        self.valid_acc.update(preds, y)\n",
-    "        self.log(\"valid_loss\", loss, prog_bar=True)\n",
-    "        return loss\n",
-    "    \n",
-    "    def validation_epoch_end(self, outs):\n",
-    "        self.log(\"valid_acc\", self.valid_acc.compute(), prog_bar=True)\n",
-    "        self.valid_acc.reset()\n",
+    "    # Conditionally define epoch end methods based on PyTorch Lightning version\n",
+    "    if parse_version(pl.__version__) >= parse_version(\"2.0\"):\n",
+    "        # For PyTorch Lightning 2.0 and above\n",
+    "        def on_training_epoch_end(self):\n",
+    "            self.log(\"train_acc\", self.train_acc.compute())\n",
+    "            self.train_acc.reset()\n",
+    "\n",
+    "        def on_validation_epoch_end(self):\n",
+    "            self.log(\"valid_acc\", self.valid_acc.compute())\n",
+    "            self.valid_acc.reset()\n",
+    "\n",
+    "        def on_test_epoch_end(self):\n",
+    "            self.log(\"test_acc\", self.test_acc.compute())\n",
+    "            self.test_acc.reset()\n",
     "\n",
+    "    else:\n",
+    "        # For PyTorch Lightning < 2.0\n",
+    "        def training_epoch_end(self, outs):\n",
+    "            self.log(\"train_acc\", self.train_acc.compute())\n",
+    "            self.train_acc.reset()\n",
+    "\n",
+    "        def validation_epoch_end(self, outs):\n",
+    "            self.log(\"valid_acc\", self.valid_acc.compute())\n",
+    "            self.valid_acc.reset()\n",
+    "\n",
+    "        def test_epoch_end(self, outs):\n",
+    "            self.log(\"test_acc\", self.test_acc.compute())\n",
+    "            self.test_acc.reset()\n",
+    "    \n",
     "    def test_step(self, batch, batch_idx):\n",
     "        x, y = batch\n",
     "        logits = self(x)\n",
@@ -1101,7 +1112,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.6"
+   "version": "3.10.14"
   }
  },
  "nbformat": 4,