5 implement continual learning benchmark evaluation

dev_notebook.ipynb

@@ -10,6 +10,8 @@
     "import torch\n",
     "from train import *\n",
     "from data.utils import *\n",
+    "from continuum import ClassIncremental\n",
+    "from continuum.tasks import split_train_val\n",
     "DEVICE = torch.device(\"cpu\")"
    ]
   },
@@ -31,16 +33,38 @@
     "    'dp': True,\n",
     "    'C': .02,\n",
     "    'sigma': .1\n",
+    "    'increment': 5\n",
     "}\n",
     "\n",
-    "model = load_model(dataset='mnist')\n",
-    "train_results = train(\n",
-    "    model, trainloader_mnist, device=DEVICE, n_epochs=10, opt_params=params\n",
-    ")\n",
     "\n",
-    "test_results = test(\n",
-    "    model, testloader_mnist, device=DEVICE\n",
-    ")"
+    "for learning in ['central', 'single', 'multi']:\n",
+    "    if learning == 'central':\n",
+    "        trainloader = torch.utils.data.DataLoader(\n",
+    "            dataset=trainset_mnist, batch_size=params['batch_size'], shuffle=True, drop_last=True)\n",
+    "\n",
+    "        testloader = torch.utils.data.DataLoader(\n",
+    "            dataset=testset_mnist, batch_size=params['batch_size'], shuffle=False, drop_last=True)\n",
+    "        model = load_model(dataset='mnist')\n",
+    "        train_results = train( model=model, trainloader=trainloader,\n",
+    "              device=DEVICE, opt_params=params)\n",
+    "        test_results = test(model, testloader_mnist, device=DEVICE)\n",
+    "\n",
+    "    else:\n",
+    "        # Incremental Learning(Changing params will change how many classes at a time to learn\n",
+    "        scenario = ClassIncremental(trainset_mnist, increment=(1 if learning == 'single' else params['increment']))\n",
+    "        print(f\"Number of classes: {scenario.nb_classes}.\")\n",
+    "        print(f\"Number of tasks: {scenario.nb_tasks}.\")\n",
+    "        for task_id, train_taskset in enumerate(scenario):\n",
+    "            train_taskset, val_taskset = split_train_val(train_taskset, val_split=0)\n",
+    "            trainloader = torch.utils.data.DataLoader(dataset=train_taskset, batch_size=params['batch_size'],\n",
+    "                                                      shuffle=True, drop_last=True)\n",
+    "            model = load_model(dataset='mnist')\n",
+    "            train_results = train( model=model, trainloader=trainloader,\n",
+    "                   device=DEVICE, opt_params=params)\n",
+    "            test_results = test(model, testloader_mnist, device=DEVICE)\n",
+    "\n",
+    "\n",
+    "\n"
    ],
    "metadata": {
     "collapsed": false,

train.py

@@ -52,27 +52,27 @@ def train(
         for b_ix, sample in enumerate(pbar):
             pbar.set_description(f'Epoch {epoch}: Training...')
 
-            image, label = sample[0].to(device), sample[1].to(device)
-
-            # forward + backward + optimize
-            optimizer.zero_grad()
-            y_hat = model(image)
-            loss = criterion(y_hat, label)
-            loss.backward()
-
-            if use_dp: 
-                for param in model.parameters():     
-                    clip_grad_norm_(param.grad, max_norm=C)
-                    param = param - lr * param.grad
-                    param += torch.normal(mean=torch.zeros(param.shape), std=sigma * C)
-            else: 
-                optimizer.step()
-
-            # Collect statistics
-            running_loss += loss.item()
-            _, predicted = torch.max(y_hat.data, 1)
-            total += label.size(0)
-            running_acc += (predicted == label).sum().item()
+            images, label = sample[0].to(device), sample[1].to(device)
+            for image in images:
+                # forward + backward + optimize
+                optimizer.zero_grad()
+                y_hat = model(image)
+                loss = criterion(y_hat, label)
+                loss.backward()
+
+                if use_dp:
+                    for param in model.parameters():
+                        clip_grad_norm_(param.grad, max_norm=C)
+                        param = param - lr * param.grad
+                        param += torch.normal(mean=torch.zeros(param.shape), std=sigma * C)
+                else:
+                    optimizer.step()
+
+                # Collect statistics
+                running_loss += loss.item()
+                _, predicted = torch.max(y_hat.data, 1)
+                total += label.size(0)
+                running_acc += (predicted == label).sum().item()
 
         results.append((running_loss/total, running_acc/total))