sparsemax.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "sparsemax.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mvj1lEZbrDgm"
      },
      "source": [
        "!pip install tensorflow==1.15 sparsemax > /dev/null"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wc_9mvO9KB2e"
      },
      "source": [
        "import torch\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "from sparsemax import sparsemax"
      ],
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0HxTwQoJwzvM"
      },
      "source": [
        "## Sample data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uDME8eImroTB"
      },
      "source": [
        "np.random.seed(42)\n",
        "data = np.random.randn(10) # numpy\n",
        "data_t = torch.tensor([data]) # pytorch tensor"
      ],
      "execution_count": 36,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1P-aqDH8r9yj"
      },
      "source": [
        "## Tensorflow\n",
        "I'll be using TF 1.5 because is the last version that includes sparsemax on the contrib module.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EDA29VCYr-7q",
        "outputId": "fe6380ee-3c14-483f-dfbd-0f5ccb81f2fe"
      },
      "source": [
        "x = tf.constant(data[None, ...])\n",
        "output = tf.contrib.sparsemax.sparsemax(x)\n",
        "tf_result = tf.Session().run(output)\n",
        "tf_result"
      ],
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0.        , 0.        , 0.        , 0.47190852, 0.        ,\n",
              "        0.        , 0.52809148, 0.        , 0.        , 0.        ]])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YN0KR-SayHUy"
      },
      "source": [
        "## Pytorch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9ODNYg-byQK-"
      },
      "source": [
        "# github repo: https://github.com/aced125/sparsemax\n",
        "\n",
        "sp = sparsemax.SparsemaxFunction()\n"
      ],
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mWWxWnSy2hbj",
        "outputId": "ac74eca7-9987-400b-8308-1a7d7b46c322"
      },
      "source": [
        "torch_result = sp.apply(data_t)\n",
        "torch_result"
      ],
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[0.0000, 0.0000, 0.0000, 0.4719, 0.0000, 0.0000, 0.5281, 0.0000, 0.0000,\n",
              "         0.0000]], dtype=torch.float64)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uUy305jVygKy"
      },
      "source": [
        "## Custom"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jhAGg0CHn90h"
      },
      "source": [
        "class MySparsemax(torch.autograd.Function):\n",
        "\n",
        "  @staticmethod\n",
        "  def forward(ctx, input):\n",
        "    sorted, _ = torch.sort(input, dim=-1, descending=True)\n",
        "    cumsum = sorted.cumsum(dim=-1)\n",
        "    col_range = torch.arange(1, input.size(-1)+1)\n",
        "    is_gt = (1+col_range*sorted) > cumsum\n",
        "    kz = is_gt.sum(dim=-1, keepdim=True)\n",
        "\n",
        "    row_range = torch.arange(input.size(0))[..., None]\n",
        "    tau_z = (cumsum[row_range, kz-1]-1) / kz\n",
        "    output = (input - tau_z).clamp(0)\n",
        "    ctx.save_for_backward(output)\n",
        "    return output\n",
        "  \n",
        "  @staticmethod\n",
        "  def backward(ctx, grad_output):\n",
        "    output, *_ = ctx.saved_tensors\n",
        "\n",
        "    nonzeros = torch.ne(output, 0)\n",
        "    support_size = nonzeros.sum(dim=-1, keepdim=True)\n",
        "    v_hat = (grad_output * nonzeros).sum(-1, keepdim=True) / support_size\n",
        "\n",
        "    return nonzeros * (grad_output - v_hat), None"
      ],
      "execution_count": 40,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xUZpjbIqxYXS",
        "outputId": "481d4021-1bf8-4595-99aa-20607e379072"
      },
      "source": [
        "my_result = MySparsemax.apply(data_t)\n",
        "my_result"
      ],
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[0.0000, 0.0000, 0.0000, 0.4719, 0.0000, 0.0000, 0.5281, 0.0000, 0.0000,\n",
              "         0.0000]], dtype=torch.float64)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 41
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UhUTckl8ytxh"
      },
      "source": [
        "## Compare my result with the other implementations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "20yTJ5gIxajJ"
      },
      "source": [
        "np.testing.assert_allclose(my_result.data.numpy(), tf_result)\n",
        "np.testing.assert_allclose(my_result.data.numpy(), torch_result.data.numpy())"
      ],
      "execution_count": 42,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kGhIY5rQyyu-"
      },
      "source": [
        "## Check gradient"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B93eIgyZy33b",
        "outputId": "e1a4817a-49dd-4969-b62f-0ab576a68b1a"
      },
      "source": [
        "torch.random.manual_seed(42)\n",
        "inputs = torch.randn(2, 5, requires_grad=True).double()\n",
        "torch.autograd.gradcheck(MySparsemax.apply, inputs=inputs)"
      ],
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 43
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wFRdyDKhy8iw"
      },
      "source": [
        ""
      ],
      "execution_count": 43,
      "outputs": []
    }
  ]
}