myTrain.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd0584de",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Class EMA  -------------------------------------------------------------------------#\n",
    "class EMA:\n",
    "    def __init__(self, model, decay):\n",
    "        self.decay = decay\n",
    "        self.shadow = {}\n",
    "        self.original = {}\n",
    "\n",
    "        for name, param in model.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                self.shadow[name] = param.data.clone()\n",
    "\n",
    "    def __call__(self, model, num_updates):\n",
    "        decay = min(self.decay, (1.0 + num_updates) / (10.0 + num_updates))\n",
    "        for name, param in model.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                assert name in self.shadow\n",
    "                new_average = (1.0 - decay) * param.data + decay * self.shadow[name]\n",
    "                self.shadow[name] = new_average.clone()\n",
    "\n",
    "    def assign(self, model):\n",
    "        for name, param in model.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                assert name in self.shadow\n",
    "                self.original[name] = param.data.clone()\n",
    "                param.data = self.shadow[name]\n",
    "\n",
    "    def resume(self, model):\n",
    "        for name, param in model.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                assert name in self.shadow\n",
    "                param.data = self.original[name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d32d1a9",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Class Datasets  -------------------------------------------------------------------------#\n",
    "import torch\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "214d19d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MnistDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, training=True, transform=None):\n",
    "        if training==True:\n",
    "            f = open('../data/MNIST/raw/train-images-idx3-ubyte', 'rb')\n",
    "            xs = np.array(np.frombuffer(f.read(), np.uint8, offset=16))\n",
    "            f.close()\n",
    "            f = open('../data/MNIST/raw/train-labels-idx1-ubyte', 'rb')\n",
    "            ys = np.array(np.frombuffer(f.read(), np.uint8, offset=8))\n",
    "            f.close()\n",
    "        else:\n",
    "            f = open('../data/MNIST/raw/t10k-images-idx3-ubyte', 'rb')\n",
    "            xs = np.array(np.frombuffer(f.read(), np.uint8, offset=16))\n",
    "            f.close()\n",
    "            f = open('../data/MNIST/raw/t10k-labels-idx1-ubyte', 'rb')\n",
    "            ys = np.array(np.frombuffer(f.read(), np.uint8, offset=8))\n",
    "            f.close()\n",
    "        xs = np.reshape(xs, (-1, 28, 28, 1)).astype(np.float32)\n",
    "        ys = ys.astype(np.int)\n",
    "        self.x_data = xs\n",
    "        self.y_data = ys\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.x_data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        x = Image.fromarray(self.x_data[idx].reshape(28, 28))\n",
    "        y = torch.tensor(np.array(self.y_data[idx]))\n",
    "        if self.transform:\n",
    "            x = self.transform(x)\n",
    "        x = transforms.ToTensor()(np.array(x)/255)\n",
    "        return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18c09a52",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Class transforms  -------------------------------------------------------------------------#\n",
    "import random\n",
    "import torchvision.transforms.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c10d0930",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "class RandomRotation(object):\n",
    "    def __init__(self, degrees, seed=1):\n",
    "        self.degrees = (-degrees, degrees)\n",
    "        random.seed(seed)\n",
    "    \n",
    "    @staticmethod\n",
    "    def get_params(degrees):\n",
    "        angle = random.uniform(degrees[0], degrees[1])\n",
    "        return angle\n",
    "\n",
    "    def __call__(self, img):\n",
    "        angle = self.get_params(self.degrees)\n",
    "        return F.rotate(img, angle, False, False, None, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66b3d5eb",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Class ModelM5  -------------------------------------------------------------------------#\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2944c3a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelM5(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ModelM5, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 5, bias=False)\n",
    "        self.conv1_bn = nn.BatchNorm2d(32)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 5, bias=False)\n",
    "        self.conv2_bn = nn.BatchNorm2d(64)\n",
    "        self.conv3 = nn.Conv2d(64, 96, 5, bias=False)\n",
    "        self.conv3_bn = nn.BatchNorm2d(96)\n",
    "        self.conv4 = nn.Conv2d(96, 128, 5, bias=False)\n",
    "        self.conv4_bn = nn.BatchNorm2d(128)\n",
    "        self.conv5 = nn.Conv2d(128, 160, 5, bias=False)\n",
    "        self.conv5_bn = nn.BatchNorm2d(160)\n",
    "        self.fc1 = nn.Linear(10240, 10, bias=False)\n",
    "        self.fc1_bn = nn.BatchNorm1d(10)\n",
    "    def get_logits(self, x):\n",
    "        x = (x - 0.5) * 2.0\n",
    "        conv1 = F.relu(self.conv1_bn(self.conv1(x)))\n",
    "        conv2 = F.relu(self.conv2_bn(self.conv2(conv1)))\n",
    "        conv3 = F.relu(self.conv3_bn(self.conv3(conv2)))\n",
    "        conv4 = F.relu(self.conv4_bn(self.conv4(conv3)))\n",
    "        conv5 = F.relu(self.conv5_bn(self.conv5(conv4)))\n",
    "        flat5 = torch.flatten(conv5.permute(0, 2, 3, 1), 1)\n",
    "        logits = self.fc1_bn(self.fc1(flat5))\n",
    "        return logits\n",
    "    def forward(self, x):\n",
    "        logits = self.get_logits(x)\n",
    "        return F.log_softmax(logits, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "493ff4e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8bf7119",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports -------------------------------------------------------------------------#\n",
    "import sys\n",
    "import os\n",
    "import argparse\n",
    "import numpy as np \n",
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torchsummary import summary\n",
    "from PIL import Image\n",
    "# from ema import EMA\n",
    "# from datasets import MnistDataset\n",
    "# from transforms import RandomRotation\n",
    "# from models.modelM3 import ModelM3\n",
    "# from models.modelM5 import ModelM5\n",
    "# from models.modelM7 import ModelM7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ceda3bb",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def run(p_seed=0, p_epochs=150, p_kernel_size=5, p_logdir=\"temp\"):\n",
    "    # random number generator seed ------------------------------------------------#\n",
    "    SEED = p_seed\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    torch.manual_seed(SEED)\n",
    "    torch.cuda.manual_seed_all(SEED)\n",
    "    np.random.seed(SEED)\n",
    "\n",
    "    # kernel size of model --------------------------------------------------------#\n",
    "    KERNEL_SIZE = p_kernel_size\n",
    "\n",
    "    # number of epochs ------------------------------------------------------------#\n",
    "    NUM_EPOCHS = p_epochs\n",
    "\n",
    "    # file names ------------------------------------------------------------------#\n",
    "    if not os.path.exists(\"../logs/%s\"%p_logdir):\n",
    "        os.makedirs(\"../logs/%s\"%p_logdir)\n",
    "    OUTPUT_FILE = str(\"../logs/%s/log%03d.out\"%(p_logdir,SEED))\n",
    "    MODEL_FILE = str(\"../logs/%s/model%03d.pth\"%(p_logdir,SEED))\n",
    "\n",
    "    # enable GPU usage ------------------------------------------------------------#\n",
    "    use_cuda = torch.cuda.is_available()\n",
    "    device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "    if use_cuda == False:\n",
    "        print(\"WARNING: CPU will be used for training.\")\n",
    "        exit(0)\n",
    "\n",
    "    # data augmentation methods ---------------------------------------------------#\n",
    "    transform = transforms.Compose([\n",
    "        RandomRotation(20, seed=SEED),\n",
    "        transforms.RandomAffine(0, translate=(0.2, 0.2)),\n",
    "        ])\n",
    "\n",
    "    # data loader -----------------------------------------------------------------#\n",
    "    train_dataset = MnistDataset(training=True, transform=transform)\n",
    "    test_dataset = MnistDataset(training=False, transform=None)\n",
    "    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=120, shuffle=True)\n",
    "    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=100, shuffle=False)\n",
    "\n",
    "    # model selection -------------------------------------------------------------#\n",
    "    if(KERNEL_SIZE == 3):\n",
    "        model = ModelM3().to(device)\n",
    "    elif(KERNEL_SIZE == 5):\n",
    "        model = ModelM5().to(device)\n",
    "    elif(KERNEL_SIZE == 7):\n",
    "        model = ModelM7().to(device)\n",
    "\n",
    "    summary(model, (1, 28, 28))\n",
    "\n",
    "    # hyperparameter selection ----------------------------------------------------#\n",
    "    ema = EMA(model, decay=0.999)\n",
    "    optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "    lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.98)\n",
    "\n",
    "    # delete result file ----------------------------------------------------------#\n",
    "    f = open(OUTPUT_FILE, 'w')\n",
    "    f.close()\n",
    "\n",
    "    # global variables ------------------------------------------------------------#\n",
    "    g_step = 0\n",
    "    max_correct = 0\n",
    "\n",
    "    # training and evaluation loop ------------------------------------------------#\n",
    "    for epoch in range(NUM_EPOCHS):\n",
    "        #--------------------------------------------------------------------------#\n",
    "        # train process                                                            #\n",
    "        #--------------------------------------------------------------------------#\n",
    "        model.train()\n",
    "        train_loss = 0\n",
    "        train_corr = 0\n",
    "        for batch_idx, (data, target) in enumerate(train_loader):\n",
    "            data, target = data.to(device), target.to(device, dtype=torch.int64)\n",
    "            optimizer.zero_grad()\n",
    "            output = model(data)\n",
    "            loss = F.nll_loss(output, target)\n",
    "            train_pred = output.argmax(dim=1, keepdim=True)\n",
    "            train_corr += train_pred.eq(target.view_as(train_pred)).sum().item()\n",
    "            train_loss += F.nll_loss(output, target, reduction='sum').item()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            g_step += 1\n",
    "            ema(model, g_step)\n",
    "            if batch_idx % 100 == 0:\n",
    "                print('Train Epoch: {} [{:05d}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                    epoch, batch_idx * len(data), len(train_loader.dataset),\n",
    "                    100. * batch_idx / len(train_loader), loss.item()))\n",
    "        train_loss /= len(train_loader.dataset)\n",
    "        train_accuracy = 100 * train_corr / len(train_loader.dataset)\n",
    "\n",
    "        #--------------------------------------------------------------------------#\n",
    "        # test process                                                             #\n",
    "        #--------------------------------------------------------------------------#\n",
    "        model.eval()\n",
    "        ema.assign(model)\n",
    "        test_loss = 0\n",
    "        correct = 0\n",
    "        total_pred = np.zeros(0)\n",
    "        total_target = np.zeros(0)\n",
    "        with torch.no_grad():\n",
    "            for data, target in test_loader:\n",
    "                data, target = data.to(device), target.to(device,  dtype=torch.int64)\n",
    "                output = model(data)\n",
    "                test_loss += F.nll_loss(output, target, reduction='sum').item()\n",
    "                pred = output.argmax(dim=1, keepdim=True)\n",
    "                total_pred = np.append(total_pred, pred.cpu().numpy())\n",
    "                total_target = np.append(total_target, target.cpu().numpy())\n",
    "                correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "            if(max_correct < correct):\n",
    "                torch.save(model.state_dict(), MODEL_FILE)\n",
    "                max_correct = correct\n",
    "                print(\"Best accuracy! correct images: %5d\"%correct)\n",
    "        ema.resume(model)\n",
    "\n",
    "        #--------------------------------------------------------------------------#\n",
    "        # output                                                                   #\n",
    "        #--------------------------------------------------------------------------#\n",
    "        test_loss /= len(test_loader.dataset)\n",
    "        test_accuracy = 100 * correct / len(test_loader.dataset)\n",
    "        best_test_accuracy = 100 * max_correct / len(test_loader.dataset)\n",
    "        print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%) (best: {:.2f}%)\\n'.format(\n",
    "            test_loss, correct, len(test_loader.dataset), test_accuracy, best_test_accuracy))\n",
    "\n",
    "        f = open(OUTPUT_FILE, 'a')\n",
    "        f.write(\" %3d %12.6f %9.3f %12.6f %9.3f %9.3f\\n\"%(epoch, train_loss, train_accuracy, test_loss, test_accuracy, best_test_accuracy))\n",
    "        f.close()\n",
    "\n",
    "        #--------------------------------------------------------------------------#\n",
    "        # update learning rate scheduler                                           #\n",
    "        #--------------------------------------------------------------------------#\n",
    "        lr_scheduler.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "400b24c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    p = argparse.ArgumentParser()\n",
    "    p.add_argument(\"--seed\", default=0, type=int)\n",
    "    p.add_argument(\"--trials\", default=15, type=int)\n",
    "    p.add_argument(\"--epochs\", default=150, type=int)    \n",
    "    p.add_argument(\"--kernel_size\", default=5, type=int)    \n",
    "    p.add_argument(\"--gpu\", default=0, type=int)\n",
    "    p.add_argument(\"--logdir\", default=\"temp\")\n",
    "    args = p.parse_args()\n",
    "    os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"]=str(args.gpu)\n",
    "    for i in range(args.trials):\n",
    "        run(p_seed = args.seed + i,\n",
    "            p_epochs = args.epochs,\n",
    "            p_kernel_size = args.kernel_size,\n",
    "            p_logdir = args.logdir)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}