docs/source/vision_transformers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e06e7b36-d1b0-4367-8ff0-68ff640649a6",
   "metadata": {},
   "source": [
    "# Sparse Vision Transformers with xformers\n",
    "\n",
    "In this notebook, we will illustrate how one can use the functionality available in xformers to\n",
    "quickly adapt an already existing Vision Transformer to use xformers efficient sparse kernels.\n",
    "\n",
    "In this notebook, you'll see:\n",
    "- how to get **memory and speed improvements** on a standard ViT model by using xformers with 80% sparsity\n",
    "- how to scale up a standard ViT model so that it can handle sequence **lengths larger than 10k**\n",
    "\n",
    "Let's start with a few imports. In this notebook, we will be using the `VisionTransformer` from timm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0ea655b8-5e01-41c1-b064-37525fe3c931",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import copy\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils import benchmark\n",
    "\n",
    "import xformers.components.attention.attention_patterns as AP\n",
    "from xformers.components.attention.core import scaled_dot_product_attention\n",
    "from xformers.components.attention._sputnik_sparse import SparseCS\n",
    "\n",
    "import timm\n",
    "from timm.models.vision_transformer import VisionTransformer\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec7139ee-c701-423a-910b-51b7c16e4b4f",
   "metadata": {},
   "source": [
    "To facilitate benchmarking and memory profiling, let's define a function that takes a generic callable and executes it, measuring the execution time and the GPU memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fef5fd04-5e83-4fb0-af21-379ab6d1f1fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def profile_model(fn, min_run_time=2):\n",
    "    torch.cuda.reset_peak_memory_stats()\n",
    "    torch.cuda.synchronize()\n",
    "    res = benchmark.Timer(\n",
    "        stmt='fn()',\n",
    "        globals={\"fn\": fn},\n",
    "        label=\"profile\",\n",
    "        sub_label=\"\",\n",
    "        description=\"\"\n",
    "    ).blocked_autorange(min_run_time=min_run_time)\n",
    "    torch.cuda.synchronize()\n",
    "    memory = torch.cuda.max_memory_allocated() / 2 ** 20\n",
    "    memory = f\"Memory used: {memory} MB\"\n",
    "    print(res)\n",
    "    print(memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b384c152-0e62-41cb-9c1e-611e8e5115fa",
   "metadata": {},
   "source": [
    "Now it comes the core of it. We will implement an `Attention` module following the same API and modules as timm's, but using our `scaled_dot_product_attention` function.\n",
    "It is worth noting that the `scaled_dot_product_attention` from xformers supports both dense and sparse tensors with the same API, so it can be a drop-in replacement\n",
    "for experimenting with both sparse and dense attention.\n",
    "\n",
    "Internally, `scaled_dot_product_attention` will dispatch to the sparse variants if the `attn_mask` is sparse, otherwise it will dispatch to the dense implementation.\n",
    "\n",
    "We provide support for both COO-sparse matrices (via PyTorch `sparse_coo_tensor`) and CSR matrices (via our custom implementation).\n",
    "The CSR implementation is significantly faster and should be preferred at pretty-much all times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1aaa3de6-4be6-4a5c-b8db-6cb6bd712984",
   "metadata": {},
   "outputs": [],
   "source": [
    "# almost drop-in replacement for timm attention\n",
    "# but using the sparsity-aware scaled_dot_product_attention from xformers\n",
    "class Attention(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        dim,\n",
    "        num_heads=8,\n",
    "        qkv_bias=False,\n",
    "        attn_drop=0.0,\n",
    "        proj_drop=0.0,\n",
    "        attn_mask=None,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.num_heads = num_heads\n",
    "\n",
    "        self.qkv = torch.nn.Linear(dim, dim * 3, bias=qkv_bias)\n",
    "        self.attn_drop = torch.nn.Dropout(attn_drop)\n",
    "        self.proj = torch.nn.Linear(dim, dim)\n",
    "        self.proj_drop = torch.nn.Dropout(proj_drop)\n",
    "        self.attn_mask = attn_mask\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, N, C = x.shape\n",
    "        qkv = (\n",
    "            self.qkv(x)\n",
    "            .reshape(B, N, 3, self.num_heads, C // self.num_heads)\n",
    "            .permute(2, 0, 3, 1, 4)\n",
    "        )\n",
    "\n",
    "        qkv = qkv.flatten(1, 2)\n",
    "\n",
    "        q, k, v = qkv.unbind()\n",
    "\n",
    "        x = scaled_dot_product_attention(q, k, v, self.attn_mask, dropout=self.attn_drop)\n",
    "        x = x.reshape(B, self.num_heads, N, C // self.num_heads)\n",
    "\n",
    "        x = x.transpose(1, 2).reshape(B, N, C)\n",
    "        x = self.proj(x)\n",
    "        x = self.proj_drop(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a853593-b3df-4526-871f-9967e1942394",
   "metadata": {},
   "source": [
    "Let's write a function that given a model, will replace all instances of `timm.models.vision_transformer.Attention` with our own implementation, which leverages `scaled_dot_product_attention` from xformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fccb2a4a-b864-452c-8649-7138855709af",
   "metadata": {},
   "outputs": [],
   "source": [
    "def replace_attn_with_xformers_one(module, att_mask):\n",
    "    module_output = module\n",
    "    if isinstance(module, timm.models.vision_transformer.Attention):\n",
    "        qkv = module.qkv\n",
    "        dim = qkv.weight.shape[1] * module.num_heads\n",
    "        module_output = Attention(dim, module.num_heads, attn_mask=att_mask)\n",
    "    for name, child in module.named_children():\n",
    "        module_output.add_module(name, replace_attn_with_xformers_one(child, att_mask))\n",
    "    del module\n",
    "    return module_output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ece835a-a715-4f44-94d4-0fdeb1bf1386",
   "metadata": {},
   "source": [
    "Now it's time to create our Vision Transformer. Nothing unusual here. Note that we will be keeping a copy of the model, which will be the model to use sparse self-attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "16934f00-841b-4bc0-a140-8c8ab7aada34",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_size = 224\n",
    "patch_size = 16\n",
    "\n",
    "model = VisionTransformer(img_size=img_size, patch_size=patch_size,\n",
    "                              embed_dim=96, depth=8, num_heads=8, mlp_ratio=3.,\n",
    "                              qkv_bias=False, norm_layer=nn.LayerNorm).cuda()\n",
    "\n",
    "model_sparse = copy.deepcopy(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0620e39a-1db5-407f-ad5d-8ddafcf9a8b9",
   "metadata": {},
   "source": [
    "What comes next is about creating a sparsity pattern that we will be using for our model.\n",
    "\n",
    "You can refer to the `2d_attention` tutorial for further information on how to build custom attention patterns with xformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ffe7f787-3398-47ba-86e6-b8d333713740",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequence length: 14x14 = 196\n",
      "Sparsity: 0.8061789721250534, nnz=1152\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/Z1A+gAAAACXBIWXMAAAsTAAALEwEAmpwYAAALfElEQVR4nO3dX6xlBXXH8e+vM8A4+AdoGyIMKTxQDCEIZkJRG9s4NIxIGB/6AJEGqsm8tBWNiYHwYPrWRGM0qdEQREkl8IBYCVFhihrTpFIHmFJgEChaGBgcWlIxEBkmrj6cQzLczIC5e599z7C+n+Tmnr3Pvnete5kf+885+65UFZLe/H5vrRuQNA3DLjVh2KUmDLvUhGGXmlg/ZbGjc0xt4NgpSy6FPz77pUFf/+gDG4/I2preb3iR/fVyDvVcpnzp7e05of4kWyartyzufGbXoK+/8KRzjsjamt49dTcv1POHDLuH8VIThl1qwrBLTQwKe5KtSX6W5PEkV4/VlKTxrTrsSdYBXwY+BJwJXJbkzLEakzSuIXv284DHq+qJqtoP3AJsG6ctSWMbEvaTgacOWt4zX/caSbYn2Zlk5yu8PKCcpCEWfoGuqq6rqs1Vtfkojll0OUmHMSTsTwOnHLS8ab5O0hIaEvafAqcnOS3J0cClwO3jtCVpbKt+b3xVHUjyt8CdwDrghqp6aLTOJI1q0I0wVfVd4Lsj9SJpgXwHndSEYZeamPR+9q68TVTLwD271IRhl5ow7FIThl1qwrBLTRh2qQnDLjVh2KUmDLvUhGGXmjDsUhOGXWrCsEtNGHapCcMuNWHYpSYMu9SEYZeaMOxSE0OmuJ6S5IdJHk7yUJKrxmxM0riG/MHJA8Cnq+q+JG8D7k2yo6oeHqk3SSNa9Z69qvZW1X3zx78GdnOIKa6SlsMof0o6yanAucA9h3huO7AdYAMbxygnaRUGX6BL8lbgW8Anq+qFlc87sllaDoPCnuQoZkG/qapuG6clSYsw5Gp8gK8Bu6vqC+O1JGkRhuzZ3w/8FfDBJLvmHxeN1JekkQ2Zz/6vQEbsRdIC+Q46qQnDLjVh2KUmDLvUhGGXmjDsUhOGXWrCsEtNGHapCcMuNWHYpSYMu9SEYZeaMOxSE4ZdasKwS00YdqkJwy41YdilJgy71IRhl5ow7FITY4x/Wpfk/iR3jNGQpMUYY89+FbMJrpKW2NBZb5uADwPXj9OOpEUZumf/IvAZ4LeH2yDJ9iQ7k+x8hZcHlpO0WkMGO14M7Kuqe19vO0c2S8th6GDHS5L8AriF2YDHb47SlaTRrTrsVXVNVW2qqlOBS4EfVNXlo3UmaVS+zi41seqRzQerqh8BPxrje0laDPfsUhOGXWpilMN46XDufGbXqr/2wpPOOWJrLyP37FIThl1qwrBLTRh2qQnDLjVh2KUmDLvUhGGXmjDsUhOGXWrCsEtNGHapCcMuNWHYpSa8xfUIMORWzbWuvZa3ir4Zb1Mdwj271IRhl5ow7FIThl1qYuhgx+OS3JrkkSS7k7x3rMYkjWvo1fgvAd+vqr9McjSwcYSeJC3AqsOe5B3AB4ArAapqP7B/nLYkjW3IYfxpwHPA15Pcn+T6JMeu3MiRzdJyGBL29cB7gK9U1bnAi8DVKzdyZLO0HIaEfQ+wp6rumS/fyiz8kpbQkJHNzwJPJTljvmoL8PAoXUka3dCr8X8H3DS/Ev8E8NfDW5K0CIPCXlW7gM3jtCJpkXwHndSEYZeaSFVNVmzzuzfUv995ymT1pG7Ou/Apdv7Hb3Ko59yzS00YdqkJwy41YdilJgy71IRhl5ow7FIThl1qwrBLTRh2qQnDLjVh2KUmDLvUhGGXmjDsUhOTzmd/9IGNzsxehbWczz7Ukfrf+0idS/9o/e9hn3PPLjVh2KUmDLvUxNCRzZ9K8lCSB5PcnGTDWI1JGteqw57kZOATwOaqOgtYB1w6VmOSxjX0MH498JYk65nNZn9meEuSFmHIrLengc8DTwJ7gV9V1V0rt3Nks7QchhzGHw9sYzan/STg2CSXr9zOkc3SchhyGH8B8POqeq6qXgFuA943TluSxjYk7E8C5yfZmCTMRjbvHqctSWMbcs5+D3ArcB/wn/Pvdd1IfUka2dCRzZ8FPjtSL5IWyHfQSU0YdqmJSW9x1eoMuV3ySL1Vc629GX9u9+xSE4ZdasKwS00YdqkJwy41YdilJgy71IRhl5ow7FIThl1qwrBLTRh2qQnDLjVh2KUmDLvUhGGXmjDsUhOGXWrCsEtNvGHYk9yQZF+SBw9ad0KSHUkem38+frFtShrqd9mzfwPYumLd1cDdVXU6cPd8WdISe8OwV9WPgedXrN4G3Dh/fCPwkXHbkjS21f4p6ROrau/88bPAiYfbMMl2YDvABjauspykoQZfoKuqAup1nndks7QEVhv2XyZ5J8D8877xWpK0CKsN++3AFfPHVwDfGacdSYvyu7z0djPwb8AZSfYk+TjwD8BfJHkMuGC+LGmJveEFuqq67DBPbRm5F0kL5DvopCYMu9SEI5sn4NhkLQP37FIThl1qwrBLTRh2qQnDLjVh2KUmDLvUhGGXmjDsUhOGXWrCsEtNGHapCcMuNWHYpSYMu9SE97NPwPvRtQzcs0tNGHapCcMuNbHakc2fS/JIkgeSfDvJcQvtUtJgqx3ZvAM4q6rOBh4Frhm5L0kjW9XI5qq6q6oOzBd/AmxaQG+SRjTGOfvHgO+N8H0kLdCg19mTXAscAG56nW2czy4tgVWHPcmVwMXAlvmM9kOqquuA6wDenhMOu52kxVpV2JNsBT4D/FlVvTRuS5IWYbUjm/8ReBuwI8muJF9dcJ+SBlrtyOavLaAXSQvkO+ikJgy71IRhl5ow7FIThl1qwrBLTRh2qQnDLjVh2KUmDLvUhGGXmjDsUhOGXWrCsEtNGHapCcMuNWHYpSYMu9REXucPw45fLHkO+O/X2eQPgP+ZqB1rW/vNWPuPquoPD/XEpGF/I0l2VtVma1vb2uPzMF5qwrBLTSxb2K+ztrWtvRhLdc4uaXGWbc8uaUEMu9TEUoQ9ydYkP0vyeJKrJ6x7SpIfJnk4yUNJrpqq9kE9rEtyf5I7Jq57XJJbkzySZHeS905Y+1Pz3/eDSW5OsmHB9W5Isi/JgwetOyHJjiSPzT8fP2Htz81/7w8k+XaS4xZRe6U1D3uSdcCXgQ8BZwKXJTlzovIHgE9X1ZnA+cDfTFj7VVcBuyeuCfAl4PtV9S7g3VP1kORk4BPA5qo6C1gHXLrgst8Atq5YdzVwd1WdDtw9X56q9g7grKo6G3gUuGZBtV9jzcMOnAc8XlVPVNV+4BZg2xSFq2pvVd03f/xrZv/gT56iNkCSTcCHgeunqjmv+w7gA8wHdFbV/qr6vwlbWA+8Jcl6YCPwzCKLVdWPgedXrN4G3Dh/fCPwkalqV9VdVXVgvvgTYNMiaq+0DGE/GXjqoOU9TBi4VyU5FTgXuGfCsl9kNuf+txPWBDgNeA74+vwU4vokx05RuKqeBj4PPAnsBX5VVXdNUXuFE6tq7/zxs8CJa9ADwMeA701RaBnCvuaSvBX4FvDJqnphopoXA/uq6t4p6q2wHngP8JWqOhd4kcUdxr7G/Nx4G7P/4ZwEHJvk8ilqH07NXn+e/DXoJNcyO5W8aYp6yxD2p4FTDlreNF83iSRHMQv6TVV121R1gfcDlyT5BbNTlw8m+eZEtfcAe6rq1aOYW5mFfwoXAD+vqueq6hXgNuB9E9U+2C+TvBNg/nnflMWTXAlcDHy0JnqzyzKE/afA6UlOS3I0s4s1t09ROEmYnbfurqovTFHzVVV1TVVtqqpTmf3MP6iqSfZwVfUs8FSSM+artgAPT1Gb2eH7+Uk2zn//W1ibC5S3A1fMH18BfGeqwkm2Mjt9u6SqXpqqLlW15h/ARcyuSv4XcO2Edf+U2eHbA8Cu+cdFa/Dz/zlwx8Q1zwF2zn/2fwaOn7D23wOPAA8C/wQcs+B6NzO7PvAKs6OajwO/z+wq/GPAvwAnTFj7cWbXqV79N/fVKX7vvl1WamIZDuMlTcCwS00YdqkJwy41YdilJgy71IRhl5r4f997Y5TiItSwAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "H, W = img_size // patch_size, img_size // patch_size\n",
    "print(f\"Sequence length: {H}x{W} = {H * W}\")\n",
    "\n",
    "axial_pattern = AP.axial_2d_pattern(H, W)\n",
    "loc_2d_dist = AP.local_2d_pattern(H, W, distance=2, p=2.0)\n",
    "rand_pattern = torch.rand((H * W) ** 2).reshape(H * W, H * W) > 0.99\n",
    "\n",
    "gaus_2d_dist = AP.local_2d_gausian_distribution(H, W, sigma=5)\n",
    "sparsity = 0.97\n",
    "num_non_zeros = int((H * W) ** 2 * (1 - sparsity))\n",
    "random_gaus_2d_pattern = AP.random_pattern_from_probability_matrix(gaus_2d_dist, num_non_zeros)\n",
    "\n",
    "\n",
    "t_mask = axial_pattern | loc_2d_dist | rand_pattern | random_gaus_2d_pattern\n",
    "\n",
    "# and let's not forget to add a global attention for the cls_token\n",
    "mask = torch.ones((H * W + 1, H * W + 1), dtype=torch.bool)\n",
    "mask[1:, 1:] = t_mask\n",
    "\n",
    "print(f\"Sparsity: {1 - mask.float().mean().item()}, nnz={num_non_zeros}\")\n",
    "\n",
    "plt.imshow(mask[H * W // 2 + W // 2][1:].reshape(H, W))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "885350c5-3f16-49cd-89c6-55c53fb3d77e",
   "metadata": {},
   "source": [
    "From the print just above, we can see that the full `attn_mask` has a sparsity level of around 80%.\n",
    "\n",
    "Once we are happy with the sparsity pattern, the next step is to create the sparse matrix that holds our sparsity pattern.\n",
    "\n",
    "Everything is handled by `SparseCS` class.\n",
    "\n",
    "**Note: the optimized kernels require that the total number of nonzero elements in the full `attn_mask` to be a multiple of 4. This is handled by default in the implementation (by removing elements until the number of non zeros is a multiple of 4)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ef202f12-7622-42d3-bbf6-01fa964c5b21",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8061789791027855\n"
     ]
    }
   ],
   "source": [
    "mask = SparseCS(mask, torch.device(\"cuda\"))\n",
    "print(1 - mask.values.shape[1] / (mask.shape[0] * mask.shape[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce894e15-97a1-4490-98a9-4c4d039e4cc3",
   "metadata": {},
   "source": [
    "Now we are ready to replace the dense attentions with with sparse versions which leverage the sparsity pattern that we have just created"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "232d065e-afc5-4d7e-88ee-7e7fad875c07",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_sparse = replace_attn_with_xformers_one(model_sparse, mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a809363-a1c4-4bda-ad0a-610c078d7e95",
   "metadata": {},
   "source": [
    "Let's new create an input tensor and benchmark both the sparse and the dense versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "62c6d031-f4a1-4b89-a388-3b8f2f3c17a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = torch.rand(64, 3, img_size, img_size).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b420967-81ae-4c2a-ae20-7690c575c392",
   "metadata": {},
   "source": [
    "## Profiling the default (dense) model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f3f150fe-5b2c-49b2-92a6-285a6f67205e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward only\n",
      "<torch.utils.benchmark.utils.common.Measurement object at 0x7fbdf375a400>\n",
      "profile\n",
      "  Median: 22.70 ms\n",
      "  IQR:    0.09 ms (22.68 to 22.77)\n",
      "  88 measurements, 1 runs per measurement, 1 thread\n",
      "Memory used: 1790.81689453125 MB\n",
      "\n",
      "Forward + backward\n",
      "<torch.utils.benchmark.utils.common.Measurement object at 0x7fbdf0464760>\n",
      "profile\n",
      "  Median: 77.05 ms\n",
      "  3 measurements, 10 runs per measurement, 1 thread\n",
      "Memory used: 1975.78173828125 MB\n"
     ]
    }
   ],
   "source": [
    "print(\"Forward only\")\n",
    "profile_model(lambda : model(i))\n",
    "print(\"\")\n",
    "print(\"Forward + backward\")\n",
    "profile_model(lambda : model(i).sum().backward())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1959b60b-aa09-4823-be4a-32ff24910d28",
   "metadata": {},
   "source": [
    "## Profiling the sparsity-aware model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fab86e1c-fff6-4457-aaee-9ad8303ee2df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward only\n",
      "<torch.utils.benchmark.utils.common.Measurement object at 0x7fbdf0406460>\n",
      "profile\n",
      "  Median: 16.38 ms\n",
      "  IQR:    0.05 ms (16.37 to 16.42)\n",
      "  13 measurements, 10 runs per measurement, 1 thread\n",
      "Memory used: 822.93701171875 MB\n",
      "\n",
      "Forward + backward\n",
      "<torch.utils.benchmark.utils.common.Measurement object at 0x7fbdf375a040>\n",
      "profile\n",
      "  Median: 53.41 ms\n",
      "  IQR:    0.01 ms (53.40 to 53.41)\n",
      "  4 measurements, 10 runs per measurement, 1 thread\n",
      "Memory used: 835.51025390625 MB\n"
     ]
    }
   ],
   "source": [
    "print(\"Forward only\")\n",
    "profile_model(lambda : model_sparse(i))\n",
    "print(\"\")\n",
    "print(\"Forward + backward\")\n",
    "profile_model(lambda : model_sparse(i).sum().backward())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65f51323-0ab6-4af2-b731-dad87a65789d",
   "metadata": {},
   "source": [
    "This shows that not only the sparse model is ~30% faster, but it also uses less than half the memory compared to the dense counterpart. Great!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9654e78-daa2-4023-848e-4c4ae8b6e500",
   "metadata": {},
   "source": [
    "# Scaling ViT to sequence lengths of size larger than 10k\n",
    "\n",
    "Getting speed-ups for training / inference on standard workloads thanks to sparsity is nice, but by leveraging the memory savings from sparse computations we can actually easily scale up ViT to sequence lengths which are much larger than what was originally possible.\n",
    "\n",
    "Let's get back to the original model, but this time scale it up so that the sequence length is `112 x 112`, and use a 99.5% sparsity pattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a6cc8a1e-18f1-423c-af1e-2f7e88d3b332",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_size = 224\n",
    "patch_size = 2\n",
    "\n",
    "model_sparse = VisionTransformer(img_size=img_size, patch_size=patch_size,\n",
    "                              embed_dim=96, depth=8, num_heads=8, mlp_ratio=3.,\n",
    "                              qkv_bias=False, norm_layer=nn.LayerNorm).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1922754f-9a0b-42b3-9819-4a729c8611b0",
   "metadata": {},
   "source": [
    "Everything was the same as before, with the difference that the patch size is now `2x2`, which means that the sequence length is `112x112`, which is much larger than what has been used so far.\n",
    "\n",
    "Let's create a sparsity pattern to have roughly 99.5% sparsity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a0ff06e2-9485-4e80-b18c-88b1531d904c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequence length: 112x112 = 12544\n",
      "Sparsity: 0.9938480411656201, nnz=786759\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQEAAAD7CAYAAABqkiE2AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/Z1A+gAAAACXBIWXMAAAsTAAALEwEAmpwYAAAkW0lEQVR4nO3deXxU5bnA8d8zS2ayryQEQtgSVmUzVVCrWGtdrla9t4tLK7Uq3LpXrVVv61Lbe2urtva2WlHx2rqhiIpatUpta6tQRPY9QEISyAaBhOwz894/ZsQACUlmyZnJPN/PJ5/MnDnLwwl58p53FWMMSqn4ZbM6AKWUtTQJKBXnNAkoFec0CSgV5zQJKBXnNAkoFecikgRE5BwR2SIipSJyRySuoZQKDwl3PwERsQNbgbOASmAFcKkxZmNYL6SUCgtHBM55IlBqjNkBICIvAhcCPSaBBHEZN8kRCEUp9ZkmGuqNMUOO3B6JJDAcqOjyvhI46cidRGQuMBfATRInyZkRCEUp9Zn3zaLy7rZbVjFojJlvjCkxxpQ4cVkVhlJxLxJJoAoY0eV9QWCbUioKRSIJrACKRWS0iCQAlwBLInAdpVQYhL1OwBjjEZHrgXcBO7DAGLMh3NdRSoVHJCoGMcb8CfhTJM6tlAov7TGoVJzTJKBUnIvI44Dqnu/06ewf42bIsnq8m7ZZHY5SgJYEBo4I278rvPWTB9lz5lGdtpSyjJYEBooxpK5yMzt5Hrk7O62ORqlDNAkMoKG/WY78VjBer9WhKHWIJoGB5PNifFYHodThtE5AqTinSUCpOKdJQKk4p0lAqTinSUCpOKdJQKk4p0lAqTinSUCpOKdJQKk4p0lAqTinScBC4nAgzgQQsTqUfhNngj92FfM0CVhEnAmU3f0Fql4qwnfaNKvD6RfH6JFsfeJ4tvx2Ko6heVaHo0KkScAiYreRNG0fz01fQPPQINZdsNmxJSdjc7vDH1wvfGlJ3Hrin7n0xOWYlKQBv74KLx1FaBFfezuZv05h3tCbyfqoEk9/T1Ayifq729m7I5Pxd67H19wciTC7JTurePFH5yFeQ/IenUg61mkSsIoxOP6ykjTofwIAOrLc3D/hJX7pPhtxDuyP0dvYSNLi5QDoyOjYp0kgRiWuLON/vj+HhMZOfAd3Wx2OimGaBGKUt64O9xt1AIR3cXkVNBHsOTkAePfuA19szCClFYNKhYmjYDhVT+Rw4I+p2ItGWR1On2kSUCpcHHam5lUxc0gZDHA9TShiJ1Klopx3dzW1NxRTYxNs22NnXQlNAkqFiWlvh0/WA7HVaqKPA0rFOS0JqKjhGFWINycNW3kN3ro6q8OJG0GXBERkhIh8ICIbRWSDiNwU2J4lIu+JyLbA98zwhasGs803D+M7z71F7UVFVocSV0J5HPAAtxpjJgEzgetEZBJwB7DUGFMMLA28V6pX9jZhU9swbB1WRxJfgk4Cxpg9xphPA6+bgE3AcOBC4JnAbs8AF4UYo4oTRQ9u5ZOLishevN7qUOJKWOoERGQUMB1YDuQZY/YEPqoGuh1rKiJzgbkAbnQkmgJv/V6o32t1GN2yTZtEW14SSWsr8eyptjqcsAq5dUBEUoBXgJuNMY1dPzPGGHro1WqMmW+MKTHGlDgJYiitUgNEHA623ubivseepP6s0VaHE3YhJQERceJPAM8ZYxYHNteISH7g83ygNrQQlbKW8RnslW4erjgbd0NsjAfoj6AfB0REgKeATcaYh7t8tASYA/w88P31kCJUymo+L2PvX0uH20Vi4+pBN2ArlDqBU4BvA+tEZHVg2134f/lfEpGrgHLgGyFFqAaOzY7njGm05DrJ+rACT2WV1RFFDV9zMwzgxC0DKegkYIz5B9DTDJlnBnteZR2b20XzrQf43cTnuem2G0jWJBAXtNuwOsR0emhYNYSr117Bvol29l05C0fBcKvDUhGm3YbVIaazg9H3rMCWnUXCSw3cVfgmt9dci0tLBIOaJoEwsE2bxO7ZGWRt7iThnRVWhxMS4/Fgmg6y6cOpXFI4j/G7m2NqRJzqP00CYbB7dgb/vO1hprx/HcXvWB1N6HwtLYy6+1/+1zEyRZYKniaBMMja3MmU968j86NBtCKP/vLHDU0CYZDwzopBUQJQ8UlbB5SKc5oElIpzmgSUinOaBJSKc5oElIpzmgRU7BLxf6mQaBOhikn2nGy23T4Ob6KPiQ/twVO2y+qQYpaWBFRMktQULjv779z/lUV4c9KsDiemaUlAxSRf/T7eefA03kqA3LKtaP/G4GkSUDHJ19RExh8/BtAEECJ9HFAqzmlJQMU2EWxJ/inrfS0tYAbbDICRpyUBFdMchQVsf7KI7U8W4SgssDqcmKQlARXTTJKbr01YBcDqpPEWRxObNAmomGbKq1j2gy8A4CrfZHE0sUmTQBjYkpKwZWVimlvwNjQc2m5PS0PSUvEdaMTX1GRhhP1nHzIESXDira3HdEbvCqG+lhac76/0v7Y4llildQJhsO8/plL8eg3bb51w2PZd3zuOSUt2U33F8RZFFhxbaiqbHxxB6ktteE+aZHU4KsI0CYSBJwnOSl9PZ/rhf4s6Uwxnpa2nM9WiwIIkImRmHeSLmdvwuu097+dMwDF6JI4RBWDreT8V3cREQZNKmmSZkyR21ytxDM2jc2w+jpoDeEt3fr59VCGdw7NwltfF1mo+Nju2ScX4UhKwbyzD29jY7W4yfTIZv93Dlvpchs1rwFNdM8CBqv543yxaaYwpOXK71gmEgae6BqmuOarnmqdsF1K2C48lUYXA58W3fjMQ6I1ns2MfOxIcdnw7dmHa2/372YVcVxO73emWhapCp0lA9cqRn0fz73yMTatl97wxmDX+WnhZX8q274wlpdODt36vxVGqYGkSUL3z+Wj32mn1Og/f3NYGgRJDd+zZWfjGDMfWcPCwxyQVXbRiUPXKU1NH1rUe9s/JwGzc3ufjGs4ex7ef/ROb7sxGHPr3JlrpTybG2bOz6DxuJI6GVnzrtkSm77zPi2dnef9jazd8enAkctCB8VlTAW3PzqJz0kgcjW341m7WsQXdCLkkICJ2EVklIm8G3o8WkeUiUioiC0VkEC3LE30OnlrENU8spuJeOzaXy+pwDpP69jo2XzSMCfdvs2xFo+ZZRVz15Gvsvg9siYmWxBDtwvE4cBPQtb/mA8CvjDFFQANwVRiuoXrgaPayqLaE5upkoqG5tytfSwueikpLKw2dzR4W1Z5AY3Uq+LRPYXdC6icgIgXAM8DPgFuAC4A6YKgxxiMis4B7jTFnH+s8sd5PwEriTMCWkQ6dHXj3H7A6nKij9+dzkeon8GvgduCzPnHZwH5jzGdN45XA8O4OFJG5wFwAN0n9vrB93FgOTM0htfQgZtWGfh8/WJjODrx1dVaHERTblAk0Tsggff0+vBu3RuQasXx/BkrQjwMicj5Qa4xZGczxxpj5xpgSY0yJk/4/y1Z8NY8XHnqIrVenBHN5FQW2X5bJwgcfZMc3c6wOJa6FUhI4BfiqiJwHuIE04BEgQ0QcgdJAARCR/rIpVT7mbLmclO3aZz3a2ceNZe/MXNJ2tGL7x+pD21PK4NtbLie1PLrqMuJNWMYOiMhs4DZjzPki8jLwijHmRRH5PbDWGPPosY4Ppk5AHA4kIQHT6Ynqoa4Kaq8/mT/d/gtO+eBGir+z6lAz3aGfYUcHxhNznatjTk91ApHoLPRD4BYRKcVfR/BUBK6B8XjwtbRoAugjM2sqtdeeDDOnDPi108o9nPPpNaSsdR/WTn/oZ6gJwFJh6SxkjPkr8NfA6x3AieE4rwqfnRcmseZbv2baH29i9LKBvbb7zRXkv+PEeHVy8GikPQbjRPZ6w/R/Xk3W+t4f/8SZwIGvz6A1x8bwJZWhL/FljJbYopgmgTiR/uwy0p+TPnWbtSUnkv+f27m38A1u3H4DLl3nb1DTJBBP+lgJbNra2fLOZC4suJEJ5Y1hm7vPnDKNijOTyP3Ug/vNf4XprCpUmgTUUXxtbRT890f+12E8b9VpSayf91uK3prHuDfDeGIVEk0CFhNnAnXfPYHm4TDmuTq8W0oBsGdmUnH1RNqy/X+93fuEEU9txrt3n5XhhiRvZQfFr36PoR9Jn48RZwJ1V57AwREw9oW9EetZGM80CVjMlugm/9IyHh69iKtXfp/ELYEPcjK5cs473JK1A4D/bRjJ269/AWI4CTj//AnFf+7fMeJ2kfPNChYVvci3V99K8sbIxBbPNAlYzNfaRvXCyVww/DbGbq77fJ7Cfft5+g/n8FiOvyTg2icU1vc8i09IbHbqrzmR/RMMxc83Y1asO+xje0Y65ddOxpNsKPp9BZ6KysjE0Q3T1k79SyM4Z8QPGLtpr65AHAE627BCnAnUvjKapdOf5uwf3UrmMx8f9rlj5AhKluzgxOTt/Oayr8O/1vVwJhXNdLZh1SPj6cT1YianrLiNkav2HVUZ6Nu3nyXzT2dx4ukUVuyIvdmT1TFpElBgDGnPLyON7lsDfE1N5P7O31rQawIQOXTOPpG+9V1QkaMTjaqwkRMms3XBDHb9eBY2t7vX/fd/exY7n59C0yUzByA61RNNAlFEHI7+z8orgjgTomIZsJaCZF6d/SiFp+9CEnqfWnLvFFj9xSfYe1zfmwxV+OnjQJSwpaay7Z7JkN9O8QOt/plx+8DMmkL5zT7YnMKo+1da2kc/dfUervqfm0mq85HcUt3r/mNebeGUspsZtfLgAESneqJJIEqIK4GSWVv5bt6HPJBzRZ9/MC3D3Lx04iPMS70ccTosTQKe8gpy5lcA0JenfPloDbkf9f384nIhIvja27UeIYw0CUQJ34Emqv97KvekjSNzw44+t4enL69i7o9uJqnWg6+tLJIhWsqWmsqWn00ipbCRYT+1Yz5Zb3VIg8agTwK2pCQkwYnvYHNUT15hOjtwvbUCF/SrQ4ynopL0Zweu845VJMHJhOMruG7EX3go61s4ez9E9dHgrhi02dlx11RsryXRcv4Mq6NRIfAdaKT9/nweuO0K3Kv6vxqS6tngTgJAZ2E7d49cQmum9bXnPbFnpGMfMsRfy38kmx17Tjb27KzP2+DjkPF4cPxlJYmv/UunEA+zwZ0EfF7GP9zG7dddS+570Tkxhi0piU2/GEfCIhvemZOP+tw+ppC6/8um6qk8HIUFFkSoBrtBXyfgW70R1+o+9HSzit1O5tBG/j1vJU8nF3JUWSDByWn5pRz0uqhwDo1cGDnZSGIivrp6/5LjIbC53diG5GDa2qPmr/ahmFpbLV0WLRoN7pJADPA1t5B3t50/XnMBictLj/rc7NjFuuuPp+zGYnwVuyMSgzgcbL63iOLXqmk+J/TZiFu/dDyjF9ex6f7RSJQsktpy1hTGvlrDpvvHdv/YFcc0CYSBPSMd+7ix2PNyj7mfY2ief7+0tM83+rz4Vm/E9uEqvA0NRx3ja2tDPl4Dy9Zi2tvDHfrnUjuZmbKdpuEO7MVjsCX1f2m4z3hdNqanlONIi57JRb0uoSRlJ7bUTqtDiTqaBMKg7uJJfHnxKrbdOrbnnUTY8oPRfHnxKuovPvrZ30rG42HCT/fz1NUX0fGlA8x4eRutpwcfY+pft/LyFWdRfG8TpiM6EkHa0q08/+1zGX/fAZ35+AiDvk5gIBg7pNtb8DmO3YvN2CHF3oaxMPU6hg/DNyQDqao77HldOj3Y2j3Y7T6yHM00jXCQOGUCbK/A19zcr2t4GxpgRUPQE4A4Ro7Al5GC7KrutnQUjFBjGsy0JBAGua9sZtHXZzPhoWO0QBjDhAcrePXrpzHktQjNENQHpd8byZyFb1N1efGhbeJwsPHuXP7jmaXYlmby7tWnknBxLac8u5qOkyYMbIA2OxvvzOfiF//GvvPGD+y145SWBMLA29AADQ29zszrqazyL9ZuIVunsLszE9uRzSUdNqo6Mkmt8mLbsJO9+8dQnZOO+CLXR98xooDOEdk4d9X7702AdAi1nWnY9M/2gNDpxeKMPScbMtOhbi/e/QcObXcMzcOkJGH21OJracExLB9cCfh2V4fcZNiT8vtO5jeXP8H3F1xzaIpzwF/BmpqMqanH19QUkWvHI51eLEbZkpLwTh8HxmD/dEuffyEdBcNpGz8UV+WBQ9OYA/428m7ayT3VNYe/r9rtn9tg+kR8SQ6cq7bjbWwM7R9zBOdBWNo0GccRVQ7emlqoFWxTJuDJTsK1blfQ/Q3s2Vl0TBmFY387ZvVGHX3YDa0TiHIyqoCp/7uGoke2IAX5fT5u90Uj+dkT89l8fXbQ17ZlZsID+7josffpnDIm6PP0ZMTj61l/wXCGLTh64lJJSGD7nU5umf8cB844RqtLL5pnFfG9xxdRc68XW2JiKOEOWloSiHLS1sHSynF0eByMbD/Q+wEBrv0+Hq+djbsmhDETNmFM6l5OcJfxqjP8fy+8jY3QU+nCZ/DuSeLpPafibAq+csDZ7OGPe2axf3caQ32f19o4RhVy8PihJFUcxLc6vhczCCkJiEgG8CRwHP55JL4LbAEWAqOAMuAbxpjwtPPEIU9ZBUPn+DsXefbv7/NxGS+voubdNEa2rAnrUmIDxXR2MP6ejbQkOHEfWNOnSUq64/hwLR3/kc7Ejs14uzxK7T6/gGdve4jz372RcfPCE3OsCrUk8AjwjjHmayKSACQBdwFLjTE/F5E7gDuAH4Z4nfjl8wbVVm7aw9Bvv72dt1cfz/KhheQ2tA14MglHHYTxeLodK5BU6+OHZf9OUnkQMxPY7HhmT6M9y0H6P8qOqk+JNUG3DohIOrAaGGO6nEREtgCzjTF7RCQf+Ksx5pgNvto6EL1sSUlgt+NrbgHf4GmzE4cDW1ISpqOj360fttRUWl/J5idFr3H3DdfgemtFhKIMr0i0DowG6oCnRWQqsBK4CcgzxuwJ7FMN5HV3sIjMBeYCuAm+n7qKLF9LS4+f2XOyafhKMY5WQ8rba8LWlOg9YwYHRrkY8o8avNt2hOWcRzIeT/AlDa+XyjX5XN9yKfl7IzieY4CEUtvjAGYAjxljpgPN+Iv+hwRKCN0WNYwx840xJcaYEifRMdJM9Y9nXAG33fs8x9+1BltO8K0Qh7HZ2TXPw6v3/ZLa2d3+/bCcr6WFortWMvySnbA89pdkC6UkUAlUGmOWB94vwp8EakQkv8vjQG2oQaro5Kg/yG3/+Aa2/Q7Gt/ZjyfCZU6g5MYUhq1qxfbjq8M+MD9fKFM6yf4/8cutH/HWcXULDuATy/7bvsGngTWcHxvrwwiLoJGCMqRaRChEZb4zZApwJbAx8zQF+Hvj+elgiVVHHu3U74+c6MD6Dtx/1BeXnprDm6keYtPAGij484kNjGPbQcsQmGK/FdRAi7PqWl0/PeIjTuZW8tdaGEymhtg7cADwXaBnYAVyJ/xHjJRG5CigHvhHiNVQYdX6lhJoTEhj+9xbkn6tDPl8wMzhnr/cy+a9zyVnVw5yJPi+ml6YIe1oatd+cjNcl5L9c6u9leKz9J42j4vwc0sp8pCxa0bdKTmNI+ziRGVxL4dbBO/w4pCRgjFkNHFXbiL9UoKJQ2VftbLv4txznvJ7Cf1oTQ8rLy0lZFNpCpJKTxQXX/42JiVUsWHYB9JIE9s3IYukNv+Tf1lyJvO7AtPetlJH76EfkPja4F03VHoNxJu8joShhHoWfWjzrYoi/VOZAIy+8fjpeF4yrqeh1Dsn0rc2c9Mb3Sd/swHRu7+fFoisBNH1zJvsm2Rj1ZhNmRegVkzqKUKlYIsLOF45n9alPcsp/30zuo31fx01HEUaBA5fPpH6aMPr1Vmz/WG11OD0Sl4vquSfQkm8oeqbusFGIymLGkPF2MsfvvpGxa3ruw9EfOopwoIhQf147Gy77DXUzortzlC3RTdE3tvL6ZQ/RUpRpdTjqCJnPfEzRLcvCUrELWhIYOMaQu8TFlJ03MnJF/+bsG2i+1jZ2PjuJC0fdiudcH7YzZjLu8dqI9d5T1tIkMIBSFy4jdaHVUfTOtLeT8/jH5OXlMuKNJq7I+Sf3vHU19m1WR6YiQZOA6pFpOsjKx6fxccZ0RpTuit5VnFRINAmoHvlaWshesAzEhqe/IwhtdjC+qGteCzsREFtMj7DUikHVI3taGjt+PpOtT03DPmlcn48zp0xj+7PHU3HXrEG95JctNZWd/zOTbQumYZ8cu9OjaxJQPUt0c+rp63n8i3+gIzel9/1FEGcCjaPd/OmU35Fych1iH7z/xcSVQMkXN7Pgi0/TnteH+xOltLOQ6pG4XBw8fxodKTaGvLuj9xl0Zk6h4jYfrU0uMla4SC/rxPXOpzFdVD4Wcblo/rdptKfZGPJOH+6PxbSzkOo3095O8ivLSaZvS7u3DXHzxIzf86PtF5N4U0PYpyiPNqa9naTFy0mib/cnWmkSUGGT8skufnDXtbj3efC1VPV+gIoKmgSijQj21FSwCd7Gg9FRlO5jTJ491aS+WA30MJ2UikqDt9YmRtlzctjy6Fjqn8vFPrHI6nAAsOcOYevvx1D3bB72CcEvBKKikyaBKCMJTr4wupzLRq3AmxwdzWvidnH55BX8fOIrtIxMxz5kiL8fgBoUNAlEGW9tPftuL+TN67+EbVOZ1eEcpsR1kK888Hf2PZOOfUyh1eGoMNE6gShjOjuQf67GDgO62Id9yBAkORFfbf3R04x7fWxozGdNaiKnpmyhOjeNUueoAYxORZImAYU4HGy6bzRfm7mCZfefSNKryw/73FtTx8Gbi7k/5UoAbK0ebDt1NNFgoUkgxtmSk5HCYUhLG57yiuBPlOhlQuIePnIePfmn6eyAlRsOe3bsWkqxud3IyAKkoxNP2a7BP15gkNE6gRjXevokTnpxA5t/moPN7Q7qHMbjYdL9e3n5W2eS/t6Wfh/vPWECRc+Vs+vhZOwZGUHFoKyjJQEFgGdHGQAD2SvBMXwY3vws7FX1ePZUD+CVVVdaEohxiX/byPJLJjPhR/VhWwuwv+wrN1N6+UgKb2nG24/l03dcNYprnl9C+ZwxkQtO9UpLAjHO19wMm8JUSSeCbfJ4vOluHBvL+7wkuq+tDYKYjNTeDpvahmEPYk1PcSbA1HEYhw3b2tJjLpyqjk1LAuoQm8tF+b0O5ix4g4OnFUf8eoXzN/HxRRMY/vSGfh9rH5ZH2q/3MOXRdVA8MgLRxQ8tCahDjDG01ifx9t7jcTRHvnbA29AAvZU2bHZkxkQ6Mly4V5fhrd/r3+7xsrF2KNXJaaR0xPIYPutpElCHmPZ2Jv5XKQ0JbhL2rY+KQUC2RDc1d3dy98TX+NUtl+F+w58EPLv3UDiv0z+oqW6vxVHGNk0C6jDevfusDuFwXi/7d2XwWPJsnE1d/uIbg7euzrq4BhFNAiqq+draGH/nRsThQJrWRkXpZLDRJKCinq+pyeoQBrWQWgdE5PsiskFE1ovICyLiFpHRIrJcREpFZKGIRMd4WKVUt4JOAiIyHLgRKDHGHAfYgUuAB4BfGWOKgAbgqnAEqlQ42McXsfeaWXR++QSrQ4kaofYTcACJIuIAkoA9wJeARYHPnwEuCvEaSoVN7WlDeOPHv6Ty6k6dGCUg6DoBY0yViDwI7AJagT8DK4H9xpjPqnErgeHdHS8ic4G5AG56XqXXXjyG3ecMJbXKS9Jrn0THnHuDnDgTaPz3GbTm2Bj2ZkVooxOjTNrODk7/6FrcK1L8KySpkB4HMoELgdHAMCAZOKevxxtj5htjSowxJU5cPe7XUJLLq7f9gs6r92JLcAYbruoHW3Iiudfu5OnbfkXz5KFWhxNWzvdXMvqyDeQ//LEOeQ4I5XHgy8BOY0ydMaYTWAycAmQEHg8ACoCQ5p5O29nCl967meYPcjEe7Rk2EExbO1vfG8vF795AYuUgrJn3eTUBdBFKE+EuYKaIJOF/HDgT+AT4APga8CIwB3g9pAiXrWXcMv9L/bENDF9bGyN++pH/tcWxqMgLuiRgjFmOvwLwU2Bd4FzzgR8Ct4hIKZANPBWGONUR7MVj2HX3ydTPm4W4en6cUqo3IXUWMsbcA9xzxOYdwImhnFf1rmVcNou/+yA/rvgqLQvdeNuDGI+rFNpjMGYllTZw4fO3krRHGNq60upwVAzTJBCjvFtKGX2XfyIPrStRodBJRZSKc5oElIpzmgSUinOaBJSKc5oElIpzmgSUinOaBJSKc5oElIpzmgSUinOaBJSKc5oElIpzmgRU1LFnZuIoGI7N7bY6lLigSUBFF5udrf81ntGv1tN4wVSro4kLmgRU1PFmeDgjfROeRP3vORB0KLGKLj4vE3/RwOMLLiZ7Wyk6t3TkaRJQUce7dTuyFU0AA0TLW0rFOU0CSsU5TQJKxTmtE4hx9uwsOo8biaOhFd+6Lbqohuo3LQnEuIOnFnHNE4upuNeOTdcfUEHQJBDjHM1eFtWW0FydjNFSgAqCPg7EuIS/rePgxelM7NyiC5CooGgSiHGmswNvXZ3VYUQNe2YmB08rxtHsJeHD9RhNjL3SxwE1qHgmjWTeL15h2E+2Y8/KtDqcmKAlATWo2A+08bN159Jan8TEjlKrw4kJmgTUoOLbsIWR33JhjNE6kj7q9XFARBaISK2IrO+yLUtE3hORbYHvmYHtIiK/EZFSEVkrIjMiGbxSRzEGX1ub1gX0Q1/qBP4POOeIbXcAS40xxcDSwHuAc4HiwNdc4LHwhKmUipRek4Ax5u/AviM2Xwg8E3j9DHBRl+1/MH7LgAwRyQ9TrEqpCAi2dSDPGLMn8LoayAu8Hg5UdNmvMrDtKCIyV0Q+EZFPOtGim1JWCbli0BhjRKTfXdWMMfOB+QBpkqVd3ZTqQft5X6CmxEnBB63YPlwV9vMHWxKo+ayYH/heG9heBYzosl9BYJtSKki7/s3Gpv98lMozEiNy/mCTwBJgTuD1HOD1LtuvCLQSzAQOdHlsiDmOkSOo+PHJVN90MrbUVKvDUXFq6N+FMa/OI++Tzoicv9fHARF5AZgN5IhIJXAP8HPgJRG5CigHvhHY/U/AeUAp0AJcGYGYB0zHqBwe/c7veW3/DLYtysXX1GR1SCoOpS5cRurCyJ2/1yRgjLm0h4/O7GZfA1wXalDRIqF8L9/7w3+ScACGNa6zOhylIkJ7DB6Dp2wXhfftAsBncSxKRYoOIFIqzmkSUCrOaRJQKs5pElAqzmkSUCrOaRJQKs5pElAqzmkSUCrOaRJQKs5pElAqzmkSUCrOaRJQKs5pElAqzukowm6IMwH7iGHg8+HZVQU+r9UhKRUxWhLohowfQ/IfmjjweweO4TpZshrcNAl0w7gdfDN3BecM2whOLSypwU2iYU17EakDmoF6q2PpQQ4aWzCiNbZojQsiG9tIY8yQIzdGRRIAEJFPjDElVsfRHY0tONEaW7TGBdbEpo8DSsU5TQJKxbloSgLzrQ7gGDS24ERrbNEaF1gQW9TUCSilrBFNJQGllAU0CSgV56IiCYjIOSKyRURKReQOC+MYISIfiMhGEdkgIjcFtmeJyHsisi3wPdPCGO0iskpE3gy8Hy0iywP3bqGIJFgUV4aILBKRzSKySURmRct9E5HvB36e60XkBRFxW3XfRGSBiNSKyPou27q9T4E1PX8TiHGtiMyIREyWJwERsQO/A84FJgGXisgki8LxALcaYyYBM4HrArHcASw1xhQDSwPvrXITsKnL+weAXxljioAG4CpLooJHgHeMMROAqfhjtPy+ichw4EagxBhzHGAHLsG6+/Z/wDlHbOvpPp0LFAe+5gKPRSQiY4ylX8As4N0u7+8E7rQ6rkAsrwNnAVuA/MC2fGCLRfEUBP6TfAl4ExD8vcsc3d3LAYwrHdhJoKK5y3bL7xswHKgAsvAPmHsTONvK+waMAtb3dp+Ax4FLu9svnF+WlwT4/If0mcrANkuJyChgOrAcyDOfL7FeDeRZFNavgdv5fGnEbGC/McYTeG/VvRsN1AFPBx5VnhSRZKLgvhljqoAHgV3AHuAAsJLouG+f6ek+DcjvRjQkgagjIinAK8DNxpjGrp8Zf0oe8HZVETkfqDXGrBzoa/eBA5gBPGaMmY5/HMhhRX8L71smcCH+RDUMSObo4njUsOI+RUMSqAJGdHlfENhmCRFx4k8AzxljFgc214hIfuDzfKDWgtBOAb4qImXAi/gfCR4BMkTks6GOVt27SqDSGLM88H4R/qQQDffty8BOY0ydMaYTWIz/XkbDfftMT/dpQH43oiEJrACKA7W1CfgrbZZYEYiICPAUsMkY83CXj5YAcwKv5+CvKxhQxpg7jTEFxphR+O/RX4wxlwMfAF+zOLZqoEJExgc2nQlsJAruG/7HgJkikhT4+X4Wm+X3rYue7tMS4IpAK8FM4ECXx4bwGeiKmh4qSs4DtgLbgf+yMI5T8RfF1gKrA1/n4X/2XgpsA94Hsiy+X7OBNwOvxwD/AkqBlwGXRTFNAz4J3LvXgMxouW/AfcBmYD3wR8Bl1X0DXsBfN9GJvwR1VU/3CX/F7+8Cvxfr8LdwhD0m7TasVJyLhscBpZSFNAkoFec0CSgV5zQJKBXnNAkoFec0CSgV5zQJKBXn/h/K3j5w6lIVDwAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "H, W = img_size // patch_size, img_size // patch_size\n",
    "print(f\"Sequence length: {H}x{W} = {H * W}\")\n",
    "\n",
    "rand_pattern = torch.rand((H * W) ** 2).reshape(H * W, H * W) > 0.999\n",
    "\n",
    "gaus_2d_dist = AP.local_2d_gausian_distribution(H, W, sigma=16)\n",
    "sparsity = 0.995\n",
    "num_non_zeros = int((H * W) ** 2 * (1 - sparsity))\n",
    "random_gaus_2d_pattern = AP.random_pattern_from_probability_matrix(gaus_2d_dist, num_non_zeros)\n",
    "\n",
    "t_mask = rand_pattern | random_gaus_2d_pattern\n",
    "\n",
    "# and let's not forget to add a global attention for the cls_token\n",
    "mask = torch.ones((H * W + 1, H * W + 1), dtype=torch.bool)\n",
    "mask[1:, 1:] = t_mask\n",
    "\n",
    "print(f\"Sparsity: {1 - mask.float().mean().item()}, nnz={num_non_zeros}\")\n",
    "\n",
    "plt.imshow(mask[H * W // 2 + W // 2][1:].reshape(H, W))\n",
    "plt.show()\n",
    "\n",
    "mask = SparseCS(mask, torch.device(\"cuda\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddf7b6e9-95c3-465e-9a56-c813c30dd943",
   "metadata": {},
   "source": [
    "And now, as before, replace the dense attentions with the sparse ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6cd67fee-1736-4d6b-b63b-0ef087be9fc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_sparse = replace_attn_with_xformers_one(model_sparse, mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67357795-a78e-43fa-a4f9-3a856d993370",
   "metadata": {},
   "source": [
    "Now create a random input and let's feed it to our model and benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6ea14f36-2384-428e-888e-26c7004748a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = torch.rand(8, 3, img_size, img_size).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d9bf472b-227f-441a-864d-8baa9658cda2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward only\n",
      "<torch.utils.benchmark.utils.common.Measurement object at 0x7fbea8033d00>\n",
      "profile\n",
      "  Median: 194.43 ms\n",
      "  IQR:    0.94 ms (193.73 to 194.67)\n",
      "  11 measurements, 1 runs per measurement, 1 thread\n",
      "Memory used: 8022.40283203125 MB\n",
      "\n",
      "Forward + backward\n",
      "<torch.utils.benchmark.utils.common.Measurement object at 0x7fbdf04af4f0>\n",
      "profile\n",
      "  Median: 633.81 ms\n",
      "  IQR:    3.66 ms (632.41 to 636.07)\n",
      "  4 measurements, 1 runs per measurement, 1 thread\n",
      "Memory used: 8207.640625 MB\n"
     ]
    }
   ],
   "source": [
    "print(\"Forward only\")\n",
    "profile_model(lambda : model_sparse(i))\n",
    "print(\"\")\n",
    "print(\"Forward + backward\")\n",
    "profile_model(lambda : model_sparse(i).sum().backward())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3703edc-d2c6-4ebd-946c-219cd8a17153",
   "metadata": {},
   "source": [
    "By using a very sparse matrix for the self-attention computation, we were able to run ViT with sequence length of 10k by using ~8GB of GPU memory for a batch size of 8.\n",
    "The speed / memory can be easily tuned by controlling the degree of sparsity of your `attn_mask`.\n",
    "\n",
    "# Wrapping up\n",
    "\n",
    "In this notebook, we've show one way of quickly getting started with using xformers in your codebase.\n",
    "We showed how to get both memory and speed savings by leveraging sparse computations, and we also showed that it is possible to runa full forward+backward on a ViT model whose sequence length is greater than 10k.\n",
    "\n",
    "We hope that xformers can enable new research directions on very large sequences.\n",
    "\n",
    "If you have questions, we will be happy to help you getting started with xformers!"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}