train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5df58edd-79d3-426b-8b5e-3d7e4fb180a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "import models.capsulecifar\n",
    "import tqdm\n",
    "import logging\n",
    "# import wandb\n",
    "from models.model import *\n",
    "from torch import optim\n",
    "from utils.get_args import get_arg\n",
    "from utils.get_datasets import get_dataset\n",
    "from utils.get_losses import CapsuleLoss, OrthogonalProjectionLoss\n",
    "from torch.utils.data import DataLoader\n",
    "from pathlib import Path\n",
    "# wandb.init(mode='offline')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f91049c2-070b-4232-8512-b6a2f22c7827",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = 10\n",
    "def test_capsule(model, device, test_loader, criterion):\n",
    "    model.eval()\n",
    "    correct, total = 0, 0\n",
    "    for images, labels in test_loader:\n",
    "        # Add channels = 1\n",
    "        images = images.to(device)\n",
    "        # Categogrical encoding\n",
    "        labels = torch.eye(num_classes).index_select(dim=0, index=labels).to(device)\n",
    "        logits, reconstructions, primary_caps_output, digit_caps_output, c, b = model(images)\n",
    "        pred_labels = torch.argmax(logits, dim=1)\n",
    "        correct += torch.sum(pred_labels == torch.argmax(labels, dim=1)).item()\n",
    "        total += len(labels)\n",
    "\n",
    "    print('Accuracy: {}'.format(correct / total))\n",
    "    return correct / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1ed825fb-9d25-4b9b-88c1-d07186c208f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "args = get_arg()\n",
    "# python train_capsule.py --dataset=cifar --epochs=60 --batch_size=32 --optimizer=sgd --lr=0.005 --momentum=0.9 --weight_decay=5e-4\n",
    "args.dataset = 'cifar'\n",
    "args.epochs = 300\n",
    "args.batch_size = 32\n",
    "args.optimizer = 'sgd'\n",
    "args.lr = 0.005\n",
    "args.momentum = 0.9\n",
    "args.weight_decay = 5e-4\n",
    "device = 'cuda'\n",
    "dir_checkpoint = Path(f'./checkpoints/')\n",
    "# 1. Dataloader\n",
    "train_dataset, test_dataset = get_dataset(args)\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, shuffle=True)\n",
    "test_loader = DataLoader(dataset=test_dataset, batch_size=args.batch_size, shuffle=False)\n",
    "# 2. Initialize logging\n",
    "logging.info(f'''Starting training:\n",
    "    Epochs:          {args.epochs}\n",
    "    Batch size:      {args.batch_size}\n",
    "    Learning rate:   {args.lr}\n",
    "    Optimizer:       {args.optimizer}\n",
    "    Sigma:           {args.sigma}\n",
    "    Momentum:        {args.momentum}\n",
    "    Weight decay:    {args.weight_decay}\n",
    "''')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e324b6eb-81a0-4a24-b493-c5a4f1238227",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Let's use 2 GPUs!\n",
      "==> Building model..\n"
     ]
    }
   ],
   "source": [
    "# 3. Obtain model\n",
    "model = models.capsulecifar.CapsNet(device=device)\n",
    "if torch.cuda.device_count() > 1:\n",
    "    print(\"Let's use\", torch.cuda.device_count(), \"GPUs!\")\n",
    "    model = nn.DataParallel(model)\n",
    "model.to(device)\n",
    "# 4. Set up the optimizer, the loss, the learning rate scheduler\n",
    "if args.optimizer == 'sgd':\n",
    "    optimizer = optim.SGD(params=model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)\n",
    "elif args.optimizer == 'adam':\n",
    "    optimizer = optim.Adam(params=model.parameters(), lr=args.lr, weight_decay=args.weight_decay)\n",
    "else:\n",
    "    print('我还没写')\n",
    "\n",
    "criterion = CapsuleLoss()\n",
    "criterion2 = OrthogonalProjectionLoss()\n",
    "op_weight = 1.0\n",
    "scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.96)\n",
    "global_step = 0\n",
    "best_accuracy = 0.0\n",
    "print('==> Building model..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce796d3b-7606-4406-812b-c2d949edbb86",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1 / 300:   9%|█▎             | 133/1563 [00:02<00:23, 61.57batch/s, Margin Loss=11]"
     ]
    }
   ],
   "source": [
    "# 5. Traning\n",
    "for epoch in range(1, args.epochs + 1):\n",
    "    model.train()\n",
    "    epoch_loss = 0\n",
    "    total_m_loss = 0.0\n",
    "    total_r_loss = 0.0\n",
    "    total_o_loss = 0.0\n",
    "    with tqdm.tqdm(total=len(train_loader), desc=f\"Epoch {epoch} / {args.epochs}\", unit=\"batch\") as pbar:\n",
    "        batch_id = 1\n",
    "        correct, total, total_loss = 0, 0, 0.\n",
    "        for batch_idx, (inputs, labels) in enumerate(train_loader):\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            labels_op = labels\n",
    "            optimizer.zero_grad()\n",
    "            eye = torch.eye(num_classes).to(device)\n",
    "            labels = eye.index_select(dim=0, index=labels)\n",
    "            logits, reconstruction, primary_caps_output, digit_caps_output, c, b = model(inputs)\n",
    "            loss = criterion(inputs, labels, logits, reconstruction)\n",
    "            labels_op = labels_op.view(-1)\n",
    "            primary_caps_output_op = primary_caps_output.view(args.batch_size, -1)\n",
    "            loss_op = criterion2(primary_caps_output_op, labels_op)\n",
    "            margin_loss = criterion.margin_loss.item()\n",
    "            reconstruction_loss = criterion.reconstruction_loss.item()\n",
    "            correct += torch.sum(torch.argmax(logits, dim=1) == torch.argmax(labels, dim=1)).item()\n",
    "            total += len(labels)\n",
    "            accuracy = correct / total\n",
    "            total_loss += loss + op_weight * loss_op\n",
    "            total_m_loss += margin_loss\n",
    "            total_r_loss += reconstruction_loss\n",
    "            total_o_loss += loss_op\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            batch_id += 1\n",
    "            pbar.set_postfix({\n",
    "                \"Margin Loss\": total_m_loss / (batch_idx + 1),\n",
    "                # \"Reconstruction Loss\": total_r_loss / (batch_idx + 1),\n",
    "                # \"OPLoss\": total_o_loss / (batch_idx + 1)\n",
    "            })\n",
    "            pbar.update(1)\n",
    "    scheduler.step()\n",
    "    print('------Test Result:------')\n",
    "    test_accuracy = test_capsule(model=model, device=device, test_loader=test_loader, criterion=criterion)\n",
    "\n",
    "    if test_accuracy > best_accuracy:\n",
    "        best_accuracy = test_accuracy\n",
    "        state_dict = model.state_dict()\n",
    "        best_model_path = str(\n",
    "            dir_checkpoint/ '{}'.format('Capsule') / '{}'.format(args.dataset) / 'common' / 'best_model_{}.pth'.format(args.dataset))\n",
    "        torch.save(state_dict, best_model_path)\n",
    "        logging.info(f'Best model (accuracy: {best_accuracy:.2f}%) saved to {best_model_path}')\n",
    "    # Saved model's weight as pth file\n",
    "    Path(dir_checkpoint).mkdir(parents=True, exist_ok=True)\n",
    "    state_dict = model.state_dict()\n",
    "    torch.save(state_dict,\n",
    "               str(dir_checkpoint / '{}'.format('Capsule') / '{}'.format(args.dataset) / 'common' / 'checkpoint_{}_epoch{}.pth'.format(args.dataset,\n",
    "                                                                                                   epoch)))\n",
    "    logging.info(f'Checkpoint {epoch} saved!')\n",
    "    print('Best Accuracy = {}'.format(best_accuracy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "670589f3-c14d-48ad-860d-76a9597af79f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 8192])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "primary_caps_output_op.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d33f4a8b-ef72-4663-850c-9cabb878fb24",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels_op.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0246ffe-8b0b-4cdc-a5b2-dc38d3bb8cef",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Py38",
   "language": "python",
   "name": "xy_py38"
  },
  "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}