genbmminside.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "genbmminside.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyOwxkm95cDgd+zXfnksN8Gy",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/lyutyuh/genbmm/blob/master/genbmminside.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install git+https://github.com/lyutyuh/genbmm"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sCM7jyOj1Nkm",
        "outputId": "ffa28c98-6d12-4bff-c250-d25b92b9e092"
      },
      "execution_count": 1,
      "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 git+https://github.com/lyutyuh/genbmm\n",
            "  Cloning https://github.com/lyutyuh/genbmm to /tmp/pip-req-build-vc88_tu8\n",
            "  Running command git clone -q https://github.com/lyutyuh/genbmm /tmp/pip-req-build-vc88_tu8\n",
            "Building wheels for collected packages: genbmm\n",
            "  Building wheel for genbmm (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for genbmm: filename=genbmm-0.1-cp37-cp37m-linux_x86_64.whl size=2126111 sha256=97df70f02a4d93a2afeea366664ac06c07252de4c5cdda6f98a2bd85c5816d38\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-a1odytd1/wheels/06/1a/23/e1223f7f8c9761cbd1e38c41fddacc6eaef55dc73e89000c44\n",
            "Successfully built genbmm\n",
            "Installing collected packages: genbmm\n",
            "Successfully installed genbmm-0.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "sLAiqg_v1Koe"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "import math\n",
        "from typing import Any, Dict, List, Tuple\n",
        "import numpy as np\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "assert torch.cuda.is_available(), \"enable CUDA to use genbmm\"\n",
        "\n",
        "from genbmm import logbmminside, logbmminside_rule"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def logsumexp(tensor: torch.Tensor, dim: int = -1, keepdim: bool = False) -> torch.Tensor:\n",
        "    max_score, _ = tensor.max(dim, keepdim=keepdim)\n",
        "    if keepdim:\n",
        "        stable_vec = tensor - max_score\n",
        "    else:\n",
        "        stable_vec = tensor - max_score.unsqueeze(dim)\n",
        "        \n",
        "    return max_score + stable_vec.logsumexp(dim, keepdim=keepdim)\n",
        "\n",
        "def stripe(x, n, w, offset=(0, 0), horizontal=1):\n",
        "    x, seq_len = x.contiguous(), x.size(1)\n",
        "    stride, numel = list(x.stride()), x[0, 0].numel()\n",
        "    stride[0] = (seq_len + 1) * numel\n",
        "    stride[1] = (1 if horizontal == 1 else seq_len) * numel\n",
        "    \n",
        "    return x.as_strided(\n",
        "        size=(n, w, *x.shape[2:]), \n",
        "        stride=stride,\n",
        "        storage_offset=(offset[0]*seq_len+offset[1])*numel\n",
        "    )"
      ],
      "metadata": {
        "id": "thtxdjwL132C"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# An example from https://github.com/lyutyuh/structured-span-selector\n",
        "LARGENUMBER = 1e4\n",
        "class CKY(torch.nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        max_span_width=30,\n",
        "    ):\n",
        "        super().__init__()\n",
        "        self.max_span_width = max_span_width\n",
        "        return\n",
        "    \n",
        "    def forward(\n",
        "        self,\n",
        "        span_mention_score_matrix: torch.FloatTensor, \n",
        "        sequence_lengths: torch.IntTensor,\n",
        "   ) -> Tuple[torch.FloatTensor]:\n",
        "        \n",
        "        with torch.autograd.enable_grad():\n",
        "            # Enable gradients during inference\n",
        "            return self.io(span_mention_score_matrix, sequence_lengths)\n",
        "        \n",
        "    def io(\n",
        "        self, \n",
        "        span_mention_score_matrix: torch.FloatTensor, \n",
        "        sequence_lengths: torch.IntTensor,\n",
        "    ) -> Tuple[torch.FloatTensor]:\n",
        "        \"\"\"\n",
        "            Parameters:\n",
        "                span_mention_score_matrix: shape (batch_size, sent_len, sent_len, score_dim)\n",
        "                    Score of each span being a span of interest. There are batch_size number\n",
        "                    of sentences in this document. And the maximum length of sentence is \n",
        "                    sent_len. \n",
        "                sequence_lengths: shape (batch_size, )\n",
        "                    The actual length of each sentence. \n",
        "        \"\"\"\n",
        "        span_mention_score_matrix.requires_grad_(True)\n",
        "        \n",
        "        batch_size, _, _, score_dim = span_mention_score_matrix.size()\n",
        "        seq_len = sequence_lengths.max()\n",
        "        # Shape: (batch_size, )\n",
        "        sequence_lengths = sequence_lengths.view(-1)\n",
        "        \n",
        "        # Shape: (seq_len, seq_len, score_dim, batch_size)\n",
        "        span_mention_score_matrix = span_mention_score_matrix.permute(1, 2, 3, 0)\n",
        "        \n",
        "        # There should be another matrix of non-mention span scores, which is full of 0s\n",
        "        # Shape: (seq_len, seq_len, score_dim + 1, batch_size), 2 for mention / non-mention\n",
        "        inside_s = span_mention_score_matrix.new_zeros(seq_len, seq_len, score_dim + 1, batch_size)\n",
        "        \n",
        "        for width in range(0, seq_len):\n",
        "            n = seq_len - width\n",
        "            if width == 0:\n",
        "                inside_s[:,:,:score_dim,:].diagonal(width).copy_(\n",
        "                    span_mention_score_matrix.diagonal(width)\n",
        "                )\n",
        "                continue\n",
        "\n",
        "            # [n, width, score_dim + 1, batch_size]\n",
        "            split_1 = stripe(inside_s, n, width)\n",
        "            split_2 = stripe(inside_s, n, width, (1, width), 0)\n",
        "\n",
        "            # [n, width, batch_size]\n",
        "            inside_s_span = logsumexp(split_1, 2) + logsumexp(split_2, 2)\n",
        "            # [1, batch_size, n]\n",
        "            inside_s_span = logsumexp(inside_s_span, 1, keepdim=True).permute(1, 2, 0)\n",
        "            \n",
        "            inside_s.diagonal(width).copy_(\n",
        "                torch.cat(\n",
        "                    [inside_s_span + span_mention_score_matrix.diagonal(width), # mention\n",
        "                     inside_s_span],                                            # non-mention\n",
        "                dim=0\n",
        "                )\n",
        "            )\n",
        "\n",
        "        inside_s = inside_s.permute(0,1,3,2) # (seq_len, seq_len, batch_size, 2), 2 for mention / non-mention\n",
        "        series_batchsize = torch.arange(0, batch_size, dtype=torch.long)\n",
        "        \n",
        "        Z = logsumexp(inside_s[0, sequence_lengths-1, series_batchsize], dim=-1) # (batch_size,)\n",
        "        \n",
        "        marginal = torch.autograd.grad(\n",
        "            Z.sum(),\n",
        "            span_mention_score_matrix,\n",
        "            create_graph=True,\n",
        "            only_inputs=True,\n",
        "            allow_unused=False,\n",
        "        )\n",
        "        marginal = marginal[0].squeeze()\n",
        "        \n",
        "        return (Z.view(-1), marginal.permute(2,0,1)) # Shape: (batch_size, seq_len, seq_len, ) \n",
        "          \n",
        "    def coolio(\n",
        "        self, \n",
        "        span_mention_score_matrix: torch.FloatTensor, \n",
        "        sequence_lengths: torch.IntTensor,\n",
        "    ) -> Tuple[torch.FloatTensor]:\n",
        "        \"\"\"\n",
        "            Parameters:\n",
        "                span_mention_score_matrix: shape (batch_size, sent_len, sent_len, score_dim)\n",
        "                    Score of each span being a span of interest. There are batch_size number\n",
        "                    of sentences in this document. And the maximum length of sentence is \n",
        "                    sent_len. \n",
        "                sequence_lengths: shape (batch_size, )\n",
        "                    The actual length of each sentence. \n",
        "        \"\"\"\n",
        "        span_mention_score_matrix.requires_grad_(True)\n",
        "        \n",
        "        batch_size, _, _, score_dim = span_mention_score_matrix.size()\n",
        "        seq_len = sequence_lengths.max()\n",
        "        # Shape: (batch_size, )\n",
        "        sequence_lengths = sequence_lengths.view(-1)\n",
        "        \n",
        "        rules = span_mention_score_matrix\n",
        "        log1p_exp_rules = torch.log1p(rules.squeeze(-1).exp())\n",
        "        \n",
        "        zero_rules = (rules.new_ones(seq_len, seq_len).tril(diagonal=-1))*(-LARGENUMBER)\n",
        "        zero_rules = zero_rules.unsqueeze(0).unsqueeze(-1).repeat(batch_size, 1,1,1)\n",
        "        \n",
        "        inside_s = torch.cat([rules.clone(), zero_rules], dim=3)\n",
        "        inside_s = inside_s.logsumexp(dim=3)\n",
        "            \n",
        "        for width in range(0, seq_len-1):\n",
        "            inside_s = logbmminside_rule(inside_s, log1p_exp_rules, width+1)\n",
        "            \n",
        "        series_batchsize = torch.arange(0, batch_size, dtype=torch.long)\n",
        "        Z = inside_s[series_batchsize, 0, sequence_lengths-1] # (batch_size, )\n",
        "        \n",
        "        marginal = torch.autograd.grad(\n",
        "            Z.sum(),\n",
        "            span_mention_score_matrix,\n",
        "            create_graph=True,\n",
        "            only_inputs=True,\n",
        "            allow_unused=False,\n",
        "        )\n",
        "        marginal = marginal[0].squeeze()\n",
        "        return (Z.view(-1), marginal)  # Shape: (batch_size, seq_len, seq_len, )"
      ],
      "metadata": {
        "id": "DlgNmgiI1jGL"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "mp = CKY()\n",
        "l = 128\n",
        "bs= 32\n",
        "lengthvec = torch.tensor([l]*bs, device=\"cuda:0\")\n",
        "example = torch.randn(bs,l,l,1,device=\"cuda:0\") + (-LARGENUMBER * (1-torch.ones(l,l, device=\"cuda:0\").triu())).unsqueeze(0).unsqueeze(-1) +\\\n",
        "(-LARGENUMBER * (torch.ones(l,l, device=\"cuda:0\").triu(31))).unsqueeze(0).unsqueeze(-1)\n",
        "\n",
        "r1 = mp.coolio(example,lengthvec)\n",
        "r2 = mp.io(example,lengthvec)\n",
        "print(torch.norm(r1[0] - r2[0]))\n",
        "print(torch.norm(r1[1] - r2[1]))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y0gBcLLk2Usw",
        "outputId": "fdaf4d39-d48f-4293-d6ab-58914f30aa41"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor(5.2858e-05, device='cuda:0', grad_fn=<CopyBackwards>)\n",
            "tensor(0.0003, device='cuda:0', grad_fn=<CopyBackwards>)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "with torch.autocast(enabled=True, device_type=\"cuda\"):\n",
        "    r2 = mp.io(example, lengthvec)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gFAs5AMG2lK6",
        "outputId": "8148880f-8bee-4bd5-8574-8a6695a9a9e4"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1 loop, best of 5: 161 ms per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "with torch.autocast(enabled=True, device_type=\"cuda\"):\n",
        "    r1 = mp.coolio(example, lengthvec)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JnXz_Pzj2diN",
        "outputId": "94bba9e6-2a8c-47fb-fa05-cb54e93742f8"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "10 loops, best of 5: 21.9 ms per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "McjDUqtq2jN6"
      },
      "execution_count": 7,
      "outputs": []
    }
  ]
}