Batch BERT.ipynb

batch-bert.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Batch BERT.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyM8n0ELZ0iyGphG/FpODcdt",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/davidefiocco/4d738ef9d3b1976187086ea31ca25ed2/batch-bert.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5DvL4a3XHuIF"
      },
      "source": [
        "# Benchmarking inference performance with an without data loaders\n",
        "\n",
        "Run this notebook on a CPU/GPU runtime and measure differences in inference execution times."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oLUvPiE1d7I0"
      },
      "source": [
        "!pip install transformers --quiet"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tCuloF_VewVC"
      },
      "source": [
        "from transformers import BertTokenizer, BertForSequenceClassification\n",
        "from torch.utils.data import TensorDataset\n",
        "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n",
        "import torch\n",
        "import numpy as np"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xMTdYo46GtdL",
        "outputId": "6389c93a-0a1f-4cf1-c5c5-fbdd9bb2a98f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "device(type='cpu')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vhO0kdhtpLJU"
      },
      "source": [
        "tokenizer = BertTokenizer.from_pretrained('textattack/bert-base-uncased-CoLA', truncation=True)\n",
        "model = BertForSequenceClassification.from_pretrained('textattack/bert-base-uncased-CoLA', return_dict=True).to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0kSpGglbpd-v"
      },
      "source": [
        "import spacy\n",
        "\n",
        "nlp = spacy.load(\"en_core_web_sm\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6RSwICRLqWce"
      },
      "source": [
        "text = \"\"\"\n",
        "They had stayed in my Madrid apartment 18 months earlier, so that Virgilio could recover from a two-week hospital stay after having fallen ill during a visit to Spain. My kitchen became a shrine as visitors trooped through, anxious to thank Virgilio for serving in a volunteer army called the International Brigades. That unit of 35,000 foreigners from 80 of today’s nations had fought against fascism in the Spanish civil war and been disbanded in 1938, a year before the short-lived democratic republic was finally extinguished. More than 50 years later, their actions still resonated.\n",
        "A 1937 poster declaring ‘All The People Of The World Are United With Spain’.\n",
        "The republic is an emotional touchstone for leftwing Spaniards, but admirers of the volunteers are spread across the world. Groups devoted to their memory exist in the US, Britain and half a dozen European countries. Mention of them can provoke sudden displays of enthusiasm, as I discovered when I began researching the group: a Spanish journalist pulled down his shirt to reveal the Brigades’ triangular symbol tattooed on his shoulder; a German in California sang their songs; and a Scottish writer at a neoliberal magazine talked wistfully about an uncle from Glasgow who had volunteered. David Simon, creator of The Wire, is now planning a drama series about the International Brigades.\n",
        "Elsewhere, opinion is dramatically opposed. In Poland, streets dedicated to the Dabrowski battalion of the International Brigades are being renamed by the Institute of National Remembrance, which oversees a controversial “decommunisation” law passed by the ultraconservative Law and Justice party in 2017. The brigaders had “served Stalinism”, their Polish critics argued. They were not entirely wrong.\n",
        "International Brigades veterans went on to serve as iron curtain prime ministers – or equivalent – in East Germany, Hungary and Albania. They provided dozens of ministers, generals, police chiefs and ambassadors across all Europe’s communist regimes – forming a potent elite, although they were mostly working-class. In East Germany, former International Brigades volunteers founded and ran the notorious Stasi. Suppressing freedom was part of their job. Little surprise, then, that some countrymen now despise them.\n",
        "History is neither neat nor clean, especially when it comes to past wars. The first casualty of war is said to be truth, but really it is nuance. War presents stark, binary choices. Kill or be killed. One side or the other. The truth is more complex than that, as the story of the International Brigades and their afterlife shows.\n",
        "\"\"\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MPQA8vYxqbP0"
      },
      "source": [
        "sents = [sent.text for sent in nlp(text).sents]\n",
        "sents = sents*2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GjM6swzK5Xzx"
      },
      "source": [
        "\n",
        "### With DataLoader\n",
        "\n",
        "This version is faster than with DataLoader-less one below on GPU, but not on CPU!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KZP2emUvhKAJ"
      },
      "source": [
        "encoded_data_val = tokenizer.batch_encode_plus(\n",
        "    sents, \n",
        "    add_special_tokens=True, \n",
        "    return_attention_mask=True, \n",
        "    padding='longest',\n",
        "    truncation=True,\n",
        "    max_length=256, \n",
        "    return_tensors='pt'\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dFCG4SShqmQl"
      },
      "source": [
        "input_ids_val = encoded_data_val['input_ids']\n",
        "attention_masks_val = encoded_data_val['attention_mask']\n",
        "\n",
        "dataset_val = TensorDataset(input_ids_val, attention_masks_val)\n",
        "\n",
        "batch_size = 4\n",
        "dataloader_val = DataLoader(dataset_val, sampler=SequentialSampler(dataset_val), batch_size=batch_size)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BZvrNV-w7A_D",
        "outputId": "51cc5186-80f0-42bb-d76d-d235551f46c4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        }
      },
      "source": [
        "%%time\n",
        "\n",
        "all_logits = np.empty([0,2])\n",
        "\n",
        "for batch in dataloader_val:\n",
        "    \n",
        "    batch = tuple(b.to(device) for b in batch)\n",
        "    inputs = {'input_ids':      batch[0],\n",
        "              'attention_mask': batch[1],\n",
        "              }\n",
        "\n",
        "    with torch.no_grad():        \n",
        "        outputs = model(**inputs)\n",
        "\n",
        "    logits = outputs[0]\n",
        "    all_logits = np.vstack([all_logits, torch.softmax(logits, dim=1).detach().cpu().numpy()])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU times: user 12.9 s, sys: 70.4 ms, total: 13 s\n",
            "Wall time: 13 s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vpKdxG6x8ePV"
      },
      "source": [
        "### No dataloaders\n",
        "\n",
        "Inference is faster than the code above on CPU"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AHmwfgc77GOx",
        "outputId": "16cad6ef-44d0-463c-a7ba-a8159987c5d0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        }
      },
      "source": [
        "%%time\n",
        "\n",
        "all_logits2 = np.empty([0,2])\n",
        "\n",
        "for sent in sents:\n",
        "    input_ids = torch.tensor(tokenizer.encode(sent, add_special_tokens=True, return_attention_mask=False, padding='longest', truncation=True, max_length=256)).unsqueeze(0).to(device)  # Batch size 1\n",
        "    labels = torch.tensor([1]).unsqueeze(0).to(device)  # Batch size 1\n",
        "    outputs = model(input_ids, labels=labels)\n",
        "    loss, logits = outputs[:2]\n",
        "    all_logits2 = np.vstack([all_logits2, torch.softmax(logits, dim=1).detach().cpu().numpy()])\n",
        "\n",
        "all_logits2"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU times: user 6.67 s, sys: 104 ms, total: 6.78 s\n",
            "Wall time: 6.79 s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GMIXBu0m-8q1",
        "outputId": "5c3bd4de-f2d0-4c34-a7ce-e54dca946be3",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "#check that results are practically identical in the two methods\n",
        "np.max(all_logits - all_logits2)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "1.1920928955078125e-07"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 30
        }
      ]
    }
  ]
}