ray_tune_ignite.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install -U ray\n",
        "!pip install pytorch-ignite"
      ],
      "metadata": {
        "id": "FGKayfIaBeNd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import argparse\n",
        "from filelock import FileLock\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",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "from ignite.contrib.handlers import ProgressBar\n",
        "from ignite.engine import Engine\n",
        "from ignite.metrics import Accuracy, Loss\n",
        "from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator\n",
        "\n",
        "\n",
        "import ray\n",
        "from ray import tune\n",
        "from ray.tune.schedulers import AsyncHyperBandScheduler"
      ],
      "metadata": {
        "id": "BHNpa-Ofb2DH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_data_loaders():\n",
        "    transform = transforms.Compose([\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize((0.5, ), (0.5, )),\n",
        "    ])\n",
        "\n",
        "    train_set = datasets.MNIST(\n",
        "        \"~/.pytorch/MNIST_data/\", train=True, download=True, transform=transform)\n",
        "    test_set = datasets.MNIST(\n",
        "        \"~/.pytorch/MNIST_data/\", train=False, download=True, transform=transform)\n",
        "\n",
        "    train_loader = torch.utils.data.DataLoader(\n",
        "        train_set, batch_size=256, shuffle=True) \n",
        "\n",
        "    test_loader = torch.utils.data.DataLoader(\n",
        "        test_set, batch_size=256, shuffle=True)\n",
        "    \n",
        "    return train_loader, test_loader\n"
      ],
      "metadata": {
        "id": "NdyStWwLeHWi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Net(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.fc1 = nn.Linear(784, 500)\n",
        "        self.fc2 = nn.Linear(500, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x.view(-1, 784)\n",
        "        x = self.fc1(x)\n",
        "        x = F.relu(x)\n",
        "        x = self.fc2(x)\n",
        "        return F.log_softmax(x, dim=1)\n",
        "        \n",
        "model = Net().to(device)"
      ],
      "metadata": {
        "id": "ENN4U1hzeWoG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_step(engine, batch=None, optimizer=None):\n",
        "    data, targets = batch\n",
        "    data, targets = data.to(device), targets.to(device)\n",
        "    model.train()\n",
        "    optimizer.zero_grad()\n",
        "    outputs = model(data)\n",
        "    loss = F.nll_loss(outputs, targets)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    # return loss.data()\n",
        "\n",
        "def train_one_epoch(config):\n",
        "    train_loader, test_loader = get_data_loaders()\n",
        "    optimizer = optim.SGD(model.parameters(), \n",
        "                          lr=config[\"lr\"], \n",
        "                          momentum=config[\"momentum\"]\n",
        "                          )\n",
        "    \n",
        "    for batch in train_loader:\n",
        "        # batch = batch.to(device)\n",
        "        train_step(Engine, batch, optimizer)\n",
        "        acc = test(model, test_loader)\n",
        "        tune.report(mean_accuracy=acc)\n",
        "\n",
        "def validation_step(engine, batch):\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "        x, y = batch\n",
        "        x, y = x.to(device), y.to(device)\n",
        "        y_pred = model(x)\n",
        "    return y_pred, y\n",
        "\n",
        "def test(model, data_loader):\n",
        "    model.eval()\n",
        "    correct = 0\n",
        "    total = 0\n",
        "    with torch.no_grad():\n",
        "        for batch_idx, (data, target) in enumerate(data_loader):\n",
        "            # if batch_idx * len(data) > TEST_SIZE:\n",
        "            #     break\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            outputs = model(data)\n",
        "            _, predicted = torch.max(outputs.data, 1)\n",
        "            total += target.size(0)\n",
        "            correct += (predicted == target).sum().item()\n",
        "\n",
        "    return correct / total"
      ],
      "metadata": {
        "id": "9MGGbhLkelJl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Find best hyperparams using ray tune\n"
      ],
      "metadata": {
        "id": "0JBXbsc6nNg7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ray.shutdown()\n",
        "ray.init()\n",
        "\n",
        "# for early stopping\n",
        "sched = AsyncHyperBandScheduler()\n",
        "\n",
        "analysis = tune.run(\n",
        "    train_one_epoch,\n",
        "    metric=\"mean_accuracy\",\n",
        "    mode=\"max\",\n",
        "    name=\"exp\",\n",
        "    scheduler=sched,\n",
        "    stop={\n",
        "        \"mean_accuracy\": 0.98,\n",
        "        \"training_iteration\": 3,\n",
        "    },\n",
        "    resources_per_trial={\"cpu\": 1, \"gpu\": 1},  # set this for GPUs\n",
        "    num_samples=3,\n",
        "    config={\n",
        "        \"lr\": tune.loguniform(1e-3, 1e-2),\n",
        "        \"momentum\": tune.uniform(0.1, 0.4),\n",
        "    },\n",
        ")\n",
        "\n",
        "print(\"Best config is:\", analysis.best_config)"
      ],
      "metadata": {
        "id": "LZU8rSJIbO2D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Train using PyTorch Ignite"
      ],
      "metadata": {
        "id": "-YAPi7uFnSW6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "best_optimizer = optim.SGD(model.parameters(), \n",
        "                           lr=analysis.best_config[\"lr\"], #using best lr from ray tune\n",
        "                           momentum=analysis.best_config[\"momentum\"]) #using best momentum from ray tune\n",
        "                                                \n",
        "criterion = nn.CrossEntropyLoss()\n",
        "trainer = create_supervised_trainer(model.to(device), best_optimizer, criterion, device=device)\n",
        "val_metrics = {\n",
        "    \"accuracy\": Accuracy(),\n",
        "    \"loss\": Loss(criterion)\n",
        "}\n",
        "\n",
        "train_evaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)\n",
        "val_evaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)\n",
        "\n",
        "train_loader, test_loader = get_data_loaders()\n",
        "ProgressBar().attach(trainer)\n",
        "\n",
        "evaluator = Engine(validation_step)\n",
        "Accuracy().attach(evaluator, \"accuracy\")\n",
        "\n",
        "validate_every = 1\n",
        "log_every=1\n",
        "\n",
        "@trainer.on(Events.EPOCH_COMPLETED(every=validate_every))\n",
        "def run_validation():\n",
        "    evaluator.run(test_loader)\n",
        "\n",
        "@trainer.on(Events.EPOCH_COMPLETED(every=log_every))\n",
        "def log_validation():\n",
        "    metrics = evaluator.state.metrics\n",
        "    # print(metrics)\n",
        "    print(f\"Epoch: {trainer.state.epoch},  Accuracy: {metrics['accuracy']}\")\n",
        "\n",
        "trainer.run(train_loader, max_epochs=3)"
      ],
      "metadata": {
        "id": "7g0Jrnwtb9QZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "kvlpYhG6bPQ4"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "ray-tune+ignite.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}