AdamSwitchSGDM.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XBehZRIAVTxh",
        "outputId": "dce469d9-df13-4701-a5d2-fb226dd860c1",
        "cellView": "form"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n",
            "Sat Jan 21 03:57:20 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   38C    P0    25W /  70W |      0MiB / 15109MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "#@title Connections\n",
        "# Load the Drive helper and mount\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "path='/content/drive/MyDrive/Colab Notebooks/PADAM/'\n",
        "\n",
        "# Check GPU Info\n",
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Not connected to a GPU')\n",
        "else:\n",
        "  print(gpu_info)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Sf8m1T47Us-4",
        "outputId": "26680131-e5aa-44dc-9bb1-0968b0a8c961"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting adabound\n",
            "  Downloading adabound-0.0.5-py3-none-any.whl (5.1 kB)\n",
            "Requirement already satisfied: torch>=0.4.0 in /usr/local/lib/python3.8/dist-packages (from adabound) (1.13.1+cu116)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch>=0.4.0->adabound) (4.4.0)\n",
            "Installing collected packages: adabound\n",
            "Successfully installed adabound-0.0.5\n"
          ]
        }
      ],
      "source": [
        "#@title imports\n",
        "!pip install adabound\n",
        "#!pip install merlin-dataloader\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import torch.backends.cudnn as cudnn\n",
        "\n",
        "from torch.optim import Adam\n",
        "\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.models as models\n",
        "\n",
        "import os\n",
        "import argparse\n",
        "\n",
        "from torch.autograd import Variable\n",
        "\n",
        "from torch.optim.lr_scheduler import MultiStepLR\n",
        "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
        "from torch.optim.lr_scheduler import CosineAnnealingLR\n",
        "\n",
        "import adabound\n",
        "\n",
        "import json\n",
        "from copy import deepcopy\n",
        "\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "GYnlEQEaV296"
      },
      "outputs": [],
      "source": [
        "#@title utils\n",
        "'''Some helper functions for PyTorch, including:\n",
        "    - get_mean_and_std: calculate the mean and std value of dataset.\n",
        "    - msr_init: net parameter initialization.\n",
        "    - progress_bar: progress bar mimic xlua.progress.\n",
        "'''\n",
        "import os\n",
        "import sys\n",
        "import time\n",
        "import math\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch.nn.init as init\n",
        "\n",
        "\n",
        "def get_mean_and_std(dataset):\n",
        "    '''Compute the mean and std value of dataset.'''\n",
        "    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2)\n",
        "    mean = torch.zeros(3)\n",
        "    std = torch.zeros(3)\n",
        "    print('==> Computing mean and std..')\n",
        "    for inputs, targets in dataloader:\n",
        "        for i in range(3):\n",
        "            mean[i] += inputs[:,i,:,:].mean()\n",
        "            std[i] += inputs[:,i,:,:].std()\n",
        "    mean.div_(len(dataset))\n",
        "    std.div_(len(dataset))\n",
        "    return mean, std\n",
        "\n",
        "def init_params(net):\n",
        "    '''Init layer parameters.'''\n",
        "    for m in net.modules():\n",
        "        if isinstance(m, nn.Conv2d):\n",
        "            init.kaiming_normal(m.weight, mode='fan_out')\n",
        "            if m.bias:\n",
        "                init.constant(m.bias, 0)\n",
        "        elif isinstance(m, nn.BatchNorm2d):\n",
        "            init.constant(m.weight, 1)\n",
        "            init.constant(m.bias, 0)\n",
        "        elif isinstance(m, nn.Linear):\n",
        "            init.normal(m.weight, std=1e-3)\n",
        "            if m.bias:\n",
        "                init.constant(m.bias, 0)\n",
        "\n",
        "\n",
        "# _, term_width = os.popen('stty size', 'r').read().split()\n",
        "# term_width = int(term_width)\n",
        "term_width=20\n",
        "\n",
        "TOTAL_BAR_LENGTH = 65.\n",
        "last_time = time.time()\n",
        "begin_time = last_time\n",
        "def progress_bar(current, total, msg=None):\n",
        "    global last_time, begin_time\n",
        "    if current == 0:\n",
        "        begin_time = time.time()  # Reset for new bar.\n",
        "\n",
        "    cur_len = int(TOTAL_BAR_LENGTH*current/total)\n",
        "    rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1\n",
        "\n",
        "    sys.stdout.write(' [')\n",
        "    for i in range(cur_len):\n",
        "        sys.stdout.write('=')\n",
        "    sys.stdout.write('>')\n",
        "    for i in range(rest_len):\n",
        "        sys.stdout.write('.')\n",
        "    sys.stdout.write(']')\n",
        "\n",
        "    cur_time = time.time()\n",
        "    step_time = cur_time - last_time\n",
        "    last_time = cur_time\n",
        "    tot_time = cur_time - begin_time\n",
        "\n",
        "    L = []\n",
        "    L.append('  Step: %s' % format_time(step_time))\n",
        "    L.append(' | Tot: %s' % format_time(tot_time))\n",
        "    if msg:\n",
        "        L.append(' | ' + msg)\n",
        "\n",
        "    msg = ''.join(L)\n",
        "    sys.stdout.write(str(msg))\n",
        "    for i in range(term_width-int(TOTAL_BAR_LENGTH)-len(msg)-3):\n",
        "        sys.stdout.write(' ')\n",
        "\n",
        "    # Go back to the center of the bar.\n",
        "    for i in range(term_width-int(TOTAL_BAR_LENGTH/2)+2):\n",
        "        sys.stdout.write('\\b')\n",
        "    sys.stdout.write(' %d/%d ' % (current+1, total))\n",
        "\n",
        "    if current < total-1:\n",
        "        sys.stdout.write('\\r')\n",
        "    else:\n",
        "        sys.stdout.write('\\n')\n",
        "    sys.stdout.flush()\n",
        "\n",
        "def format_time(seconds):\n",
        "    days = int(seconds / 3600/24)\n",
        "    seconds = seconds - days*3600*24\n",
        "    hours = int(seconds / 3600)\n",
        "    seconds = seconds - hours*3600\n",
        "    minutes = int(seconds / 60)\n",
        "    seconds = seconds - minutes*60\n",
        "    secondsf = int(seconds)\n",
        "    seconds = seconds - secondsf\n",
        "    millis = int(seconds*1000)\n",
        "\n",
        "    f = ''\n",
        "    i = 1\n",
        "    if days > 0:\n",
        "        f += str(days) + 'D'\n",
        "        i += 1\n",
        "    if hours > 0 and i <= 2:\n",
        "        f += str(hours) + 'h'\n",
        "        i += 1\n",
        "    if minutes > 0 and i <= 2:\n",
        "        f += str(minutes) + 'm'\n",
        "        i += 1\n",
        "    if secondsf > 0 and i <= 2:\n",
        "        f += str(secondsf) + 's'\n",
        "        i += 1\n",
        "    if millis > 0 and i <= 2:\n",
        "        f += str(millis) + 'ms'\n",
        "        i += 1\n",
        "    if f == '':\n",
        "        f = '0ms'\n",
        "    return f\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wKeN_HDbVTnJ"
      },
      "outputs": [],
      "source": [
        "#@title VGG\n",
        "'''VGG11/13/16/19 in Pytorch.'''\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.autograd import Variable\n",
        "\n",
        "\n",
        "cfg = {\n",
        "    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],\n",
        "    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],\n",
        "    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],\n",
        "    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],\n",
        "}\n",
        "\n",
        "\n",
        "class VGG(nn.Module):\n",
        "    def __init__(self, vgg_name, num_classes = 10):\n",
        "        super(VGG, self).__init__()\n",
        "        self.features = self._make_layers(cfg[vgg_name])\n",
        "        self.classifier = nn.Linear(512, num_classes)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.features(x)\n",
        "        out = out.view(out.size(0), -1)\n",
        "        out = self.classifier(out)\n",
        "        return out\n",
        "\n",
        "    def _make_layers(self, cfg):\n",
        "        layers = []\n",
        "        in_channels = 3\n",
        "        for x in cfg:\n",
        "            if x == 'M':\n",
        "                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]\n",
        "            else:\n",
        "                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),\n",
        "                           nn.BatchNorm2d(x),\n",
        "                           nn.ReLU(inplace=True)]\n",
        "                in_channels = x\n",
        "        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]\n",
        "        return nn.Sequential(*layers)\n",
        "\n",
        "# net = VGG('VGG11')\n",
        "# x = torch.randn(2,3,32,32)\n",
        "# print(net(Variable(x)).size())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZDOtd_Kg2fNG"
      },
      "outputs": [],
      "source": [
        "#@title train_loop( ) and test_loop( )\n",
        "def train_loop(trainloader, model, criterion, optimizer, use_cuda, train_errs, train_losses):\n",
        "    # Training\n",
        "    model.train()  \n",
        "    train_loss = 0\n",
        "    correct = 0\n",
        "    total = 0\n",
        "    for batch_idx, (inputs, targets) in enumerate(trainloader):\n",
        "        if use_cuda:\n",
        "            inputs, targets = inputs.cuda(), targets.cuda()\n",
        "\n",
        "        def closure():\n",
        "            outputs = model(inputs)\n",
        "            loss = criterion(outputs, targets)\n",
        "            return loss\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(inputs)\n",
        "        loss = criterion(outputs, targets)\n",
        "        loss.backward()\n",
        "        optimizer.step(closure)\n",
        "\n",
        "        train_loss += loss.item()\n",
        "        _, predicted = torch.max(outputs.data, 1)\n",
        "        total += targets.size(0)\n",
        "        correct += predicted.eq(targets.data).cpu().sum().item()\n",
        "\n",
        "        progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'\n",
        "                      % (train_loss/(batch_idx+1), 100.0/total*(correct), correct, total))\n",
        "\n",
        "    # Compute training error\n",
        "    train_loss = 0\n",
        "    correct = 0\n",
        "    total = 0\n",
        "    for batch_idx, (inputs, targets) in enumerate(trainloader):\n",
        "        if use_cuda:\n",
        "            inputs, targets = inputs.cuda(), targets.cuda()\n",
        "        outputs = model(inputs)\n",
        "        loss = criterion(outputs, targets)\n",
        "\n",
        "        train_loss += loss.item()\n",
        "        _, predicted = torch.max(outputs.data, 1)\n",
        "        total += targets.size(0)\n",
        "        correct += predicted.eq(targets.data).cpu().sum().item()\n",
        "\n",
        "        progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'\n",
        "                      % (train_loss/(batch_idx+1), 100.0/total*(correct), correct, total))\n",
        "    train_errs.append(1 - correct/total)\n",
        "    train_losses.append(train_loss/(batch_idx+1))\n",
        "\n",
        "\n",
        "def test_loop(testloader, model, criterion, use_cuda, test_errs, test_losses):\n",
        "    # Testing\n",
        "    model.eval()  \n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    total = 0\n",
        "    for batch_idx, (inputs, targets) in enumerate(testloader):\n",
        "        if use_cuda:\n",
        "            inputs, targets = inputs.cuda(), targets.cuda()\n",
        "        outputs = model(inputs)\n",
        "        loss = criterion(outputs, targets)\n",
        "\n",
        "        test_loss += loss.item()\n",
        "        _, predicted = torch.max(outputs.data, 1)\n",
        "        total += targets.size(0)\n",
        "        correct += predicted.eq(targets.data).cpu().sum().item()\n",
        "\n",
        "        progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'\n",
        "                      % (test_loss/(batch_idx+1), 100.0 / total * (correct), correct, total))\n",
        "    test_errs.append(1 - correct/total)\n",
        "    test_losses.append(test_loss/(batch_idx+1))\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Data Loader\n",
        "\n",
        "def load_cifar10():\n",
        "    print('==> Preparing cifar10 data..')\n",
        "    transform_train = transforms.Compose([\n",
        "        transforms.RandomCrop(32, padding=4),\n",
        "        transforms.RandomHorizontalFlip(),\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize((0.4914, 0.4822, 0.4465),\n",
        "                              (0.2023, 0.1994, 0.2010)),\n",
        "    ])\n",
        "\n",
        "    transform_test = transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize((0.4914, 0.4822, 0.4465),\n",
        "                              (0.2023, 0.1994, 0.2010)),\n",
        "    ])\n",
        "\n",
        "    trainset = torchvision.datasets.CIFAR10(\n",
        "        root=f'{path}/data/', train=True, download=True, transform=transform_train)\n",
        "    trainloader = torch.utils.data.DataLoader(\n",
        "        trainset, batch_size=128, shuffle=True, num_workers=2)\n",
        "\n",
        "    testset = torchvision.datasets.CIFAR10(\n",
        "        root=f'{path}/data/', train=False, download=True, transform=transform_test)\n",
        "    testloader = torch.utils.data.DataLoader(\n",
        "        testset, batch_size=100, shuffle=False, num_workers=2)\n",
        "\n",
        "    classes = ('plane', 'car', 'bird', 'cat', 'deer',\n",
        "                'dog', 'frog', 'horse', 'ship', 'truck')\n",
        "    return trainloader, testloader\n",
        "\n",
        "\n",
        "def load_cifar100():\n",
        "    print('==> Preparing cifar100 data..')\n",
        "    transform_train = transforms.Compose([\n",
        "        transforms.RandomCrop(32, padding=4),\n",
        "        transforms.RandomHorizontalFlip(),\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize(mean=[0.507, 0.487, 0.441], std=[0.267, 0.256, 0.276]),\n",
        "    ])\n",
        "    transform_test = transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize(mean=[0.507, 0.487, 0.441], std=[0.267, 0.256, 0.276]),\n",
        "    ])\n",
        "\n",
        "    trainset = torchvision.datasets.CIFAR100(root=path+'/data/CIFAR100', train=True, download=True, transform=transform_train)\n",
        "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)\n",
        "\n",
        "    testset = torchvision.datasets.CIFAR100(root=path+'/data/CIFAR100', train=False, download=True, transform=transform_test)\n",
        "    testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)\n",
        "    return trainloader, testloader\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "wXbkEt3vT9gR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title def adam_switch_sgdm(switch,lr_sgdm,dataset)\n",
        "def adam_switch_sgdm(switch,lr_sgdm,dataset):\n",
        "    # Hyperparameters ##############################################\n",
        "    # switch=switch\n",
        "    # lr=0.001\n",
        "    lr_adam=0.001\n",
        "    # lr_sgdm=lr_sgdm\n",
        "    net=\"vggnet\"\n",
        "    method=f\"adam_switch_sgdm_switch={switch}_lr={lr_adam}_{lr_sgdm}\" #{lr_sgdm}\n",
        "    wd = 5e-4\n",
        "    beta1 = 0.5\n",
        "    beta2 = 0.999\n",
        "    Nepoch=200\n",
        "    resume=True\n",
        "\n",
        "    print(method)\n",
        "\n",
        "    # init #########################################################\n",
        "    results_path=f'{path}/checkpoint/{net}/{dataset}/{method}'\n",
        "    use_cuda = torch.cuda.is_available()\n",
        "    best_acc = 0  # best test accuracy\n",
        "    start_epoch = 0  # start from epoch 0 or last checkpoint epoch\n",
        "    train_errs = []\n",
        "    test_errs = []\n",
        "    train_losses = []\n",
        "    test_losses = []\n",
        "    lrs = []\n",
        "\n",
        "    # Data #########################################################\n",
        "    if dataset == 'CIFAR10':\n",
        "        trainloader, testloader = load_cifar10()\n",
        "    elif dataset == 'CIFAR100':\n",
        "        trainloader, testloader = load_cifar100()\n",
        "    else:\n",
        "        print(f'{dataset} Not Available!')\n",
        "        return 0\n",
        "\n",
        "    # Model #########################################################\n",
        "    if resume:\n",
        "        # Load checkpoint.\n",
        "        print('==> Resuming from checkpoint..')\n",
        "        try:\n",
        "            checkpoint = torch.load(results_path)\n",
        "            model = checkpoint['model']\n",
        "            start_epoch = checkpoint['epoch']\n",
        "            train_losses = checkpoint['train_losses']\n",
        "            train_errs = checkpoint['train_errs']\n",
        "            test_losses = checkpoint['test_losses']\n",
        "            test_errs = checkpoint['test_errs']\n",
        "            best_acc = checkpoint['best_acc']\n",
        "            lrs = checkpoint['lrs']\n",
        "            lr=lrs[-1]\n",
        "        except Exception as e:\n",
        "            print(e)\n",
        "            print('==> Building model..')\n",
        "            if net == 'vggnet':\n",
        "                model = VGG('VGG16', num_classes=10)\n",
        "            # #     model = models.vgg16_bn(num_classes=10)\n",
        "            # elif net == 'resnet':\n",
        "            #     model = ResNet18(num_classes=10)\n",
        "            # #     model = models.resnet18(num_classes=10)\n",
        "            # elif net == 'wideresnet':\n",
        "            #     model = WResNet_cifar10(\n",
        "            #         num_classes=10, depth=16, multiplier=4)\n",
        "            else:\n",
        "                print('Network undefined!')\n",
        "    else:\n",
        "        print('==> Building model..')\n",
        "        if net == 'vggnet':\n",
        "            model = VGG('VGG16', num_classes=10)\n",
        "        # #     model = models.vgg16_bn(num_classes=10)\n",
        "        # elif net == 'resnet':\n",
        "        #     model = ResNet18(num_classes=10)\n",
        "        # #     model = models.resnet18(num_classes=10)\n",
        "        # elif net == 'wideresnet':\n",
        "        #     model = WResNet_cifar10(\n",
        "        #         num_classes=10, depth=16, multiplier=4)\n",
        "        else:\n",
        "            print('Network undefined!')\n",
        "\n",
        "    if use_cuda:\n",
        "        model.cuda()\n",
        "        model = torch.nn.DataParallel(\n",
        "            model, device_ids=range(torch.cuda.device_count()))\n",
        "        cudnn.benchmark = True\n",
        "\n",
        "\n",
        "    # Optimize/Train #####################################\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "    ### Adam\n",
        "    betas = (beta1, beta2)\n",
        "    optimizer = Adam(model.parameters(), lr=lr_adam, weight_decay = wd, betas = betas)\n",
        "    milestones=[150-start_epoch]\n",
        "    scheduler = MultiStepLR(optimizer, milestones=[i for i in milestones if i>0], gamma=0.1)\n",
        "\n",
        "    for epoch in range(start_epoch+1, switch+1):\n",
        "        print('\\nEpoch: %d' % epoch, ' Learning rate:', scheduler.get_last_lr(), ' best_acc: ', best_acc)\n",
        "        # Train\n",
        "        train_loop(trainloader=trainloader, model=model, \n",
        "                    criterion=criterion, optimizer=optimizer,\n",
        "                    use_cuda=use_cuda,train_errs=train_errs,train_losses=train_losses)\n",
        "        # Test\n",
        "        test_loop(testloader=testloader, model=model, criterion=criterion, \n",
        "                  use_cuda=use_cuda, test_errs=test_errs, test_losses=test_losses)\n",
        "        # Save checkpoint\n",
        "        acc = 100.0 * (1-test_errs[-1])\n",
        "        if acc > best_acc:\n",
        "            best_acc = acc\n",
        "        lrs.extend(scheduler.get_last_lr())\n",
        "        print('Saving..')\n",
        "        state = {\n",
        "            'model': model,\n",
        "            'epoch': epoch,\n",
        "            'best_acc': best_acc,\n",
        "            'train_errs': train_errs,\n",
        "            'train_losses': train_losses,\n",
        "            'test_errs': test_errs,\n",
        "            'test_losses': test_losses,\n",
        "            'lrs': lrs\n",
        "        }\n",
        "        torch.save(state, results_path)\n",
        "\n",
        "        scheduler.step()\n",
        "\n",
        "    if lr_sgdm > 0:\n",
        "        # Switch to SGDM\n",
        "        print('============ Switch to sgdm ==========')\n",
        "        optimizer = optim.SGD(model.parameters(), lr=lr_sgdm, momentum= 0.9,weight_decay = wd)\n",
        "        milestones=[100-max(switch,start_epoch),150-max(switch,start_epoch)]\n",
        "        # milestones=[150-max(switch,start_epoch)]\n",
        "        print([i for i in milestones if i>0])\n",
        "        scheduler = MultiStepLR(optimizer, milestones=[i for i in milestones if i>0], gamma=0.1)\n",
        "\n",
        "        for epoch in range(max(switch+1,start_epoch+1), 200+1):\n",
        "            print('\\nEpoch: %d' % epoch, ' Learning rate:', scheduler.get_last_lr(), ' best_acc: ', best_acc)\n",
        "            # Train\n",
        "            train_loop(trainloader=trainloader, model=model, \n",
        "                        criterion=criterion, optimizer=optimizer,\n",
        "                        use_cuda=use_cuda,train_errs=train_errs,train_losses=train_losses)\n",
        "            # Test\n",
        "            test_loop(testloader=testloader, model=model, criterion=criterion, \n",
        "                      use_cuda=use_cuda, test_errs=test_errs, test_losses=test_losses)\n",
        "            # Save checkpoint\n",
        "            acc = 100.0 * (1-test_errs[-1])\n",
        "            if acc > best_acc:\n",
        "                best_acc = acc\n",
        "            lrs.extend(scheduler.get_last_lr())\n",
        "            print('Saving..')\n",
        "            state = {\n",
        "                'model': model,\n",
        "                'epoch': epoch,\n",
        "                'best_acc': best_acc,\n",
        "                'train_errs': train_errs,\n",
        "                'train_losses': train_losses,\n",
        "                'test_errs': test_errs,\n",
        "                'test_losses': test_losses,\n",
        "                'lrs': lrs\n",
        "            }\n",
        "            torch.save(state, results_path)\n",
        "\n",
        "            scheduler.step()\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "B5Sj1757bldF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "LRS = [0.001,0.002,0.003,0.004,0.005,0.006,0.007,0.008,0.009,\n",
        "       0.01,0.015,0.02,0.025,0.03,0.035,0.04,0.045,0.05,0.055,\n",
        "       0.06,0.065,0.07,0.075,0.08,0.085,0.09,0.095,0.1,0.11,\n",
        "       0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1]"
      ],
      "metadata": {
        "id": "5AyDHoOD7VJH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "for switch in [10,15,25,50,100]:\n",
        "    for lr_sgdm in LRS:\n",
        "        adam_switch_sgdm(switch,lr_sgdm,'CIFAR10')\n"
      ],
      "metadata": {
        "id": "Rz4Yx65ycBeF"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true
    },
    "gpuClass": "premium",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}