From c422df51a286928b3c34b445b8f835c83d5902fa Mon Sep 17 00:00:00 2001 From: Gao Enhao Date: Fri, 22 Dec 2023 21:36:39 +0800 Subject: [PATCH] [ENH] improve mnist performance --- abl/learning/basic_nn.py | 7 +++--- docs/Intro/Datasets.rst | 2 +- examples/mnist_add/main.py | 32 +++++++++++++++--------- examples/mnist_add/mnist_add.ipynb | 39 +++++++++++++++++------------- 4 files changed, 47 insertions(+), 33 deletions(-) diff --git a/abl/learning/basic_nn.py b/abl/learning/basic_nn.py index 641bd7f..721576a 100644 --- a/abl/learning/basic_nn.py +++ b/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() diff --git a/docs/Intro/Datasets.rst b/docs/Intro/Datasets.rst index 732039c..4deb25d 100644 --- a/docs/Intro/Datasets.rst +++ b/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 -------------- diff --git a/examples/mnist_add/main.py b/examples/mnist_add/main.py index 0bf4a35..53f4f05 100644 --- a/examples/mnist_add/main.py +++ b/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, diff --git a/examples/mnist_add/mnist_add.ipynb b/examples/mnist_add/mnist_add.ipynb index 4760d06..6161131 100644 --- a/examples/mnist_add/mnist_add.ipynb +++ b/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)" ] }