{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# Multi-objective Optimization with Optuna\n\nThis tutorial showcases Optuna's multi-objective optimization feature by\noptimizing the validation accuracy of Fashion MNIST dataset and the FLOPS of the model implemented in PyTorch.\n\nWe use [thop](https://github.com/Lyken17/pytorch-OpCounter) to measure FLOPS.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import thop\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\n\nimport optuna\n\n\nDEVICE = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\nDIR = \"..\"\nBATCHSIZE = 128\nN_TRAIN_EXAMPLES = BATCHSIZE * 30\nN_VALID_EXAMPLES = BATCHSIZE * 10\n\n\ndef define_model(trial):\n    n_layers = trial.suggest_int(\"n_layers\", 1, 3)\n    layers = []\n\n    in_features = 28 * 28\n    for i in range(n_layers):\n        out_features = trial.suggest_int(\"n_units_l{}\".format(i), 4, 128)\n        layers.append(nn.Linear(in_features, out_features))\n        layers.append(nn.ReLU())\n        p = trial.suggest_float(\"dropout_{}\".format(i), 0.2, 0.5)\n        layers.append(nn.Dropout(p))\n\n        in_features = out_features\n\n    layers.append(nn.Linear(in_features, 10))\n    layers.append(nn.LogSoftmax(dim=1))\n\n    return nn.Sequential(*layers)\n\n\n# Defines training and evaluation.\ndef train_model(model, optimizer, train_loader):\n    model.train()\n    for batch_idx, (data, target) in enumerate(train_loader):\n        data, target = data.view(-1, 28 * 28).to(DEVICE), target.to(DEVICE)\n        optimizer.zero_grad()\n        F.nll_loss(model(data), target).backward()\n        optimizer.step()\n\n\ndef eval_model(model, valid_loader):\n    model.eval()\n    correct = 0\n    with torch.no_grad():\n        for batch_idx, (data, target) in enumerate(valid_loader):\n            data, target = data.view(-1, 28 * 28).to(DEVICE), target.to(DEVICE)\n            pred = model(data).argmax(dim=1, keepdim=True)\n            correct += pred.eq(target.view_as(pred)).sum().item()\n\n    accuracy = correct / N_VALID_EXAMPLES\n\n    flops, _ = thop.profile(model, inputs=(torch.randn(1, 28 * 28).to(DEVICE),), verbose=False)\n    return flops, accuracy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define multi-objective objective function.\nObjectives are FLOPS and accuracy.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def objective(trial):\n    train_dataset = torchvision.datasets.FashionMNIST(\n        DIR, train=True, download=True, transform=torchvision.transforms.ToTensor()\n    )\n    train_loader = torch.utils.data.DataLoader(\n        torch.utils.data.Subset(train_dataset, list(range(N_TRAIN_EXAMPLES))),\n        batch_size=BATCHSIZE,\n        shuffle=True,\n    )\n\n    val_dataset = torchvision.datasets.FashionMNIST(\n        DIR, train=False, transform=torchvision.transforms.ToTensor()\n    )\n    val_loader = torch.utils.data.DataLoader(\n        torch.utils.data.Subset(val_dataset, list(range(N_VALID_EXAMPLES))),\n        batch_size=BATCHSIZE,\n        shuffle=True,\n    )\n    model = define_model(trial).to(DEVICE)\n\n    optimizer = torch.optim.Adam(\n        model.parameters(), trial.suggest_float(\"lr\", 1e-5, 1e-1, log=True)\n    )\n\n    for epoch in range(10):\n        train_model(model, optimizer, train_loader)\n    flops, accuracy = eval_model(model, val_loader)\n    return flops, accuracy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Run multi-objective optimization\n\nIf your optimization problem is multi-objective,\nOptuna assumes that you will specify the optimization direction for each objective.\nSpecifically, in this example, we want to minimize the FLOPS (we want a faster model)\nand maximize the accuracy. So we set ``directions`` to ``[\"minimize\", \"maximize\"]``.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "study = optuna.create_study(directions=[\"minimize\", \"maximize\"])\nstudy.optimize(objective, n_trials=30, timeout=300)\n\nprint(\"Number of finished trials: \", len(study.trials))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Check trials on pareto front visually\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "optuna.visualization.plot_pareto_front(study, target_names=[\"FLOPS\", \"accuracy\"])"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "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.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}