notebook_test_gram_iteration.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e388dc74",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from bounds import *\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "torch.manual_seed(42)\n",
    "torch.cuda.manual_seed_all(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49722cb7",
   "metadata": {},
   "source": [
    "# Spectral norm estimation on dense matrix\n",
    "\n",
    "This code is related to dense spectral norm estimation, see Section 5.5 in paper. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "81928fab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dense sigma_1_reference 41.99017911516535\n"
     ]
    }
   ],
   "source": [
    "n, m = 500, 400\n",
    "with_cuda = False\n",
    "G = torch.randn(n, m).double().cuda()\n",
    "sigma_1_reference = torch.linalg.matrix_norm(G, ord=2).item()\n",
    "print(\"dense sigma_1_reference\", sigma_1_reference)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af1afefa",
   "metadata": {},
   "source": [
    "# Test Power iteration on dense matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a5b1a4a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Diff Power iteration -1.247428826900432e-10 Mean time 0.05896735191345215\n"
     ]
    }
   ],
   "source": [
    "n_iter_pi = 1000\n",
    "nb_reps = 10\n",
    "burn = 10\n",
    "sigmas_1_pi, times = [], []\n",
    "for _ in range(nb_reps + burn):\n",
    "    sigma_1_pi, time = estimate_dense(G, \n",
    "                                      n_iter=n_iter_pi, \n",
    "                                      name_func=\"pi\", \n",
    "                                      return_time=True)\n",
    "    sigmas_1_pi.append(sigma_1_pi.item())\n",
    "    times.append(time)\n",
    "print(\"Diff Power iteration\", np.mean(sigmas_1_pi[burn:]) - sigma_1_reference, \n",
    "      \"Mean time\", np.mean(times[burn:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9cb2f61",
   "metadata": {},
   "source": [
    "# Test Gram iteration on dense matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e9df7299",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Diff Gram iteration -3.574029960873304e-12 Mean time 0.0009333372116088867\n"
     ]
    }
   ],
   "source": [
    "n_iter_gram = 15\n",
    "nb_reps = 10\n",
    "burn = 10\n",
    "sigmas_1_gram, times = [], []\n",
    "for _ in range(nb_reps + burn):\n",
    "    sigma_1_gram, time = estimate_dense(G, \n",
    "                                        n_iter=n_iter_gram, \n",
    "                                        name_func=\"gi\", \n",
    "                                        return_time=True)\n",
    "    sigmas_1_gram.append(sigma_1_gram.item())\n",
    "    times.append(time)\n",
    "print(\"Diff Gram iteration\", np.mean(sigmas_1_gram[burn:]) - sigma_1_reference, \n",
    "      \"Mean time\", np.mean(times[burn:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a75f3ff8",
   "metadata": {},
   "source": [
    "# Spectral norm estimation on convolutional layer\n",
    "\n",
    "This code is related to convolutional layer spectral norm estimation, see Section 5.2 in paper.# Define convolution kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1af4f4db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kernel sigma_1_reference_sedghi2019 83.09203338623047\n",
      "kernel sigma_1_reference_ryu2019 77.56085968017578\n"
     ]
    }
   ],
   "source": [
    "cout = 64\n",
    "cin = 64\n",
    "input_size_n = 12\n",
    "kernel_size = 5\n",
    "\n",
    "\n",
    "kernel = torch.randn(cout, cin, kernel_size, kernel_size).cuda()\n",
    "\n",
    "sigma_1_reference_sedghi2019 = estimate(kernel, \n",
    "                                        n=input_size_n, \n",
    "                                        name_func=\"sedghi2019\").item()\n",
    "print(\"kernel sigma_1_reference_sedghi2019\", sigma_1_reference_sedghi2019)\n",
    "\n",
    "n_iter_ryu2019_ref = 100\n",
    "sigma_1_reference_ryu2019 = estimate(kernel, \n",
    "                                     n=input_size_n, \n",
    "                                     n_iter=n_iter_ryu2019_ref, \n",
    "                                     name_func=\"ryu2019\").item()\n",
    "print(\"kernel sigma_1_reference_ryu2019\", sigma_1_reference_ryu2019)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8f805ee",
   "metadata": {},
   "source": [
    "# Test Araujo2021 on convolutional layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f8e74be9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Diff Sedghi2019 with Araujo2021 conv  273.63285064697266 \n",
      "Diff Ryu2019 with Araujo2021 conv 279.16402435302734 \n",
      "Mean time 0.0010010957717895507\n"
     ]
    }
   ],
   "source": [
    "nb_samples_araujo2021 = 50\n",
    "nb_reps = 10\n",
    "burn = 10\n",
    "sigmas_1_araujo2021, times = [], []\n",
    "for _ in range(nb_reps + burn):\n",
    "    sigma_1_araujo2021, time = estimate(kernel, \n",
    "                                  n=input_size_n, \n",
    "                                  n_iter=nb_samples_araujo2021, \n",
    "                                  name_func=\"araujo2021\", \n",
    "                                  return_time=True)\n",
    "    sigmas_1_araujo2021.append(sigma_1_araujo2021.item())\n",
    "    times.append(time)\n",
    "print(\"Diff Sedghi2019 with Araujo2021 conv \", np.mean(sigmas_1_araujo2021[burn:]) - sigma_1_reference_sedghi2019,\n",
    "      \"\\nDiff Ryu2019 with Araujo2021 conv\", np.mean(sigmas_1_araujo2021[burn:]) - sigma_1_reference_ryu2019, \n",
    "      \"\\nMean time\", np.mean(times[burn:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80f261bd",
   "metadata": {},
   "source": [
    "# Test Singla2021 on convolutional layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bced0cae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Diff Sedghi2019 with Singla2021conv 93.96885681152344 \n",
      "Diff Ryu2019 with Singla2021 conv 99.50003051757812 \n",
      "Mean time 0.01792454719543457\n"
     ]
    }
   ],
   "source": [
    "n_iter_singla2021 = 50\n",
    "nb_reps = 10\n",
    "burn = 10\n",
    "sigmas_1_singla2021, times = [], []\n",
    "for _ in range(nb_reps + burn):\n",
    "    sigma_1_singla2021, time = estimate(kernel, \n",
    "                                  n=input_size_n, \n",
    "                                  n_iter=n_iter_singla2021, \n",
    "                                  name_func=\"singla2021\", \n",
    "                                  return_time=True)\n",
    "    sigmas_1_singla2021.append(sigma_1_singla2021.item())\n",
    "    times.append(time)\n",
    "print(\"Diff Sedghi2019 with Singla2021conv\", np.mean(sigmas_1_singla2021[burn:]) - sigma_1_reference_sedghi2019,\n",
    "      \"\\nDiff Ryu2019 with Singla2021 conv\", np.mean(sigmas_1_singla2021[burn:]) - sigma_1_reference_ryu2019, \n",
    "      \"\\nMean time\", np.mean(times[burn:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb996ff5",
   "metadata": {},
   "source": [
    "# Test Gram iteration on convolution layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5b78d520",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Diff Sedghi2019 with Gram iteration conv 0.025726318359375 \n",
      "Diff Ryu2019 with Gram iteration conv 3.7672805786132812 \n",
      "Mean time 0.000528717041015625\n"
     ]
    }
   ],
   "source": [
    "n_iter_gram = 5\n",
    "nb_reps = 10\n",
    "burn = 10\n",
    "sigmas_1_gram, times = [], []\n",
    "for _ in range(nb_reps + burn):\n",
    "    sigma_1_gram, time = estimate(kernel, \n",
    "                                  n=input_size_n, \n",
    "                                  n_iter=n_iter_gram, \n",
    "                                  name_func=\"delattre2023\", \n",
    "                                  return_time=True)\n",
    "    sigmas_1_gram.append(sigma_1_gram.item())\n",
    "    times.append(time)\n",
    "print(\"Diff Sedghi2019 with Gram iteration conv\", np.mean(sigmas_1_gram[burn:]) - sigma_1_reference_sedghi2019,\n",
    "      \"\\nDiff Ryu2019 with Gram iteration conv\", np.mean(sigmas_1_gram[burn:]) - sigma_1_reference_ryu2019, \n",
    "      \"\\nMean time\", np.mean(times[burn:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b08c21bc",
   "metadata": {},
   "source": [
    "# Test norm2toep (Gram iteration on convolution layer with zero padding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "00c75cc0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Diff Sedghi2019 with norm2toep 3.3376846313476562 \n",
      "Diff Ryu2019 with norm2toep 8.868858337402344 \n",
      "Mean time 0.0007239341735839844\n"
     ]
    }
   ],
   "source": [
    "n_iter_gram = 5\n",
    "nb_reps = 10\n",
    "burn = 10\n",
    "sigmas_1_gram, times = [], []\n",
    "for _ in range(nb_reps + burn):\n",
    "    sigma_1_gram, time = estimate(kernel, \n",
    "                                  n=input_size_n, \n",
    "                                  n_iter=n_iter_gram, \n",
    "                                  name_func=\"delattre2024\", \n",
    "                                  return_time=True)\n",
    "    sigmas_1_gram.append(sigma_1_gram.item())\n",
    "    times.append(time)\n",
    "print(\"Diff Sedghi2019 with norm2toep\", np.mean(sigmas_1_gram[burn:]) - sigma_1_reference_sedghi2019,\n",
    "      \"\\nDiff Ryu2019 with norm2toep\", np.mean(sigmas_1_gram[burn:]) - sigma_1_reference_ryu2019, \n",
    "      \"\\nMean time\", np.mean(times[burn:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b83bc0ab",
   "metadata": {},
   "source": [
    "# Compute spectral norm of convolutional layers of ResNet18\n",
    "\n",
    "This code is related to Section 5.3 in paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc9c2252",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import models\n",
    "\n",
    "n_iter_name = {\n",
    "               \"delattre2023\":7,\n",
    "               \"singla2021\":50,\n",
    "               \"ryu2019\" : 100,\n",
    "               \"araujo2021\" : 20,\n",
    "               \"sedghi2019\":None}\n",
    "func_names = [\n",
    "              \"ryu2019\",\n",
    "              #\"sedghi2019\", # commented because it takes a while\n",
    "              \"araujo2021\",\n",
    "              \"singla2021\",\n",
    "              \"delattre2023\",\n",
    "]\n",
    "model_resnet_18 = models.resnet18(pretrained=True).cuda().eval()\n",
    "\n",
    "\n",
    "\n",
    "inp_shape = (224, 224)\n",
    "lip_tot = {name : 1 for name in func_names}\n",
    "times_tot = {name : 0 for name in func_names}\n",
    "\n",
    "with torch.no_grad():\n",
    "    for name, module in model_resnet_18.named_modules():\n",
    "        print()\n",
    "        is_downsample = name.endswith(\"downsample\")\n",
    "        is_regular_conv = \"conv\" in name\n",
    "        is_max_pool = \"MaxPool2d\" in module.__class__.__name__\n",
    "        if is_max_pool:\n",
    "            stride = module.stride\n",
    "            inp_shape = (inp_shape[0] // stride, inp_shape[1] // stride)\n",
    "        if is_downsample:\n",
    "            # dowsampling layer in residual connection\n",
    "            conv, bn = module[0], module[1]\n",
    "            lip_bn = (bn.weight.detach() / bn.running_var).max().item()\n",
    "            inp_shape = (inp_shape[0] * stride, inp_shape[1] * stride)\n",
    "        elif is_regular_conv:\n",
    "            conv = module\n",
    "            stride = conv.stride[0]\n",
    "            lip_bn = 1.0\n",
    "        if is_downsample or is_regular_conv:\n",
    "            param = conv.weight.clone().detach()\n",
    "            out_channels, in_channels, H, W = param.shape\n",
    "            for name in func_names:\n",
    "                bound, curr_time = estimate(param,\n",
    "                                            inp_shape[0],\n",
    "                                            n_iter_name[name],\n",
    "                                            name,\n",
    "                                            return_time=True)\n",
    "                bound = bound.item()\n",
    "                print(name, \"conv weight dim\", param.shape, \"n\", inp_shape[0], \"bound\", bound, \"time\", curr_time)\n",
    "                if is_downsample:\n",
    "                    lip_tot[name] += bound* lip_bn\n",
    "                else:\n",
    "                    # Lipschiz of batch norm cancels in ratio\n",
    "                    lip_tot[name] *= bound\n",
    "                times_tot[name] += curr_time\n",
    "                print(\"Total Lipschitz bound\", lip_tot[name], \"total time\", times_tot[name], \"\\n\")\n",
    "            inp_shape = (inp_shape[0] // stride, inp_shape[1] // stride)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abd15c66",
   "metadata": {},
   "source": [
    "# Total Lipschitz ratio bound "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26b99c0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "lip_tot_ref = lip_tot[\"ryu2019\"]\n",
    "for name in func_names:\n",
    "    lip_tot[name] /= lip_tot_ref\n",
    "print(\"Total Lipschitz ratio\", lip_tot)\n",
    "print(\"Total times\", times_tot)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "ad2bdc8ecc057115af97d19610ffacc2b4e99fae6737bb82f5d7fb13d2f2c186"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}