jan-torch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Numpy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ -4.1467563 , -20.45317729,   8.27463481, ...,  32.17259803,\n",
       "        11.65379888, -22.2881978 ])"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "x = np.random.randn(1024)\n",
    "w = np.random.randn(1024, 1024)\n",
    "b = np.random.randn(1024)\n",
    "\n",
    "w @ x + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def new_multiplication(x, w, b, aggregate=np.sum):\n",
    "    return aggregate((x[np.newaxis, ...] * w), axis=1) + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.isclose(\n",
    "    new_multiplication(x, w, b, aggregate=np.sum),\n",
    "    w @ x + b\n",
    ").all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "71.7 µs ± 1.65 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "w @ x + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.08 ms ± 3.86 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "new_multiplication(x, w, b, aggregate=np.max)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_size = 128\n",
    "out_size = 256\n",
    "x = torch.randn(8, in_size)\n",
    "w = torch.randn(in_size, out_size)\n",
    "b = torch.randn(out_size)\n",
    "\n",
    "def new_multiplication(x, w, b, aggregate=torch.sum):\n",
    "    return aggregate((x.unsqueeze(1) * w.t()), dim=-1) + b\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.0756e+00,  1.3063e-03, -5.8510e-01,  ..., -5.8305e-03,\n",
       "         -3.4652e-01, -1.1912e+00],\n",
       "        [ 1.2265e+00, -9.0259e-02, -3.9968e-01,  ..., -1.5300e-01,\n",
       "         -5.2789e-01, -1.2939e+00],\n",
       "        [ 1.1770e+00,  1.6524e-01, -5.1461e-01,  ..., -6.5718e-02,\n",
       "         -4.9589e-01, -1.3619e+00],\n",
       "        ...,\n",
       "        [ 1.2839e+00,  1.0722e-02, -3.6673e-01,  ..., -7.4722e-02,\n",
       "         -4.0168e-01, -1.3429e+00],\n",
       "        [ 1.1328e+00,  1.0426e-01, -4.5403e-01,  ..., -1.0995e-01,\n",
       "         -6.0258e-01, -1.4791e+00],\n",
       "        [ 1.2728e+00,  7.4930e-02, -3.3275e-01,  ...,  2.4351e-02,\n",
       "         -4.1659e-01, -1.2888e+00]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_multiplication(x, w, b, aggregate=torch.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ -7.7497,  -4.6337, -20.3994,  ...,  11.0727,  14.2768,  16.8557],\n",
       "        [ 11.5624, -16.3541,   3.3344,  ...,  -7.7647,  -8.9386,   3.7175],\n",
       "        [  5.2206,  16.3496, -11.3767,  ...,   3.4070,  -4.8417,  -4.9850],\n",
       "        ...,\n",
       "        [ 18.9017,  -3.4285,   7.5522,  ...,   2.2546,   7.2165,  -2.5537],\n",
       "        [ -0.4373,   8.5438,  -3.6233,  ...,  -2.2544, -18.4990, -19.9883],\n",
       "        [ 17.4838,   4.7901,  11.9011,  ...,  14.9358,   5.3080,   4.3627]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x @ w + b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check it's approximately similar\n",
    "bool(((new_multiplication(x, w, b) - (x @ w + b)) < 0.0001).all())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "24.9 µs ± 1.07 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "(x @ w + b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "61.6 µs ± 2.17 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "new_multiplication(x, w, b, aggregate=torch.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.arange(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])"
      ]
     },
     "execution_count": 135,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha=5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9.001352149889899"
      ]
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.max() + (1/alpha)*np.log(np.sum(np.exp( alpha * (a - a.max()))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def matmul_abs_max_with_sign(x, w, b):\n",
    "    intermediate = x.unsqueeze(1) * w.t()\n",
    "    indxs = intermediate.abs().argmax(dim=-1)\n",
    "    return torch.gather(intermediate, -1, indxs.unsqueeze(-1)).squeeze()\n",
    "\n",
    "def matmul_soft_abs_max_with_sign(x, w, b, alpha=10):\n",
    "    intermediate = x.unsqueeze(1) * w.t()\n",
    "    max_vals, max_idxs = intermediate.abs().max(dim=-1)\n",
    "    sign = torch.gather(intermediate, -1, max_idxs.unsqueeze(-1)).squeeze().sign()\n",
    "    return (max_vals + (1 / alpha) * torch.log(torch.exp(alpha * (intermediate.abs() - max_vals.unsqueeze(-1))).sum(-1))) * sign\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256])"
      ]
     },
     "execution_count": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "intermediate = x.unsqueeze(1) * w.t()\n",
    "x_star = intermediate.max(dim=-1).values\n",
    "x_star.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 3.9366,  4.8260, -3.1061,  ...,  4.2637, -6.3641, -4.9272],\n",
       "        [ 2.8481, -2.8832, -2.9196,  ..., -3.7980,  5.1746,  3.6667],\n",
       "        [-4.6657, -6.2825,  5.2085,  ...,  4.4811,  5.8092, -6.8542],\n",
       "        ...,\n",
       "        [-3.8787, -3.8792, -3.4387,  ..., -4.1604,  4.0529,  3.7508],\n",
       "        [ 3.5586,  2.9782, -3.5428,  ...,  5.6816, -3.7871,  3.9682],\n",
       "        [-3.4155,  4.1684, -4.6598,  ..., -3.4170,  6.9878, -4.3139]])"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matmul_abs_max_with_sign(x, w, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 3.9407,  4.8260, -3.1115,  ...,  4.2717, -6.3758, -4.9297],\n",
       "        [ 2.8502, -2.8873, -2.9216,  ..., -3.8021,  5.1746,  3.6708],\n",
       "        [-4.6657, -6.2825,  5.2085,  ...,  4.4819,  5.8092, -6.8542],\n",
       "        ...,\n",
       "        [-3.8787, -3.8792, -3.4390,  ..., -4.1606,  4.0576,  3.7509],\n",
       "        [ 3.5617,  2.9807, -3.5460,  ...,  5.6816, -3.8231,  3.9682],\n",
       "        [-3.4157,  4.1684, -4.6598,  ..., -3.4340,  6.9878, -4.3141]])"
      ]
     },
     "execution_count": 172,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matmul_soft_abs_max_with_sign(x, w, b, alpha=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NewLinearLayer(torch.nn.Module):\n",
    "    \n",
    "    def __init__(self, in_features, out_features):\n",
    "        super().__init__()\n",
    "        self.register_parameter('weights', torch.nn.Parameter(torch.randn(in_features, out_features)))\n",
    "        self.register_parameter('bias', torch.nn.Parameter(torch.randn(out_features)))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return matmul_soft_abs_max_with_sign(x, self.weights, self.bias)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "import torch.utils.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_linear_model = torch.nn.Sequential(\n",
    "    NewLinearLayer(784, 256),\n",
    "    torch.nn.Hardtanh(),\n",
    "    NewLinearLayer(256, 128),\n",
    "    torch.nn.Hardtanh(),\n",
    "    NewLinearLayer(128, 64),\n",
    "    torch.nn.Hardtanh(),\n",
    "    NewLinearLayer(64, 10),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "metadata": {},
   "outputs": [],
   "source": [
    "linear_model = torch.nn.Sequential(\n",
    "    torch.nn.Linear(784, 256),\n",
    "    torch.nn.Hardtanh(),\n",
    "    torch.nn.Linear(256, 128),\n",
    "    torch.nn.Hardtanh(),\n",
    "    torch.nn.Linear(128, 64),\n",
    "    torch.nn.Hardtanh(),\n",
    "    torch.nn.Linear(64, 10),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = torchvision.datasets.MNIST(\"/Users/jan/datasets/mnist/\", transform=torchvision.transforms.ToTensor())\n",
    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=16)\n",
    "x, y = next(iter(dataloader))\n",
    "x = x.reshape(x.size(0), -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = torch.nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16dfebeb2f0b4d3e8f0ac3321178f18a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=3750), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.3180199432373048\n",
      "2.2612040305137633\n",
      "2.222285380363464\n",
      "2.1685437893867503\n",
      "2.12642452955246\n",
      "2.0391437649726867\n",
      "1.899020750522613\n",
      "1.7774287736415864\n",
      "1.6100853729248048\n",
      "1.4525646042823788\n",
      "1.2753579592704778\n",
      "1.1391883230209354\n",
      "0.9826255416870119\n",
      "0.9034445458650587\n",
      "0.8632834327220914\n",
      "0.7839462408423423\n",
      "0.7535398200154305\n",
      "0.702422222197056\n",
      "0.7132131338119506\n",
      "0.7136265051364897\n",
      "0.6865319877862932\n",
      "0.5714167705178261\n",
      "0.5325418104231356\n",
      "0.6202109879255291\n",
      "0.5728331223130224\n",
      "0.5429678289592267\n",
      "0.5726944382488729\n",
      "0.48500355511903764\n",
      "0.5558524391055109\n",
      "0.4953328366577627\n",
      "0.49067893758416187\n",
      "0.5232392819225788\n",
      "0.4796754842996597\n",
      "0.4696946971118451\n",
      "0.42503870829939844\n",
      "0.402633321136236\n",
      "0.3190304088592529\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3036b091a9fb4c03a66ad7beb3136877",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=3750), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.98548649072647\n",
      "4.126735692024231\n",
      "4.012626132965087\n",
      "3.960223226547241\n",
      "3.984845917224885\n",
      "3.9781308436393745\n",
      "3.8746247375011427\n",
      "3.881452393531798\n",
      "3.7416932892799366\n",
      "3.8226396346092217\n",
      "3.7050264263153077\n",
      "3.772139885425567\n",
      "3.7775453400611876\n",
      "3.7736988162994383\n",
      "3.6761367630958564\n",
      "3.746480813026428\n",
      "3.813709783554078\n",
      "3.750118653774261\n",
      "3.637923209667205\n",
      "3.604153666496277\n",
      "3.629136486053468\n"
     ]
    }
   ],
   "source": [
    "\n",
    "for model in (linear_model, new_linear_model):\n",
    "    optim = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.8)\n",
    "    avg = 0\n",
    "    for i, (x, y) in enumerate(tqdm.tqdm_notebook(dataloader)):\n",
    "        x = x.reshape(x.size(0), -1)\n",
    "        predictions = model(x)\n",
    "        loss = criterion(predictions, y)\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "        optim.zero_grad()\n",
    "        avg += loss.item() / 100\n",
    "        if i % 100 == 0 and i > 0:\n",
    "            print(avg)\n",
    "            avg = 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "vals, indxs = w.max(dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "intermediate = x.unsqueeze(1) * w.t()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "vals, indxs = intermediate.abs().max(dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256])"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indxs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[3.2035, 5.4254, 5.3714,  ..., 5.3970, 5.2756, 4.8329],\n",
       "        [4.7636, 4.2375, 4.4597,  ..., 2.7539, 4.1946, 3.9239],\n",
       "        [4.5746, 3.1270, 6.6946,  ..., 4.5977, 6.5752, 3.0569],\n",
       "        ...,\n",
       "        [3.7134, 5.4075, 6.0133,  ..., 2.7904, 5.9060, 4.0532],\n",
       "        [5.4211, 5.0082, 5.0784,  ..., 4.6823, 6.4084, 3.8672],\n",
       "        [4.9760, 6.0081, 4.3977,  ..., 2.9710, 3.4506, 3.8452]])"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 3.2035, -5.4254, -5.3714,  ...,  5.3970,  5.2756, -4.8329],\n",
       "        [-4.7636, -4.2375, -4.4597,  ...,  2.7539, -4.1946, -3.9239],\n",
       "        [ 4.5746,  3.1270, -6.6946,  ..., -4.5977,  6.5752, -3.0569],\n",
       "        ...,\n",
       "        [ 3.7134,  5.4075,  6.0133,  ..., -2.7904, -5.9060, -4.0532],\n",
       "        [ 5.4211, -5.0082, -5.0784,  ..., -4.6823,  6.4084,  3.8672],\n",
       "        [-4.9760, -6.0081, -4.3977,  ..., -2.9710, -3.4506,  3.8452]])"
      ]
     },
     "execution_count": 107,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.gather(intermediate, -1, indxs.unsqueeze(-1)).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}