[ENH] improve mnist performance

abl/learning/basic_nn.py

@@ -24,6 +24,9 @@ class BasicNN:
         The loss function used for training.
     optimizer : torch.optim.Optimizer
         The optimizer used for training.
+    scheduler : torch.optim.lr_scheduler.LRScheduler
+        The learning rate scheduler used for training, which will be called
+        at the end of each run of the ``fit`` method, by default None.
     device : torch.device, optional
         The device on which the model will be trained or used for prediction,
         by default torch.device("cpu").
@@ -216,9 +219,7 @@ class BasicNN:
         for data, target in data_loader:
             data, target = data.to(device), target.to(device)
             out = model(data)
-            proba = torch.nn.functional.softmax(out, dim=1)
-            entropy = -torch.sum(proba * torch.log(proba + 1e-5), dim=1).mean()
-            loss = loss_fn(out, target) - 0.3 * entropy
+            loss = loss_fn(out, target)
 
             optimizer.zero_grad()
             loss.backward()

docs/Intro/Datasets.rst

@@ -58,7 +58,7 @@ As an illustration, in the MNIST Addition task, the data are organized as follow
    :alt: alternate text
    :scale: 55%
 
-where each sublist in ``X``, e.g., |data_example|, is a data example and each image in the sublist, e.g. |instance|, is a data instance.
+where each sublist in ``X``, e.g., |data_example|, is a data example and each image in the sublist, e.g., |instance|, is a data instance.
 
 Data Structure
 --------------

examples/mnist_add/main.py

@@ -1,17 +1,17 @@
-import os
-import os.path as osp
 import argparse
+import os.path as osp
 
 import torch
 from torch import nn
+from torch.optim import RMSprop, lr_scheduler
 
-from examples.mnist_add.datasets import get_dataset
-from examples.models.nn import LeNet5
-from abl.learning import ABLModel, BasicNN
-from abl.reasoning import KBBase, GroundKB, PrologKB, Reasoner
+from abl.bridge import SimpleBridge
 from abl.data.evaluation import ReasoningMetric, SymbolAccuracy
+from abl.learning import ABLModel, BasicNN
+from abl.reasoning import GroundKB, KBBase, PrologKB, Reasoner
 from abl.utils import ABLLogger, print_log
-from abl.bridge import SimpleBridge
+from examples.mnist_add.datasets import get_dataset
+from examples.models.nn import LeNet5
 
 
 class AddKB(KBBase):
@@ -49,10 +49,10 @@ def main():
         "--batch-size", type=int, default=32, help="base model batch size (default : 32)"
     )
     parser.add_argument(
-        "--loops", type=int, default=5, help="number of loop iterations (default : 5)"
+        "--loops", type=int, default=1, help="number of loop iterations (default : 1)"
     )
     parser.add_argument(
-        "--segment_size", type=int or float, default=1 / 3, help="segment size (default : 1/3)"
+        "--segment_size", type=int or float, default=0.01, help="segment size (default : 0.01)"
     )
     parser.add_argument("--save_interval", type=int, default=1, help="save interval (default : 1)")
     parser.add_argument(
@@ -64,7 +64,7 @@ def main():
     parser.add_argument(
         "--require-more-revision",
         type=int,
-        default=5,
+        default=0,
         help="require more revision in reasoner (default : 0)",
     )
     kb_type = parser.add_mutually_exclusive_group()
@@ -84,16 +84,24 @@ def main():
     ### Building the Learning Part
     # Build necessary components for BasicNN
     cls = LeNet5(num_classes=10)
-    loss_fn = nn.CrossEntropyLoss()
-    optimizer = torch.optim.RMSprop(cls.parameters(), lr=args.lr, alpha=args.alpha)
+    loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)
+    optimizer = RMSprop(cls.parameters(), lr=args.lr, alpha=args.alpha)
     use_cuda = not args.no_cuda and torch.cuda.is_available()
     device = torch.device("cuda" if use_cuda else "cpu")
+    scheduler = lr_scheduler.OneCycleLR(
+        optimizer,
+        max_lr=args.lr,
+        pct_start=0.2,
+        epochs=args.loops,
+        steps_per_epoch=int(1 / args.segment_size),
+    )
 
     # Build BasicNN
     base_model = BasicNN(
         cls,
         loss_fn,
         optimizer,
+        scheduler=scheduler,
         device=device,
         batch_size=args.batch_size,
         num_epochs=args.epochs,

examples/mnist_add/mnist_add.ipynb

@@ -13,7 +13,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -22,6 +22,9 @@
     "import torch\n",
     "import torch.nn as nn\n",
     "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from torch.optim import RMSprop, lr_scheduler\n",
+    "\n",
     "from examples.mnist_add.datasets import get_dataset\n",
     "from examples.models.nn import LeNet5\n",
     "from abl.learning import ABLModel, BasicNN\n",
@@ -42,7 +45,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -61,7 +64,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -110,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -170,20 +173,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "cls = LeNet5(num_classes=10)\n",
-    "loss_fn = nn.CrossEntropyLoss()\n",
-    "optimizer = torch.optim.RMSprop(cls.parameters(), lr=0.001, alpha=0.9)\n",
+    "loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)\n",
+    "optimizer = RMSprop(cls.parameters(), lr=0.001, alpha=0.9)\n",
     "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=0.001, pct_start=0.1, total_steps=100)\n",
     "\n",
     "base_model = BasicNN(\n",
     "    cls,\n",
     "    loss_fn,\n",
     "    optimizer,\n",
-    "    device,\n",
+    "    scheduler=scheduler,\n",
+    "    device=device,\n",
     "    batch_size=32,\n",
     "    num_epochs=1,\n",
     ")"
@@ -198,7 +203,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -229,7 +234,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -245,7 +250,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -295,7 +300,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -319,7 +324,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -352,7 +357,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -385,7 +390,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -404,7 +409,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -429,7 +434,7 @@
     "log_dir = ABLLogger.get_current_instance().log_dir\n",
     "weights_dir = osp.join(log_dir, \"weights\")\n",
     "\n",
-    "bridge.train(train_data, loops=5, segment_size=1/3, save_interval=1, save_dir=weights_dir)\n",
+    "bridge.train(train_data, loops=1, segment_size=0.01, save_interval=1, save_dir=weights_dir)\n",
     "bridge.test(test_data)"
    ]
   }