Answer is: (10, 100, 200, 32) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = rearrange(x, \"b c h w -> b h w c\")\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Worked!\n", "\n", "Did you notice? Code above worked for you backend of choice.\n", "Einops functions work with any tensor like they are native to the framework.\n", "\n", "## Backpropagation\n", "\n", "- gradients are a corner stone of deep learning\n", "- You can back-propagate through einops operations
\n", " (just as with framework native operations)" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "tensor(320., dtype=torch.float64)\n" ] } ], "source": [ "y0 = x\n", "y1 = reduce(y0, \"b c h w -> b c\", \"max\")\n", "y2 = rearrange(y1, \"b c -> c b\")\n", "y3 = reduce(y2, \"c b -> \", \"sum\")\n", "\n", "if flavour == \"tensorflow\":\n", " print(reduce(tape.gradient(y3, x), \"b c h w -> \", \"sum\"))\n", "else:\n", " y3.backward()\n", " print(reduce(x.grad, \"b c h w -> \", \"sum\"))" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## Meet `einops.asnumpy`\n", "\n", "Just converts tensors to numpy (and pulls from gpu if necessary)" ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "
Answer is: (10, 640000) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = rearrange(x, \"b c h w -> b (c h w)\")\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "**space-to-depth**" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 128, 50, 100) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = rearrange(x, \"b c (h h1) (w w1) -> b (h1 w1 c) h w\", h1=2, w1=2)\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**depth-to-space** (notice that it's reverse of the previous)" ] }, { "cell_type": "code", "execution_count": 12, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 8, 200, 400) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = rearrange(x, \"b (h1 w1 c) h w -> b c (h h1) (w w1)\", h1=2, w1=2)\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Reductions\n", "\n", "Simple **global average pooling**." ] }, { "cell_type": "code", "execution_count": 13, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 32) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = reduce(x, \"b c h w -> b c\", reduction=\"mean\")\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "**max-pooling** with a kernel 2x2" ] }, { "cell_type": "code", "execution_count": 14, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 32, 50, 100) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = reduce(x, \"b c (h h1) (w w1) -> b c h w\", reduction=\"max\", h1=2, w1=2)\n", "guess(y.shape)" ] }, { "cell_type": "code", "execution_count": 15, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 32, 50, 100) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# you can skip names for reduced axes\n", "y = reduce(x, \"b c (h 2) (w 2) -> b c h w\", reduction=\"max\")\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1d, 2d and 3d pooling are defined in a similar way\n", "\n", "for sequential 1-d models, you'll probably want pooling over time\n", "```python\n", "reduce(x, '(t 2) b c -> t b c', reduction='max')\n", "```\n", "\n", "for volumetric models, all three dimensions are pooled\n", "```python\n", "reduce(x, 'b c (x 2) (y 2) (z 2) -> b c x y z', reduction='max')\n", "```\n", "\n", "Uniformity is a strong point of `einops`, and you don't need specific operation for each particular case.\n", "\n", "\n", "### Good exercises \n", "\n", "- write a version of space-to-depth for 1d and 3d (2d is provided above)\n", "- write an average / max pooling for 1d models. " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Squeeze and unsqueeze (expand_dims)" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [], "source": [ "# models typically work only with batches,\n", "# so to predict a single image ...\n", "image = rearrange(x[0, :3], \"c h w -> h w c\")\n", "# ... create a dummy 1-element axis ...\n", "y = rearrange(image, \"h w c -> () c h w\")\n", "# ... imagine you predicted this with a convolutional network for classification,\n", "# we'll just flatten axes ...\n", "predictions = rearrange(y, \"b c h w -> b (c h w)\")\n", "# ... finally, decompose (remove) dummy axis\n", "predictions = rearrange(predictions, \"() classes -> classes\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## keepdims-like behavior for reductions\n", "\n", "- empty composition `()` provides dimensions of length 1, which are broadcastable.\n", "- alternatively, you can use just `1` to introduce new axis, that's a synonym to `()`\n", "\n", "per-channel mean-normalization for *each image*:" ] }, { "cell_type": "code", "execution_count": 17, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 32, 100, 200) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = x - reduce(x, \"b c h w -> b c 1 1\", \"mean\")\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "per-channel mean-normalization for *whole batch*:" ] }, { "cell_type": "code", "execution_count": 18, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 32, 100, 200) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = x - reduce(y, \"b c h w -> 1 c 1 1\", \"mean\")\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Stacking\n", "\n", "let's take a list of tensors" ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [], "source": [ "list_of_tensors = list(x)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "New axis (one that enumerates tensors) appears *first* on the left side of expression.\n", "Just as if you were indexing list - first you'd get tensor by index" ] }, { "cell_type": "code", "execution_count": 20, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 100, 200, 32) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "tensors = rearrange(list_of_tensors, \"b c h w -> b h w c\")\n", "guess(tensors.shape)" ] }, { "cell_type": "code", "execution_count": 21, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (100, 200, 32, 10) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# or maybe stack along last dimension?\n", "tensors = rearrange(list_of_tensors, \"b c h w -> h w c b\")\n", "guess(tensors.shape)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Concatenation\n", "\n", "concatenate over the first dimension?" ] }, { "cell_type": "code", "execution_count": 22, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (1000, 200, 32) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "tensors = rearrange(list_of_tensors, \"b c h w -> (b h) w c\")\n", "guess(tensors.shape)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "or maybe concatenate along last dimension?" ] }, { "cell_type": "code", "execution_count": 23, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (100, 200, 320) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "tensors = rearrange(list_of_tensors, \"b c h w -> h w (b c)\")\n", "guess(tensors.shape)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Shuffling within a dimension\n", "\n", "**channel shuffle** (as it is drawn in shufflenet paper)" ] }, { "cell_type": "code", "execution_count": 24, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 32, 100, 200) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = rearrange(x, \"b (g1 g2 c) h w-> b (g2 g1 c) h w\", g1=4, g2=4)\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "simpler version of **channel shuffle**" ] }, { "cell_type": "code", "execution_count": 25, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [ { "data": { "text/html": [ "\n", "Answer is: (10, 32, 100, 200) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y = rearrange(x, \"b (g c) h w-> b (c g) h w\", g=4)\n", "guess(y.shape)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Split a dimension\n", "\n", "Here's a super-convenient trick.\n", "\n", "Example: when a network predicts several bboxes for each position\n", "\n", "Assume we got 8 bboxes, 4 coordinates each.\n", "To get coordinated into 4 separate variables, you move corresponding dimension to front and unpack tuple." ] }, { "cell_type": "code", "execution_count": 26, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", "
Answer is: (10, 8, 100, 200) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "\n", "Answer is: (10, 100, 200) (hover to see)
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "bbox_x, bbox_y, bbox_w, bbox_h = rearrange(x, \"b (coord bbox) h w -> coord b bbox h w\", coord=4, bbox=8)\n", "# now you can operate on individual variables\n", "max_bbox_area = reduce(bbox_w * bbox_h, \"b bbox h w -> b h w\", \"max\")\n", "guess(bbox_x.shape)\n", "guess(max_bbox_area.shape)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## Getting into the weeds of tensor packing\n", "\n", "you can skip this part - it explains why taking a habit of defining splits and packs explicitly\n", "\n", "when implementing custom gated activation (like GLU), split is needed:\n", "```python\n", "y1, y2 = rearrange(x, 'b (split c) h w -> split b c h w', split=2)\n", "result = y2 * sigmoid(y2) # or tanh\n", "```\n", "... but we could split differently\n", "```python\n", "y1, y2 = rearrange(x, 'b (c split) h w -> split b c h w', split=2)\n", "```\n", "\n", "- first one splits channels into consequent groups: `y1 = x[:, :x.shape[1] // 2, :, :]`\n", "- while second takes channels with a step: `y1 = x[:, 0::2, :, :]`\n", "\n", "This may drive to very *surprising* results when input is\n", "- a result of group convolution\n", "- a result of bidirectional LSTM/RNN\n", "- multi-head attention\n", "\n", "Let's focus on the second case (LSTM/RNN), since it is less obvious.\n", "\n", "For instance, cudnn concatenates LSTM outputs for forward-in-time and backward-in-time\n", "\n", "Also in pytorch GLU splits channels into consequent groups (first way)\n", "So when LSTM's output comes to GLU,\n", "- forward-in-time produces linear part, and backward-in-time produces activation ...\n", "- and role of directions is different, and gradients coming to two parts are different\n", " - that's not what you expect from simple `GLU(BLSTM(x))`, right?\n", "\n", "`einops` notation makes such inconsistencies explicit and easy-detectable" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Shape parsing\n", "just a handy utility" ] }, { "cell_type": "code", "execution_count": 27, "metadata": { "jupyter": { "outputs_hidden": false }, "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "from einops import parse_shape" ] }, { "cell_type": "code", "execution_count": 28, "metadata": {}, "outputs": [], "source": [ "def convolve_2d(x):\n", " # imagine we have a simple 2d convolution with padding,\n", " # so output has same shape as input.\n", " # Sorry for laziness, use imagination!\n", " return x" ] }, { "cell_type": "code", "execution_count": 29, "metadata": {}, "outputs": [], "source": [ "# imagine we are working with 3d data\n", "x_5d = rearrange(x, \"b c x (y z) -> b c x y z\", z=20)\n", "# but we have only 2d convolutions.\n", "# That's not a problem, since we can apply\n", "y = rearrange(x_5d, \"b c x y z -> (b z) c x y\")\n", "y = convolve_2d(y)\n", "# not just specifies additional information, but verifies that all dimensions match\n", "y = rearrange(y, \"(b z) c x y -> b c x y z\", **parse_shape(x_5d, \"b c x y z\"))" ] }, { "cell_type": "code", "execution_count": 30, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{'b': 10, 'c': 32, 'x': 100, 'y': 10, 'z': 20}" ] }, "execution_count": 30, "metadata": {}, "output_type": "execute_result" } ], "source": [ "parse_shape(x_5d, \"b c x y z\")" ] }, { "cell_type": "code", "execution_count": 31, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{'batch': 10, 'c': 32}" ] }, "execution_count": 31, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# we can skip some dimensions by writing underscore\n", "parse_shape(x_5d, \"batch c _ _ _\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Striding anything\n", "\n", "Finally, how to convert any operation into a strided operation?\n", "(like convolution with strides, aka dilated/atrous convolution)" ] }, { "cell_type": "code", "execution_count": 32, "metadata": {}, "outputs": [], "source": [ "# each image is split into subgrids, each subgrid now is a separate \"image\"\n", "y = rearrange(x, \"b c (h hs) (w ws) -> (hs ws b) c h w\", hs=2, ws=2)\n", "y = convolve_2d(y)\n", "# pack subgrids back to an image\n", "y = rearrange(y, \"(hs ws b) c h w -> b c (h hs) (w ws)\", hs=2, ws=2)\n", "\n", "assert y.shape == x.shape" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Layers\n", "\n", "For frameworks that prefer operating with layers, layers are available.\n", "\n", "You'll need to import a proper one depending on your backend:\n", "\n", "```python\n", "from einops.layers.torch import Rearrange, Reduce\n", "from einops.layers.flax import Rearrange, Reduce\n", "from einops.layers.tensorflow import Rearrange, Reduce\n", "from einops.layers.chainer import Rearrange, Reduce\n", "```\n", "\n", "`Einops` layers are identical to operations, and have same parameters.
\n", "(for the exception of first argument, which should be passed during call)\n", "\n", "```python\n", "layer = Rearrange(pattern, **axes_lengths)\n", "layer = Reduce(pattern, reduction, **axes_lengths)\n", "\n", "# apply layer to tensor\n", "x = layer(x)\n", "```\n", "\n", "Usually it is more convenient to use layers, not operations, to build models\n", "```python\n", "# example given for pytorch, but code in other frameworks is almost identical\n", "from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU\n", "from einops.layers.torch import Reduce\n", "\n", "model = Sequential(\n", " Conv2d(3, 6, kernel_size=5),\n", " MaxPool2d(kernel_size=2),\n", " Conv2d(6, 16, kernel_size=5),\n", " # combined pooling and flattening in a single step\n", " Reduce('b c (h 2) (w 2) -> b (c h w)', 'max'), \n", " Linear(16*5*5, 120), \n", " ReLU(),\n", " Linear(120, 10), \n", " # In flax, the {'axis': value} syntax for specifying values for axes is mandatory:\n", " # Rearrange('(b1 b2) d -> b1 b2 d', {'b1': 12}), \n", ")\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## What's now?\n", "\n", "- rush through [writing better code with einops+pytorch](https://arogozhnikov.github.io/einops/pytorch-examples.html)\n", "\n", "Use different framework? Not a big issue, most recommendations transfer well to other frameworks.
\n", "`einops` works the same way in any framework.\n", "\n", "Finally - just write your code with einops!" ] } ], "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.6" } }, "nbformat": 4, "nbformat_minor": 4 } arogozhnikov-einops-ad2c8d6/docs/3-einmix-layer.ipynb 0000664 0000000 0000000 00000065502 14752016746 0023004 0 ustar 00root root 0000000 0000000 { "cells": [ { "cell_type": "markdown", "id": "0c6e8fa3-2900-4879-a610-3c038f8c8d23", "metadata": {}, "source": [ "# EinMix: universal toolkit for advanced MLP architectures\n", "\n", "Recent progress in MLP-based architectures demonstrated that *very specific* MLPs can compete with convnets and transformers (and even outperform them).\n", "\n", "EinMix allows writing such architectures in a more uniform and readable way." ] }, { "cell_type": "markdown", "id": "dc9a3913-5d5f-41da-8946-f17b3e537ec4", "metadata": {}, "source": [ "## EinMix — building block of MLPs" ] }, { "cell_type": "code", "execution_count": 1, "id": "5f24a8f2-6680-4298-9736-081719c53f88", "metadata": {}, "outputs": [], "source": [ "from einops.layers.torch import EinMix as Mix" ] }, { "cell_type": "markdown", "id": "eacd97e9-7ae2-4321-a3ad-3fb9b122fd81", "metadata": {}, "source": [ "Logic of EinMix is very close to the one of `einsum`. \n", "If you're not familiar with einsum, follow these guides first:\n", "\n", "- https://rockt.github.io/2018/04/30/einsum\n", "- https://towardsdatascience.com/einsum-an-underestimated-function-99ca96e2942e\n", "- https://theaisummer.com/einsum-attention/\n", "\n", "Einsum uniformly describes a number of operations, however `EinMix` is defined slightly differently.\n", "\n", "Here is a linear layer, a common block in sequence modelling (e.g. in NLP/speech), written with einsum\n", "```python\n", "weight = <...create tensor...>\n", "result = torch.einsum('tbc,cd->tbd', embeddings, weight)\n", "```\n", "\n", "EinMix counter-part is:\n", "```python\n", "mix_channels = Mix('t b c -> t b c_out', weight_shape='c c_out', ...)\n", "result = mix_channels(embeddings)\n", "```\n", "\n", "Main differences compared to plain `einsum` are:\n", "\n", "- layer takes care of the weight initialization & management hassle\n", "- weight is not in the comprehension\n", "\n", "We'll discuss other changes a bit later, now let's implement ResMLP" ] }, { "cell_type": "code", "execution_count": 2, "id": "d6983748-20d9-4825-a79f-a0641e58085c", "metadata": {}, "outputs": [], "source": [ "# let's start\n", "import torch\n", "from torch import nn" ] }, { "cell_type": "markdown", "id": "928fc321-7b41-4737-bbf4-644021c1c209", "metadata": {}, "source": [ "## ResMLP — original implementation\n", "\n", "Building blocks of ResMLP consist only of linear/affine layers and one activation (GELU).
\n", "Let's see how we can rewrite all of the components with Mix. \n", "\n", "We start from a reference code for ResMLP block published in the [paper](https://arxiv.org/pdf/2105.03404.pdf):" ] }, { "cell_type": "code", "execution_count": 3, "id": "eed8edf4-0024-482a-bba2-03e84d4086ab", "metadata": {}, "outputs": [], "source": [ "# No norm layer\n", "class Affine(nn.Module):\n", " def __init__(self, dim):\n", " super().__init__()\n", " self.alpha = nn.Parameter(torch.ones(dim))\n", " self.beta = nn.Parameter(torch.zeros(dim))\n", "\n", " def forward(self, x):\n", " return self.alpha * x + self.beta\n", "\n", "\n", "class Mlp(nn.Module):\n", " def __init__(self, dim):\n", " super().__init__()\n", " self.fc1 = nn.Linear(dim, 4 * dim)\n", " self.act = nn.GELU()\n", " self.fc2 = nn.Linear(4 * dim, dim)\n", "\n", " def forward(self, x):\n", " x = self.fc1(x)\n", " x = self.act(x)\n", " x = self.fc2(x)\n", " return x\n", "\n", "class ResMLP_Blocks(nn.Module):\n", " def __init__(self, nb_patches, dim, layerscale_init):\n", " super().__init__()\n", " self.affine_1 = Affine(dim)\n", " self.affine_2 = Affine(dim)\n", " self.linear_patches = nn.Linear(nb_patches, nb_patches) #Linear layer on patches\n", " self.mlp_channels = Mlp(dim) #MLP on channels\n", " self.layerscale_1 = nn.Parameter(layerscale_init * torch.ones((dim))) # LayerScale\n", " self.layerscale_2 = nn.Parameter(layerscale_init * torch.ones((dim))) # parameters\n", "\n", " def forward(self, x):\n", " res_1 = self.linear_patches(self.affine_1(x).transpose(1,2)).transpose(1,2)\n", " x = x + self.layerscale_1 * res_1\n", " res_2 = self.mlp_channels(self.affine_2(x))\n", " x = x + self.layerscale_2 * res_2\n", " return x" ] }, { "cell_type": "markdown", "id": "7496cf7b-a3e1-4648-b62b-30de26c2e718", "metadata": {}, "source": [ "## ResMLP — rewritten\n", "\n", "Code below is the result of first rewriting: \n", "- combination [transpose -> linear -> transpose back] got nicely packed into a single `EinMix` (`mix_patches`)
\n", " `Mix('b t c -> b t0 c', weight_shape='t t0', bias_shape='t0', t=nb_patches, t0=nb_patches)` \n", " - pattern `'b t c -> b t0 c'` tells that `b` and `c` are unperturbed, while tokens `t->t0` were mixed\n", " - explicit parameter shapes are also quite insightful\n", " \n", "- In new implementation affine layer is also handled by `EinMix`:
\n", " `Mix('b t c -> b t c', weight_shape='c', bias_shape='c', c=dim)`\n", " - from the pattern you can see that there is no mixing at all, only multiplication and shift\n", " - multiplication and shift are defined by weight and bias - and those depend only on a channel\n", " - thus affine transform is per-channel\n", " \n", "- Linear layer is also handled by EinMix, the only difference compared to affine layer is absence of bias\n", "- We specified that input is 3d and order is `btc`, not `tbc` - this is not written explicitly in the original code\n", "\n", "The only step back that we had to do is change an initialization schema for EinMix for affine and linear layers" ] }, { "cell_type": "code", "execution_count": 4, "id": "a9cc8c91-d80a-4c35-a010-dbafad22b8c7", "metadata": {}, "outputs": [], "source": [ "def Mlp(dim):\n", " return nn.Sequential(\n", " nn.Linear(dim, 4 * dim),\n", " nn.GELU(),\n", " nn.Linear(4 * dim, dim),\n", " )\n", "\n", "def init(Mix_layer, scale=1.):\n", " Mix_layer.weight.data[:] = scale\n", " if Mix_layer.bias is not None:\n", " Mix_layer.bias.data[:] = 0\n", " return Mix_layer\n", "\n", "class ResMLP_Blocks2(nn.Module):\n", " def __init__(self, nb_patches, dim, layerscale_init):\n", " super().__init__()\n", "\n", " self.affine1 = init(Mix('b t c -> b t c', weight_shape='c', bias_shape='c', c=dim))\n", " self.affine2 = init(Mix('b t c -> b t c', weight_shape='c', bias_shape='c', c=dim))\n", " self.mix_patches = Mix('b t c -> b t0 c', weight_shape='t t0', bias_shape='t0', t=nb_patches, t0=nb_patches)\n", " self.mlp_channels = Mlp(dim)\n", " self.linear1 = init(Mix('b t c -> b t c', weight_shape='c', c=dim), scale=layerscale_init)\n", " self.linear2 = init(Mix('b t c -> b t c', weight_shape='c', c=dim), scale=layerscale_init)\n", "\n", " def forward(self, x):\n", " res1 = self.mix_patches(self.affine1(x))\n", " x = x + self.linear1(res1)\n", " res2 = self.mlp_channels(self.affine2(x))\n", " x = x + self.linear2(res2)\n", " return x" ] }, { "cell_type": "markdown", "id": "761d0d39-b041-4389-a19d-c3d512d64624", "metadata": {}, "source": [ "## ResMLP — rewritten more\n", "\n", "Since here in einops-land we care about code being easy to follow, let's make one more transformation.\n", "\n", "We group layers from both branches, and now the order of operations matches the order as they are written in the code.\n", "\n", "Could we go further? Actually, yes - `nn.Linear` layers can also be replaced by EinMix,\n", "however they are very organic here since first and last operations in `branch_channels` show components.\n", "\n", "Brevity of `nn.Linear` is benefitial when the context specifies tensor shapes.\n", "\n", "Other interesing observations:\n", "- hard to notice in the original code `nn.Linear` is preceded by a linear layer (thus latter is redundant or can be fused in the former)\n", "- hard to notice in the original code second `nn.Linear` is followed by an affine layer (thus latter is again redundant)\n", "\n", "Take time to reorganize your code. This may be quite insightful." ] }, { "cell_type": "code", "execution_count": 5, "id": "6fb50fac-1ef4-45a2-acdb-48f1050df7ca", "metadata": {}, "outputs": [], "source": [ "def init(layer: Mix, scale=1.):\n", " layer.weight.data[:] = scale\n", " if layer.bias is not None:\n", " layer.bias.data[:] = 0\n", " return layer\n", "\n", "class ResMLP_Blocks3(nn.Module):\n", " def __init__(self, nb_patches, dim, layerscale_init):\n", " super().__init__()\n", " self.branch_patches = nn.Sequential(\n", " init(Mix('b t c -> b t c', weight_shape='c', c=dim), scale=layerscale_init),\n", " Mix('b t c -> b t0 c', weight_shape='t t0', bias_shape='t0', t=nb_patches, t0=nb_patches),\n", " init(Mix('b t c -> b t c', weight_shape='c', bias_shape='c', c=dim)),\n", " )\n", "\n", " self.branch_channels = nn.Sequential(\n", " init(Mix('b t c -> b t c', weight_shape='c', c=dim), scale=layerscale_init),\n", " nn.Linear(dim, 4 * dim),\n", " nn.GELU(),\n", " nn.Linear(4 * dim, dim),\n", " init(Mix('b t c -> b t c', weight_shape='c', bias_shape='c', c=dim)),\n", " )\n", "\n", " def forward(self, x):\n", " x = x + self.branch_patches(x)\n", " x = x + self.branch_channels(x)\n", " return x" ] }, { "cell_type": "markdown", "id": "ad7f343e-b81d-4b25-96bd-4c48ff33f9c8", "metadata": {}, "source": [ "## ResMLP — performance\n", "\n", "There is some fear of using einsum because historically it lagged in performance.\n", "\n", "Below we run a test and verify that performace didn't change after transition to `EinMix`" ] }, { "cell_type": "code", "execution_count": 6, "id": "3987bbbc-ccc2-4374-8979-fd3fb29c4910", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "28.1 ms ± 1.61 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "26.3 ms ± 620 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "25.9 ms ± 706 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "26.8 ms ± 2.99 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "25.9 ms ± 794 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "25.6 ms ± 723 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" ] } ], "source": [ "x = torch.zeros([32, 128, 128])\n", "for layer in [\n", " ResMLP_Blocks(128, dim=128, layerscale_init=1.),\n", " ResMLP_Blocks2(128, dim=128, layerscale_init=1.),\n", " ResMLP_Blocks3(128, dim=128, layerscale_init=1.),\n", " # scripted versions\n", " torch.jit.script(ResMLP_Blocks(128, dim=128, layerscale_init=1.)),\n", " torch.jit.script(ResMLP_Blocks2(128, dim=128, layerscale_init=1.)),\n", " torch.jit.script(ResMLP_Blocks3(128, dim=128, layerscale_init=1.)),\n", "]:\n", " %timeit -n 10 y = layer(x)" ] }, { "cell_type": "markdown", "id": "05428b35-5788-4e7f-867e-aa742d7215e7", "metadata": {}, "source": [ "## TokenMixer from MLPMixer — original code\n", "\n", "Let's now delve into MLPMixer. We start from pytorch [implementation](https://github.com/jaketae/mlp-mixer/blob/e7d68dfc31e94721724689e6ec90f05806b50124/mlp_mixer/core.py) by Jake Tae.\n", "\n", "We'll focus on two components of MLPMixer that don't exist in convnets. First component is TokenMixer:" ] }, { "cell_type": "code", "execution_count": 7, "id": "8f6a6073-44af-40ec-8154-02ae8370f749", "metadata": {}, "outputs": [], "source": [ "from torch.nn import functional as F\n", "\n", "class MLP(nn.Module):\n", " def __init__(self, num_features, expansion_factor, dropout):\n", " super().__init__()\n", " num_hidden = num_features * expansion_factor\n", " self.fc1 = nn.Linear(num_features, num_hidden)\n", " self.dropout1 = nn.Dropout(dropout)\n", " self.fc2 = nn.Linear(num_hidden, num_features)\n", " self.dropout2 = nn.Dropout(dropout)\n", "\n", " def forward(self, x):\n", " x = self.dropout1(F.gelu(self.fc1(x)))\n", " x = self.dropout2(self.fc2(x))\n", " return x\n", "\n", "\n", "class TokenMixer(nn.Module):\n", " def __init__(self, num_features, num_patches, expansion_factor, dropout):\n", " super().__init__()\n", " self.norm = nn.LayerNorm(num_features)\n", " self.mlp = MLP(num_patches, expansion_factor, dropout)\n", "\n", " def forward(self, x):\n", " # x.shape == (batch_size, num_patches, num_features)\n", " residual = x\n", " x = self.norm(x)\n", " x = x.transpose(1, 2)\n", " # x.shape == (batch_size, num_features, num_patches)\n", " x = self.mlp(x)\n", " x = x.transpose(1, 2)\n", " # x.shape == (batch_size, num_patches, num_features)\n", " out = x + residual\n", " return out" ] }, { "cell_type": "markdown", "id": "362fa6ac-7146-4f7b-9e1d-0765486c26da", "metadata": {}, "source": [ "## TokenMixer from MLPMixer — reimplemented\n", "\n", "We can significantly reduce amount of code by using `EinMix`. \n", "\n", "- Main caveat addressed by original code is that `nn.Linear` mixes only last axis. `EinMix` can mix any axis.\n", "- Sequential structure is always preferred as it is easier to follow\n", "- Intentionally there is no residual connection in `TokenMixer`, because honestly it's not work of Mixer and should be done by caller\n" ] }, { "cell_type": "code", "execution_count": 8, "id": "f447c6c5-932f-4929-991a-f97b14776618", "metadata": {}, "outputs": [], "source": [ "def TokenMixer(num_features: int, n_patches: int, expansion_factor: int, dropout: float):\n", " n_hidden = n_patches * expansion_factor\n", " return nn.Sequential(\n", " nn.LayerNorm(num_features),\n", " Mix('b hw c -> b hid c', weight_shape='hw hid', bias_shape='hid', hw=n_patches, hidden=n_hidden),\n", " nn.GELU(),\n", " nn.Dropout(dropout),\n", " Mix('b hid c -> b hw c', weight_shape='hid hw', bias_shape='hw', hw=n_patches, hidden=n_hidden),\n", " nn.Dropout(dropout),\n", " )" ] }, { "cell_type": "markdown", "id": "19caacfc-a066-4727-918e-b3a98a190014", "metadata": {}, "source": [ "You may also like independent [implementation](https://github.com/lucidrains/mlp-mixer-pytorch/blob/main/mlp_mixer_pytorch/mlp_mixer_pytorch.py) of MLPMixer from Phil Wang.
\n", "Phil solves the issue by repurposing `nn.Conv1d` to mix on the second dimension. Hacky, but does the job\n" ] }, { "cell_type": "markdown", "id": "5210dfd1-0195-458f-a8b9-819547123bad", "metadata": {}, "source": [ "## MLPMixer's patch embeddings — original\n", "\n", "Second interesting part of MLPMixer is derived from vision transformers.\n", "\n", "In the very beginning an image is split into patches, and each patch is linearly projected into embedding.\n", "\n", "I've taken the part of Jake's code responsible for embedding patches:" ] }, { "cell_type": "code", "execution_count": 9, "id": "24752920-006c-49e1-acb9-227543f84afb", "metadata": {}, "outputs": [], "source": [ "def check_sizes(image_size, patch_size):\n", " sqrt_num_patches, remainder = divmod(image_size, patch_size)\n", " assert remainder == 0, \"`image_size` must be divisibe by `patch_size`\"\n", " num_patches = sqrt_num_patches ** 2\n", " return num_patches\n", "\n", "class Patcher(nn.Module):\n", " def __init__(\n", " self,\n", " image_size=256,\n", " patch_size=16,\n", " in_channels=3,\n", " num_features=128,\n", " ):\n", " _num_patches = check_sizes(image_size, patch_size)\n", " super().__init__()\n", " # per-patch fully-connected is equivalent to strided conv2d\n", " self.patcher = nn.Conv2d(\n", " in_channels, num_features, kernel_size=patch_size, stride=patch_size\n", " )\n", "\n", " def forward(self, x):\n", " patches = self.patcher(x)\n", " batch_size, num_features, _, _ = patches.shape\n", " patches = patches.permute(0, 2, 3, 1)\n", " patches = patches.view(batch_size, -1, num_features)\n", "\n", " return patches" ] }, { "cell_type": "markdown", "id": "b5ecaac5-d5fe-48b7-9339-fe84fdccd9ac", "metadata": {}, "source": [ "## MLPMixer's patch embeddings — reimplemented\n", "\n", "`EinMix` does this in a single operation. This may require some training at first to understand.\n", "\n", "Let's go step-by-step:\n", "\n", "- `b c_in (h hp) (w wp) ->` - 4-dimensional input tensor (BCHW-ordered) is split into patches of shape `hp x wp`\n", "- `weight_shape='c_in hp wp c'`. Axes `c_in`, `hp` and `wp` are all absent in the output: three dimensional patch tensor was *mixed* to produce a vector of length `c`\n", "- `-> b (h w) c` - output is 3-dimensional. All patches were reorganized from `h x w` grid to one-dimensional sequence of vectors\n", "\n", "\n", "We don't need to provide image_size beforehead, new implementation handles images of different dimensions as long as they can be divided into patches" ] }, { "cell_type": "code", "execution_count": 10, "id": "2b786a62-a742-49d3-9e66-d02eb2f4e384", "metadata": {}, "outputs": [], "source": [ "def patcher(patch_size=16, in_channels=3, num_features=128):\n", " return Mix('b c_in (h hp) (w wp) -> b (h w) c', weight_shape='c_in hp wp c', bias_shape='c',\n", " c=num_features, hp=patch_size, wp=patch_size, c_in=in_channels)" ] }, { "cell_type": "markdown", "id": "a26b2f34-07cd-4fc2-bdf9-0043dd3451ea", "metadata": {}, "source": [ "## Vision Permutator\n", "\n", "As a third example we consider pytorch-like code from [ViP paper](https://arxiv.org/pdf/2106.12368.pdf).\n", "\n", "Vision permutator is only slightly more nuanced than previous models, because \n", "1. it operates on spatial dimensions separately, while MLPMixer and its friends just pack all spatial info into one axis. \n", "2. it splits channels into groups called 'segments'\n", "\n", "Paper provides pseudo-code, so I reworked that to complete module with minimal changes. Enjoy:" ] }, { "cell_type": "code", "execution_count": 11, "id": "2c94453b-bc78-4f06-8d51-f110e4b6501a", "metadata": {}, "outputs": [], "source": [ "class WeightedPermuteMLP(nn.Module):\n", " def __init__(self, H, W, C, S):\n", " super().__init__()\n", "\n", " self.proj_h = nn.Linear(H * S, H * S)\n", " self.proj_w = nn.Linear(W * S, W * S)\n", " self.proj_c = nn.Linear(C, C)\n", " self.proj = nn.Linear(C, C)\n", " self.S = S\n", "\n", " def forward(self, x):\n", " B, H, W, C = x.shape\n", " S = self.S\n", " N = C // S\n", " x_h = x.reshape(B, H, W, N, S).permute(0, 3, 2, 1, 4).reshape(B, N, W, H*S)\n", " x_h = self.proj_h(x_h).reshape(B, N, W, H, S).permute(0, 3, 2, 1, 4).reshape(B, H, W, C)\n", "\n", " x_w = x.reshape(B, H, W, N, S).permute(0, 1, 3, 2, 4).reshape(B, H, N, W*S)\n", " x_w = self.proj_w(x_w).reshape(B, H, N, W, S).permute(0, 1, 3, 2, 4).reshape(B, H, W, C)\n", "\n", " x_c = self.proj_c(x)\n", "\n", " x = x_h + x_w + x_c\n", " x = self.proj(x)\n", " return x" ] }, { "cell_type": "markdown", "id": "27a4aa53-e719-4079-973a-83b07d507c1f", "metadata": {}, "source": [ "That didn't look readable, right? \n", "\n", "This code is also very inflexible: code in the paper did not support batch dimension, and multiple changes were necessary to allow batch processing.
\n", "This process is fragile and easily can result in virtually uncatchable bugs.\n", "\n", "Now good news: each of these long method chains can be replaced with a single `EinMix` layer:" ] }, { "cell_type": "code", "execution_count": 12, "id": "a0a699c0-f0b9-4e2c-8dc1-b33f5a7dfa1e", "metadata": {}, "outputs": [], "source": [ "class WeightedPermuteMLP_new(nn.Module):\n", " def __init__(self, H, W, C, seg_len):\n", " super().__init__()\n", " assert C % seg_len == 0, f\"can't divide {C} into segments of length {seg_len}\"\n", " self.mlp_c = Mix('b h w c -> b h w c0', weight_shape='c c0', bias_shape='c0', c=C, c0=C)\n", " self.mlp_h = Mix('b h w (n c) -> b h0 w (n c0)', weight_shape='h c h0 c0', bias_shape='h0 c0',\n", " h=H, h0=H, c=seg_len, c0=seg_len)\n", " self.mlp_w = Mix('b h w (n c) -> b h w0 (n c0)', weight_shape='w c w0 c0', bias_shape='w0 c0',\n", " w=W, w0=W, c=seg_len, c0=seg_len)\n", " self.proj = nn.Linear(C, C)\n", "\n", " def forward(self, x):\n", " x = self.mlp_c(x) + self.mlp_h(x) + self.mlp_w(x)\n", " return self.proj(x)" ] }, { "cell_type": "markdown", "id": "7a9a7165-6c58-409f-b277-f4f7314e57fe", "metadata": {}, "source": [ "Great, now let's confirm that performance did not deteriorate." ] }, { "cell_type": "code", "execution_count": 13, "id": "f375fcd7-8238-470d-aeaf-976892062911", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "90.5 ms ± 1.22 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "91.5 ms ± 616 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "91.8 ms ± 626 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", "87.4 ms ± 3.59 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" ] } ], "source": [ "x = torch.zeros([32, 32, 32, 128])\n", "\n", "for layer in [\n", " WeightedPermuteMLP(H=32, W=32, C=128, S=4),\n", " WeightedPermuteMLP_new(H=32, W=32, C=128, seg_len=4),\n", " # scripted versions\n", " torch.jit.script(WeightedPermuteMLP(H=32, W=32, C=128, S=4)),\n", " torch.jit.script(WeightedPermuteMLP_new(H=32, W=32, C=128, seg_len=4)),\n", "]:\n", " %timeit -n 10 y = layer(x)" ] }, { "cell_type": "markdown", "id": "6694ade8-557d-4461-8f3a-4bd5f74dbea2", "metadata": {}, "source": [ "## Final remarks\n", "\n", "`EinMix` has an incredible potential: \n", "it helps with MLPs that don't fit into a limited 'mix all in the last axis' paradigm.\n", "\n", "However existing research is ... very limited, it does not cover real possibilities of densely connected architectures.\n", "\n", "Most of its *systematic* novelty is \"mix along spacial axes too\". \n", "But `EinMix` provides **an astonishing amount of other possibilities!**" ] }, { "cell_type": "markdown", "id": "a1ee15e6-933a-4c3e-842e-e908a0aeee15", "metadata": {}, "source": [ "### Groups of mixers\n", "\n", "You can find two settings compared in the MLPMixer paper (Supplementary A1)\n", "```python\n", "'b hw c -> b hw_out c', weight_shape='hw hw_out'\n", "```\n", "and\n", "```python\n", "'b hw c -> b hw_out c', weight_shape='c hw hw_out'\n", "```\n", "While latter makes more sense (why mixing should work similarly for all channels?), the former performs better.\n", "\n", "So one more question is reasonable: what if channels are split into groups, and mixing is defined for each group?\n", "```python\n", "'b hw (group c) -> b hw_out (group c)', weight_shape='group hw hw_out'\n", "```\n", "Implementing such setting without einops is considerably harder.\n", "\n", "### Mixing within patch on a grid\n", "\n", "What if you make mixing 'local' in space? Completely doable:\n", "\n", "```python\n", "'b c (h hI) (w wI) -> b c (h hO) (w wO)', weight_shape='c hI wI hO wO'\n", "```\n", "\n", "We split tensor into patches of shape `hI wI` and mixed things within channel.\n", " \n", "### Mixing in subgrids\n", "\n", "Ok, done with local mixing. How to collect information from the whole image?
\n", "Well, you can again densely connect all the tokens, but all-to-all connection is too expensive.\n", "\n", "\n", "\n", "TODO need some image here to show sub-grids and information exhange. \n", "\n", "\n", "Here is EinMix-way: split the image into subgrids (each subgrid has steps `h` and `w`), and connect densely tokens within each subgrid\n", "\n", "```python\n", "'b c (hI h) (wI w) -> b c (hO h) (wO w)', weight_shape='c hI wI hO wO'\n", "```\n", "\n", "\n", "### Going deeper\n", "And that's very top of the iceberg.
\n", "\n", "- Want to mix part of axis? — No problems!\n", "- ... in a grid-like manner — Supported! \n", "- ... while mixing channels within group? — Welcome! \n", "- In 2d/3d/4d? — Sure!\n", "- I don't use pytorch — EinMix is available for multiple frameworks!\n", "\n", "Hopefully this guide helped you to find MLPs more interesting and intriguing. And simpler to experiment with." ] } ], "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.9.0" } }, "nbformat": 4, "nbformat_minor": 5 } arogozhnikov-einops-ad2c8d6/docs/4-pack-and-unpack.ipynb 0000664 0000000 0000000 00000041274 14752016746 0023337 0 ustar 00root root 0000000 0000000 { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# einops.pack and einops.unpack\n", "\n", "einops 0.6 introduces two more functions to the family: `pack` and `unpack`.\n", "\n", "Here is what they do:\n", "\n", "- `unpack` reverses `pack`\n", "- `pack` reverses `unpack`\n", "\n", "Enlightened with this exhaustive description, let's move to examples.\n", "\n" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "# we'll use numpy for demo purposes\n", "# operations work the same way with other frameworks\n", "import numpy as np" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Stacking data layers\n", "\n", "Assume we have RGB image along with a corresponding depth image that we want to stack:" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "from einops import pack, unpack\n", "\n", "h, w = 100, 200\n", "# image_rgb is 3-dimensional (h, w, 3) and depth is 2-dimensional (h, w)\n", "image_rgb = np.random.random([h, w, 3])\n", "image_depth = np.random.random([h, w])\n", "# but we can stack them\n", "image_rgbd, ps = pack([image_rgb, image_depth], 'h w *')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## How to read packing patterns\n", "\n", "pattern `h w *` means that\n", "- output is 3-dimensional\n", "- first two axes (`h` and `w`) are shared across all inputs and also shared with output\n", "- inputs, however do not have to be 3-dimensional. They can be 2-dim, 3-dim, 4-dim, etc.
\n", " Regardless of inputs dimensionality, they all will be packed into 3-dim output, and information about how they were packed is stored in `PS`" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [ { "data": { "text/plain": [ "((100, 200, 3), (100, 200), (100, 200, 4))" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# as you see, pack properly appended depth as one more layer\n", "# and correctly aligned axes!\n", "# this won't work off the shelf with np.concatenate or torch.cat or alike\n", "image_rgb.shape, image_depth.shape, image_rgbd.shape" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [ { "data": { "text/plain": [ "[(3,), ()]" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# now let's see what PS keeps.\n", "# PS means Packed Shapes, not PlayStation or Post Script\n", "ps" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "which reads: first tensor had shape `h, w, *and 3*`, while second tensor had shape `h, w *and nothing more*`.\n", "That's just enough to reverse packing:" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [ { "data": { "text/plain": [ "((100, 200, 3), (100, 200))" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# remove 1-axis in depth image during unpacking. Results are (h, w, 3) and (h, w)\n", "unpacked_rgb, unpacked_depth = unpack(image_rgbd, ps, 'h w *')\n", "unpacked_rgb.shape, unpacked_depth.shape" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "we can unpack tensor in different ways manually:" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "# simple unpack by splitting the axis. Results are (h, w, 3) and (h, w, 1)\n", "rgb, depth = unpack(image_rgbd, [[3], [1]], 'h w *')\n", "# different split, both outputs have shape (h, w, 2)\n", "rg, bd = unpack(image_rgbd, [[2], [2]], 'h w *')\n", "# unpack to 4 tensors of shape (h, w). More like 'unstack over last axis'\n", "[r, g, b, d] = unpack(image_rgbd, [[], [], [], []], 'h w *')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Short summary so far\n", "\n", "- `einops.pack` is a 'more generic concatenation' (that can stack too)\n", "- `einops.unpack` is a 'more generic split'\n", "\n", "And, of course, `einops` functions are more verbose, and *reversing* concatenation now is *dead simple*\n", "\n", "Compared to other `einops` functions, `pack` and `unpack` have a compact pattern without arrow, and the same pattern can be used in `pack` and `unpack`. These patterns are very simplistic: just a sequence of space-separated axes names.\n", "One axis is `*`, all other axes are valid identifiers.\n", "\n", "Now let's discuss some practical cases" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Auto-batching\n", "\n", "ML models by default accept batches: batch of images, or batch of sentences, or batch of audios, etc.\n", "\n", "During debugging or inference, however, it is common to pass a single image instead (and thus output should be a single prediction)
\n", "In this example we'll write `universal_predict` that can handle both cases." ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "from einops import reduce\n", "def image_classifier(images_bhwc):\n", " # mock for image classifier\n", " predictions = reduce(images_bhwc, 'b h w c -> b c', 'mean', h=100, w=200, c=3)\n", " return predictions\n", "\n", "\n", "def universal_predict(x):\n", " x_packed, ps = pack([x], '* h w c')\n", " predictions_packed = image_classifier(x_packed)\n", " [predictions] = unpack(predictions_packed, ps, '* cls')\n", " return predictions" ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "(3,)\n", "(5, 3)\n", "(5, 7, 3)\n" ] } ], "source": [ "# works with a single image\n", "print(universal_predict(np.zeros([h, w, 3])).shape)\n", "# works with a batch of images\n", "batch = 5\n", "print(universal_predict(np.zeros([batch, h, w, 3])).shape)\n", "# or even a batch of videos\n", "n_frames = 7\n", "print(universal_predict(np.zeros([batch, n_frames, h, w, 3])).shape)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**what we can learn from this example**:\n", "\n", "- `pack` and `unpack` play nicely together. That's not a coincidence :)\n", "- patterns in `pack` and `unpack` may differ, and that's quite common for applications\n", "- unlike other operations in `einops`, `(un)pack` does not provide arbitrary reordering of axes" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Class token in VIT\n", "\n", "Let's assume we have a simple transformer model that works with `BTC`-shaped tensors." ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "def transformer_mock(x_btc):\n", " # imagine this is a transformer model, a very efficient one\n", " assert len(x_btc.shape) == 3\n", " return x_btc" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's implement vision transformer (ViT) with a class token (i.e. static token, corresponding output is used to classify an image)" ] }, { "cell_type": "code", "execution_count": 10, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "# below it is assumed that you already\n", "# 1) split batch of images into patches 2) applied linear projection and 3) used positional embedding.\n", "\n", "# We'll skip that here. But hey, here is an einops-style way of doing all of that in a single shot!\n", "# from einops.layers.torch import EinMix\n", "# patcher_and_posembedder = EinMix('b (h h2) (w w2) c -> b h w c_out', weight_shape='h2 w2 c c_out',\n", "# bias_shape='h w c_out', h2=..., w2=...)\n", "# patch_tokens_bhwc = patcher_and_posembedder(images_bhwc)" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "# preparations\n", "batch, height, width, c = 6, 16, 16, 256\n", "patch_tokens = np.random.random([batch, height, width, c])\n", "class_tokens = np.zeros([batch, c])" ] }, { "cell_type": "code", "execution_count": 12, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [ { "data": { "text/plain": [ "((6, 256), (6, 16, 16, 256))" ] }, "execution_count": 12, "metadata": {}, "output_type": "execute_result" } ], "source": [ "def vit_einops(class_tokens, patch_tokens):\n", " input_packed, ps = pack([class_tokens, patch_tokens], 'b * c')\n", " output_packed = transformer_mock(input_packed)\n", " return unpack(output_packed, ps, 'b * c_out')\n", "\n", "class_token_emb, patch_tokens_emb = vit_einops(class_tokens, patch_tokens)\n", "\n", "class_token_emb.shape, patch_tokens_emb.shape" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "At this point, let's make a small pause and understand conveniences of this pipeline, by contrasting it to more 'standard' code" ] }, { "cell_type": "code", "execution_count": 13, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "def vit_vanilla(class_tokens, patch_tokens):\n", " b, h, w, c = patch_tokens.shape\n", " class_tokens_b1c = class_tokens[:, np.newaxis, :]\n", " patch_tokens_btc = np.reshape(patch_tokens, [b, -1, c])\n", " input_packed = np.concatenate([class_tokens_b1c, patch_tokens_btc], axis=1)\n", " output_packed = transformer_mock(input_packed)\n", " class_token_emb = np.squeeze(output_packed[:, :1, :], 1)\n", " patch_tokens_emb = np.reshape(output_packed[:, 1:, :], [b, h, w, -1])\n", " return class_token_emb, patch_tokens_emb\n", "\n", "class_token_emb2, patch_tokens_emb2 = vit_vanilla(class_tokens, patch_tokens)\n", "assert np.allclose(class_token_emb, class_token_emb2)\n", "assert np.allclose(patch_tokens_emb, patch_tokens_emb2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Notably, we have put all packing and unpacking, reshapes, adding and removing of dummy axes into a couple of lines." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Packing different modalities together\n", "\n", "We can extend the previous example: it is quite common to mix elements of different types of inputs in transformers.\n", "\n", "The simples one is to mix tokens from all inputs:\n", "\n", "```python\n", "all_inputs = [text_tokens_btc, image_bhwc, task_token_bc, static_tokens_bnc]\n", "inputs_packed, ps = pack(all_inputs, 'b * c')\n", "```\n", "\n", "and you can `unpack` resulting tokens to the same structure." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Packing data coming from different sources together\n", "\n", "Most notable example is of course GANs:\n", "\n", "```python\n", "input_ims, ps = pack([true_images, fake_images], '* h w c')\n", "true_pred, fake_pred = unpack(model(input_ims), ps, '* c')\n", "```\n", "`true_pred` and `fake_pred` are handled differently, that's why we separated them" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Predicting multiple outputs at the same time\n", "\n", "It is quite common to pack prediction of multiple target values into a single layer.\n", "\n", "This is more efficient, but code is less readable. For example, that's how detection code may look like:" ] }, { "cell_type": "code", "execution_count": 14, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "def loss_detection(model_output_bhwc, mask_h: int, mask_w: int, n_classes: int):\n", " output = model_output_bhwc\n", "\n", " confidence = output[..., 0].sigmoid()\n", " bbox_x_shift = output[..., 1].sigmoid()\n", " bbox_y_shift = output[..., 2].sigmoid()\n", " bbox_w = output[..., 3]\n", " bbox_h = output[..., 4]\n", " mask_logits = output[..., 5: 5 + mask_h * mask_w]\n", " mask_logits = mask_logits.reshape([*mask_logits.shape[:-1], mask_h, mask_w])\n", " class_logits = output[..., 5 + mask_h * mask_w:]\n", " assert class_logits.shape[-1] == n_classes, class_logits.shape[-1]\n", "\n", " # downstream computations\n", " return confidence, bbox_x_shift, bbox_y_shift, bbox_h, bbox_w, mask_logits, class_logits" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "When the same logic is implemented in einops, there is no need to memorize offsets.
\n", "Additionally, reshapes and shape checks are automatic:" ] }, { "cell_type": "code", "execution_count": 15, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "def loss_detection_einops(model_output, mask_h: int, mask_w: int, n_classes: int):\n", " confidence, bbox_x_shift, bbox_y_shift, bbox_w, bbox_h, mask_logits, class_logits \\\n", " = unpack(model_output, [[]] * 5 + [[mask_h, mask_w], [n_classes]], 'b h w *')\n", "\n", " confidence = confidence.sigmoid()\n", " bbox_x_shift = bbox_x_shift.sigmoid()\n", " bbox_y_shift = bbox_y_shift.sigmoid()\n", "\n", " # downstream computations\n", " return confidence, bbox_x_shift, bbox_y_shift, bbox_h, bbox_w, mask_logits, class_logits" ] }, { "cell_type": "code", "execution_count": 16, "metadata": { "collapsed": false, "jupyter": { "outputs_hidden": false } }, "outputs": [], "source": [ "# check that results are identical\n", "import torch\n", "dims = dict(mask_h=6, mask_w=8, n_classes=19)\n", "model_output = torch.randn([3, 5, 7, 5 + dims['mask_h'] * dims['mask_w'] + dims['n_classes']])\n", "for a, b in zip(loss_detection(model_output, **dims), loss_detection_einops(model_output, **dims)):\n", " assert torch.allclose(a, b)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Or maybe **reinforcement learning** is closer to your mind?\n", "\n", "If so, predicting multiple outputs is valuable there too:\n", "\n", "```python\n", "action_logits, reward_expectation, q_values, expected_entropy_after_action = \\\n", " unpack(predictions_btc, [[n_actions], [], [n_actions], [n_actions]], 'b step *')\n", "\n", "\n", "```\n", "\n", "\n", "## That's all for today!\n", "\n", "happy packing and unpacking!" ] }, { "cell_type": "markdown", "metadata": {}, "source": [] } ], "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.10.6" } }, "nbformat": 4, "nbformat_minor": 4 } arogozhnikov-einops-ad2c8d6/docs/README.md 0000664 0000000 0000000 00000001415 14752016746 0020446 0 ustar 00root root 0000000 0000000 # Einops tutorial You can read notebooks from github by clicking on files above. If you prefer nbviewer, use the links below: - Part 1 notebook at [nbviewer](https://nbviewer.jupyter.org/github/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb) - Part 2 notebook at [nbviewer](https://nbviewer.jupyter.org/github/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb) - Part 3: EinMix for great MLPs at [nbviewer](https://nbviewer.jupyter.org/github/arogozhnikov/einops/blob/main/docs/3-einmix-layer.ipynb) - Part 4: einops.pack and einops.unpack [nbviewer](https://nbviewer.jupyter.org/github/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb) - Writing better code with einops and pytorch: [web link](https://einops.rocks/pytorch-examples.html) arogozhnikov-einops-ad2c8d6/docs/pytorch-examples.html 0000664 0000000 0000000 00000772117 14752016746 0023377 0 ustar 00root root 0000000 0000000
Writing a better code with pytorch and einops
Rewriting building blocks of deep learning
Now let's get to examples from real world. These code fragments taken from official tutorials and popular repositories.
Learn how to improve code and how einops can help you.
Left: as it was, Right: improved version
# start from importing some stuff
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
from einops import rearrange, reduce, asnumpy, parse_shape
from einops.layers.torch import Rearrange, Reduce
Simple ConvNet
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
conv_net_old = Net()
conv_net_new = nn.Sequential(
nn.Conv2d(1, 10, kernel_size=5),
nn.MaxPool2d(kernel_size=2),
nn.ReLU(),
nn.Conv2d(10, 20, kernel_size=5),
nn.MaxPool2d(kernel_size=2),
nn.ReLU(),
nn.Dropout2d(),
Rearrange('b c h w -> b (c h w)'),
nn.Linear(320, 50),
nn.ReLU(),
nn.Dropout(),
nn.Linear(50, 10),
nn.LogSoftmax(dim=1)
)
Reasons to prefer new implementation:
- in the original code (to the left) if input size is changed and batch size is divisible by 16 (that's usually so), we'll get something senseless after reshaping
- new code will explicitly raise an error in this case
- we won't forget to use dropout with flag self.training with new version
- code is straightforward to read and analyze
- sequential makes printing / saving / passing trivial. And there is no need in your code to load a model (which also has a number of benefits)
- don't need logsoftmax? Now you can use
conv_net_new[:-1]. One more reason to prefernn.Sequential - ... and we could also add inplace for ReLU
Super-resolution
class SuperResolutionNetOld(nn.Module):
def __init__(self, upscale_factor):
super(SuperResolutionNetOld, self).__init__()
self.relu = nn.ReLU()
self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
def forward(self, x):
x = self.relu(self.conv1(x))
x = self.relu(self.conv2(x))
x = self.relu(self.conv3(x))
x = self.pixel_shuffle(self.conv4(x))
return x
def SuperResolutionNetNew(upscale_factor):
return nn.Sequential(
nn.Conv2d(1, 64, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(64, 32, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(32, upscale_factor ** 2, kernel_size=3, padding=1),
Rearrange('b (h2 w2) h w -> b (h h2) (w w2)', h2=upscale_factor, w2=upscale_factor),
)
Here is the difference:
- no need in special instruction pixel_shuffle (and result is transferrable between frameworks)
- output doesn't contain a fake axis (and we could do the same for the input)
- inplace ReLU used now, for high resolution pictures that becomes critical and saves us much memory
- and all the benefits of nn.Sequential again
Restyling Gram matrix for style transfer
Original code is already good - first line shows what kind of input is expected
- einsum operation should be read like:
- for each batch and for each pair of channels, we sum over h and w.
- I've also changed normalization, because that's how Gram matrix is defined, otherwise we should call it normalized Gram matrix or alike
def gram_matrix_old(y):
(b, ch, h, w) = y.size()
features = y.view(b, ch, w * h)
features_t = features.transpose(1, 2)
gram = features.bmm(features_t) / (ch * h * w)
return gram
def gram_matrix_new(y):
b, ch, h, w = y.shape
return torch.einsum('bchw,bdhw->bcd', [y, y]) / (h * w)
It would be great to use just 'b c1 h w,b c2 h w->b c1 c2', but einsum supports only one-letter axes
Recurrent model
All we did here is just made information about shapes explicit to skip deciphering
class RNNModelOld(nn.Module):
"""Container module with an encoder, a recurrent module, and a decoder."""
def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):
super(RNNModel, self).__init__()
self.drop = nn.Dropout(dropout)
self.encoder = nn.Embedding(ntoken, ninp)
self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
self.decoder = nn.Linear(nhid, ntoken)
def forward(self, input, hidden):
emb = self.drop(self.encoder(input))
output, hidden = self.rnn(emb, hidden)
output = self.drop(output)
decoded = self.decoder(output.view(output.size(0)*output.size(1), output.size(2)))
return decoded.view(output.size(0), output.size(1), decoded.size(1)), hidden
class RNNModelNew(nn.Module):
"""Container module with an encoder, a recurrent module, and a decoder."""
def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):
super(RNNModel, self).__init__()
self.drop = nn.Dropout(p=dropout)
self.encoder = nn.Embedding(ntoken, ninp)
self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
self.decoder = nn.Linear(nhid, ntoken)
def forward(self, input, hidden):
t, b = input.shape
emb = self.drop(self.encoder(input))
output, hidden = self.rnn(emb, hidden)
output = rearrange(self.drop(output), 't b nhid -> (t b) nhid')
decoded = rearrange(self.decoder(output), '(t b) token -> t b token', t=t, b=b)
return decoded, hidden
Channel shuffle (from shufflenet)
def channel_shuffle_old(x, groups):
batchsize, num_channels, height, width = x.data.size()
channels_per_group = num_channels // groups
# reshape
x = x.view(batchsize, groups,
channels_per_group, height, width)
# transpose
# - contiguous() required if transpose() is used before view().
# See https://github.com/pytorch/pytorch/issues/764
x = torch.transpose(x, 1, 2).contiguous()
# flatten
x = x.view(batchsize, -1, height, width)
return x
def channel_shuffle_new(x, groups):
return rearrange(x, 'b (c1 c2) h w -> b (c2 c1) h w', c1=groups)
While progress is obvious, this is not the limit. As you'll see below, we don't even need to write these couple of lines.
Shufflenet
from collections import OrderedDict
def channel_shuffle(x, groups):
batchsize, num_channels, height, width = x.data.size()
channels_per_group = num_channels // groups
# reshape
x = x.view(batchsize, groups,
channels_per_group, height, width)
# transpose
# - contiguous() required if transpose() is used before view().
# See https://github.com/pytorch/pytorch/issues/764
x = torch.transpose(x, 1, 2).contiguous()
# flatten
x = x.view(batchsize, -1, height, width)
return x
class ShuffleUnitOld(nn.Module):
def __init__(self, in_channels, out_channels, groups=3,
grouped_conv=True, combine='add'):
super(ShuffleUnitOld, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.grouped_conv = grouped_conv
self.combine = combine
self.groups = groups
self.bottleneck_channels = self.out_channels // 4
# define the type of ShuffleUnit
if self.combine == 'add':
# ShuffleUnit Figure 2b
self.depthwise_stride = 1
self._combine_func = self._add
elif self.combine == 'concat':
# ShuffleUnit Figure 2c
self.depthwise_stride = 2
self._combine_func = self._concat
# ensure output of concat has the same channels as
# original output channels.
self.out_channels -= self.in_channels
else:
raise ValueError("Cannot combine tensors with \"{}\"" \
"Only \"add\" and \"concat\" are" \
"supported".format(self.combine))
# Use a 1x1 grouped or non-grouped convolution to reduce input channels
# to bottleneck channels, as in a ResNet bottleneck module.
# NOTE: Do not use group convolution for the first conv1x1 in Stage 2.
self.first_1x1_groups = self.groups if grouped_conv else 1
self.g_conv_1x1_compress = self._make_grouped_conv1x1(
self.in_channels,
self.bottleneck_channels,
self.first_1x1_groups,
batch_norm=True,
relu=True
)
# 3x3 depthwise convolution followed by batch normalization
self.depthwise_conv3x3 = conv3x3(
self.bottleneck_channels, self.bottleneck_channels,
stride=self.depthwise_stride, groups=self.bottleneck_channels)
self.bn_after_depthwise = nn.BatchNorm2d(self.bottleneck_channels)
# Use 1x1 grouped convolution to expand from
# bottleneck_channels to out_channels
self.g_conv_1x1_expand = self._make_grouped_conv1x1(
self.bottleneck_channels,
self.out_channels,
self.groups,
batch_norm=True,
relu=False
)
@staticmethod
def _add(x, out):
# residual connection
return x + out
@staticmethod
def _concat(x, out):
# concatenate along channel axis
return torch.cat((x, out), 1)
def _make_grouped_conv1x1(self, in_channels, out_channels, groups,
batch_norm=True, relu=False):
modules = OrderedDict()
conv = conv1x1(in_channels, out_channels, groups=groups)
modules['conv1x1'] = conv
if batch_norm:
modules['batch_norm'] = nn.BatchNorm2d(out_channels)
if relu:
modules['relu'] = nn.ReLU()
if len(modules) > 1:
return nn.Sequential(modules)
else:
return conv
def forward(self, x):
# save for combining later with output
residual = x
if self.combine == 'concat':
residual = F.avg_pool2d(residual, kernel_size=3,
stride=2, padding=1)
out = self.g_conv_1x1_compress(x)
out = channel_shuffle(out, self.groups)
out = self.depthwise_conv3x3(out)
out = self.bn_after_depthwise(out)
out = self.g_conv_1x1_expand(out)
out = self._combine_func(residual, out)
return F.relu(out)
class ShuffleUnitNew(nn.Module):
def __init__(self, in_channels, out_channels, groups=3,
grouped_conv=True, combine='add'):
super().__init__()
first_1x1_groups = groups if grouped_conv else 1
bottleneck_channels = out_channels // 4
self.combine = combine
if combine == 'add':
# ShuffleUnit Figure 2b
self.left = Rearrange('...->...') # identity
depthwise_stride = 1
else:
# ShuffleUnit Figure 2c
self.left = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)
depthwise_stride = 2
# ensure output of concat has the same channels as original output channels.
out_channels -= in_channels
assert out_channels > 0
self.right = nn.Sequential(
# Use a 1x1 grouped or non-grouped convolution to reduce input channels
# to bottleneck channels, as in a ResNet bottleneck module.
conv1x1(in_channels, bottleneck_channels, groups=first_1x1_groups),
nn.BatchNorm2d(bottleneck_channels),
nn.ReLU(inplace=True),
# channel shuffle
Rearrange('b (c1 c2) h w -> b (c2 c1) h w', c1=groups),
# 3x3 depthwise convolution followed by batch
conv3x3(bottleneck_channels, bottleneck_channels,
stride=depthwise_stride, groups=bottleneck_channels),
nn.BatchNorm2d(bottleneck_channels),
# Use 1x1 grouped convolution to expand from
# bottleneck_channels to out_channels
conv1x1(bottleneck_channels, out_channels, groups=groups),
nn.BatchNorm2d(out_channels),
)
def forward(self, x):
if self.combine == 'add':
combined = self.left(x) + self.right(x)
else:
combined = torch.cat([self.left(x), self.right(x)], dim=1)
return F.relu(combined, inplace=True)
Rewriting the code helped to identify:
- There is no sense in doing reshuffling and not using groups in the first convolution (indeed, in the paper it is not so). However, result is an equivalent model.
- It is also strange that the first convolution may be not grouped, while the last convolution is always grouped (and that is different from the paper)
Other comments:
- There is an identity layer for pytorch introduced here
- The last thing left is get rid of conv1x1 and conv3x3 in the code - those are not better than standard
Simplifying ResNet
class ResNetOld(nn.Module):
def __init__(self, block, layers, num_classes=1000):
self.inplanes = 64
super(ResNetOld, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.avgpool = nn.AvgPool2d(7, stride=1)
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(self.inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * block.expansion),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
def make_layer(inplanes, planes, block, n_blocks, stride=1):
downsample = None
if stride != 1 or inplanes != planes * block.expansion:
# output size won't match input, so adjust residual
downsample = nn.Sequential(
nn.Conv2d(inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * block.expansion),
)
return nn.Sequential(
block(inplanes, planes, stride, downsample),
*[block(planes * block.expansion, planes) for _ in range(1, n_blocks)]
)
def ResNetNew(block, layers, num_classes=1000):
e = block.expansion
resnet = nn.Sequential(
Rearrange('b c h w -> b c h w', c=3, h=224, w=224),
nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
make_layer(64, 64, block, layers[0], stride=1),
make_layer(64 * e, 128, block, layers[1], stride=2),
make_layer(128 * e, 256, block, layers[2], stride=2),
make_layer(256 * e, 512, block, layers[3], stride=2),
# combined AvgPool and view in one averaging operation
Reduce('b c h w -> b c', 'mean'),
nn.Linear(512 * e, num_classes),
)
# initialization
for m in resnet.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
return resnet
Changes:
- explicit check for input shape
- no views and simple sequential structure, output is just nn.Sequential, so can always be saved/passed/etc
- no need in AvgPool and additional views, this place is much clearer now
make_layerdoesn't use internal state (that's quite faulty place)
Improving RNN language modelling
class RNNOld(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, bidirectional, dropout):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.rnn = nn.LSTM(embedding_dim, hidden_dim, num_layers=n_layers,
bidirectional=bidirectional, dropout=dropout)
self.fc = nn.Linear(hidden_dim*2, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
#x = [sent len, batch size]
embedded = self.dropout(self.embedding(x))
#embedded = [sent len, batch size, emb dim]
output, (hidden, cell) = self.rnn(embedded)
#output = [sent len, batch size, hid dim * num directions]
#hidden = [num layers * num directions, batch size, hid dim]
#cell = [num layers * num directions, batch size, hid dim]
#concat the final forward (hidden[-2,:,:]) and backward (hidden[-1,:,:]) hidden layers
#and apply dropout
hidden = self.dropout(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1))
#hidden = [batch size, hid dim * num directions]
return self.fc(hidden.squeeze(0))
class RNNNew(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, bidirectional, dropout):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.rnn = nn.LSTM(embedding_dim, hidden_dim, num_layers=n_layers,
bidirectional=bidirectional, dropout=dropout)
self.dropout = nn.Dropout(dropout)
self.directions = 2 if bidirectional else 1
self.fc = nn.Linear(hidden_dim * self.directions, output_dim)
def forward(self, x):
#x = [sent len, batch size]
embedded = self.dropout(self.embedding(x))
#embedded = [sent len, batch size, emb dim]
output, (hidden, cell) = self.rnn(embedded)
hidden = rearrange(hidden, '(layer dir) b c -> layer b (dir c)',
dir=self.directions)
# take the final layer's hidden
return self.fc(self.dropout(hidden[-1]))
- original code misbehaves for non-bidirectional models
- ... and fails when bidirectional = False, and there is only one layer
- modification of the code shows both how hidden is structured and how it is modified
Writing FastText faster
class FastTextOld(nn.Module):
def __init__(self, vocab_size, embedding_dim, output_dim):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.fc = nn.Linear(embedding_dim, output_dim)
def forward(self, x):
#x = [sent len, batch size]
embedded = self.embedding(x)
#embedded = [sent len, batch size, emb dim]
embedded = embedded.permute(1, 0, 2)
#embedded = [batch size, sent len, emb dim]
pooled = F.avg_pool2d(embedded, (embedded.shape[1], 1)).squeeze(1)
#pooled = [batch size, embedding_dim]
return self.fc(pooled)
def FastTextNew(vocab_size, embedding_dim, output_dim):
return nn.Sequential(
Rearrange('t b -> t b'),
nn.Embedding(vocab_size, embedding_dim),
Reduce('t b c -> b c', 'mean'),
nn.Linear(embedding_dim, output_dim),
Rearrange('b c -> b c'),
)
Some comments on new code:
- first and last operations do nothing and can be removed
- but were added to explicitly show expected input and output
- this also gives you a flexibility of changing interface by editing a single line. Should you need to accept inputs as (batch, time),
you just change first line to
Rearrange('b t -> t b'),
CNNs for text classification
class CNNOld(nn.Module):
def __init__(self, vocab_size, embedding_dim, n_filters, filter_sizes, output_dim, dropout):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.conv_0 = nn.Conv2d(in_channels=1, out_channels=n_filters, kernel_size=(filter_sizes[0],embedding_dim))
self.conv_1 = nn.Conv2d(in_channels=1, out_channels=n_filters, kernel_size=(filter_sizes[1],embedding_dim))
self.conv_2 = nn.Conv2d(in_channels=1, out_channels=n_filters, kernel_size=(filter_sizes[2],embedding_dim))
self.fc = nn.Linear(len(filter_sizes)*n_filters, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
#x = [sent len, batch size]
x = x.permute(1, 0)
#x = [batch size, sent len]
embedded = self.embedding(x)
#embedded = [batch size, sent len, emb dim]
embedded = embedded.unsqueeze(1)
#embedded = [batch size, 1, sent len, emb dim]
conved_0 = F.relu(self.conv_0(embedded).squeeze(3))
conved_1 = F.relu(self.conv_1(embedded).squeeze(3))
conved_2 = F.relu(self.conv_2(embedded).squeeze(3))
#conv_n = [batch size, n_filters, sent len - filter_sizes[n]]
pooled_0 = F.max_pool1d(conved_0, conved_0.shape[2]).squeeze(2)
pooled_1 = F.max_pool1d(conved_1, conved_1.shape[2]).squeeze(2)
pooled_2 = F.max_pool1d(conved_2, conved_2.shape[2]).squeeze(2)
#pooled_n = [batch size, n_filters]
cat = self.dropout(torch.cat((pooled_0, pooled_1, pooled_2), dim=1))
#cat = [batch size, n_filters * len(filter_sizes)]
return self.fc(cat)
class CNNNew(nn.Module):
def __init__(self, vocab_size, embedding_dim, n_filters, filter_sizes, output_dim, dropout):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.convs = nn.ModuleList([
nn.Conv1d(embedding_dim, n_filters, kernel_size=size) for size in filter_sizes
])
self.fc = nn.Linear(len(filter_sizes) * n_filters, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
x = rearrange(x, 't b -> t b')
emb = rearrange(self.embedding(x), 't b c -> b c t')
pooled = [reduce(conv(emb), 'b c t -> b c', 'max') for conv in self.convs]
concatenated = rearrange(pooled, 'filter b c -> b (filter c)')
return self.fc(self.dropout(F.relu(concatenated)))
- Original code misuses Conv2d, while Conv1d is the right choice
- Fixed code can work with any number of filter_sizes (and won't fail)
- First line in new code does nothing, but was added for simplicity
Highway convolutions
- Highway convolutions are common in TTS systems. Code below makes splitting a bit more explicit.
- Splitting policy may eventually turn out to be important if input had previously groups over channel axes (group convolutions or bidirectional LSTMs/GRUs)
- Same applies to GLU and gated units in general
class HighwayConv1dOld(nn.Conv1d):
def forward(self, inputs):
L = super(HighwayConv1dOld, self).forward(inputs)
H1, H2 = torch.chunk(L, 2, 1) # chunk at the feature dim
torch.sigmoid_(H1)
return H1 * H2 + (1.0 - H1) * inputs
class HighwayConv1dNew(nn.Conv1d):
def forward(self, inputs):
L = super().forward(inputs)
H1, H2 = rearrange(L, 'b (split c) t -> split b c t', split=2)
torch.sigmoid_(H1)
return H1 * H2 + (1.0 - H1) * inputs
Tacotron's CBHG module
class CBHG_Old(nn.Module):
"""CBHG module: a recurrent neural network composed of:
- 1-d convolution banks
- Highway networks + residual connections
- Bidirectional gated recurrent units
"""
def __init__(self, in_dim, K=16, projections=[128, 128]):
super(CBHG, self).__init__()
self.in_dim = in_dim
self.relu = nn.ReLU()
self.conv1d_banks = nn.ModuleList(
[BatchNormConv1d(in_dim, in_dim, kernel_size=k, stride=1,
padding=k // 2, activation=self.relu)
for k in range(1, K + 1)])
self.max_pool1d = nn.MaxPool1d(kernel_size=2, stride=1, padding=1)
in_sizes = [K * in_dim] + projections[:-1]
activations = [self.relu] * (len(projections) - 1) + [None]
self.conv1d_projections = nn.ModuleList(
[BatchNormConv1d(in_size, out_size, kernel_size=3, stride=1,
padding=1, activation=ac)
for (in_size, out_size, ac) in zip(
in_sizes, projections, activations)])
self.pre_highway = nn.Linear(projections[-1], in_dim, bias=False)
self.highways = nn.ModuleList(
[Highway(in_dim, in_dim) for _ in range(4)])
self.gru = nn.GRU(
in_dim, in_dim, 1, batch_first=True, bidirectional=True)
def forward_old(self, inputs):
# (B, T_in, in_dim)
x = inputs
# Needed to perform conv1d on time-axis
# (B, in_dim, T_in)
if x.size(-1) == self.in_dim:
x = x.transpose(1, 2)
T = x.size(-1)
# (B, in_dim*K, T_in)
# Concat conv1d bank outputs
x = torch.cat([conv1d(x)[:, :, :T] for conv1d in self.conv1d_banks], dim=1)
assert x.size(1) == self.in_dim * len(self.conv1d_banks)
x = self.max_pool1d(x)[:, :, :T]
for conv1d in self.conv1d_projections:
x = conv1d(x)
# (B, T_in, in_dim)
# Back to the original shape
x = x.transpose(1, 2)
if x.size(-1) != self.in_dim:
x = self.pre_highway(x)
# Residual connection
x += inputs
for highway in self.highways:
x = highway(x)
# (B, T_in, in_dim*2)
outputs, _ = self.gru(x)
return outputs
def forward_new(self, inputs, input_lengths=None):
x = rearrange(inputs, 'b t c -> b c t')
_, _, T = x.shape
# Concat conv1d bank outputs
x = rearrange([conv1d(x)[:, :, :T] for conv1d in self.conv1d_banks],
'bank b c t -> b (bank c) t', c=self.in_dim)
x = self.max_pool1d(x)[:, :, :T]
for conv1d in self.conv1d_projections:
x = conv1d(x)
x = rearrange(x, 'b c t -> b t c')
if x.size(-1) != self.in_dim:
x = self.pre_highway(x)
# Residual connection
x += inputs
for highway in self.highways:
x = highway(x)
# (B, T_in, in_dim*2)
outputs, _ = self.gru(self.highways(x))
return outputs
There is still a large room for improvements, but in this example only forward function was changed
Simple attention
Good news: there is no more need to guess order of dimensions. Neither for inputs nor for outputs
class Attention(nn.Module):
def __init__(self):
super(Attention, self).__init__()
def forward(self, K, V, Q):
A = torch.bmm(K.transpose(1,2), Q) / np.sqrt(Q.shape[1])
A = F.softmax(A, 1)
R = torch.bmm(V, A)
return torch.cat((R, Q), dim=1)
def attention(K, V, Q):
_, n_channels, _ = K.shape
A = torch.einsum('bct,bcl->btl', [K, Q])
A = F.softmax(A * n_channels ** (-0.5), 1)
R = torch.einsum('bct,btl->bcl', [V, A])
return torch.cat((R, Q), dim=1)
Transformer's attention needs more attention
class ScaledDotProductAttention(nn.Module):
''' Scaled Dot-Product Attention '''
def __init__(self, temperature, attn_dropout=0.1):
super().__init__()
self.temperature = temperature
self.dropout = nn.Dropout(attn_dropout)
self.softmax = nn.Softmax(dim=2)
def forward(self, q, k, v, mask=None):
attn = torch.bmm(q, k.transpose(1, 2))
attn = attn / self.temperature
if mask is not None:
attn = attn.masked_fill(mask, -np.inf)
attn = self.softmax(attn)
attn = self.dropout(attn)
output = torch.bmm(attn, v)
return output, attn
class MultiHeadAttentionOld(nn.Module):
''' Multi-Head Attention module '''
def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
super().__init__()
self.n_head = n_head
self.d_k = d_k
self.d_v = d_v
self.w_qs = nn.Linear(d_model, n_head * d_k)
self.w_ks = nn.Linear(d_model, n_head * d_k)
self.w_vs = nn.Linear(d_model, n_head * d_v)
nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v)))
self.attention = ScaledDotProductAttention(temperature=np.power(d_k, 0.5))
self.layer_norm = nn.LayerNorm(d_model)
self.fc = nn.Linear(n_head * d_v, d_model)
nn.init.xavier_normal_(self.fc.weight)
self.dropout = nn.Dropout(dropout)
def forward(self, q, k, v, mask=None):
d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
sz_b, len_q, _ = q.size()
sz_b, len_k, _ = k.size()
sz_b, len_v, _ = v.size()
residual = q
q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)
q = q.permute(2, 0, 1, 3).contiguous().view(-1, len_q, d_k) # (n*b) x lq x dk
k = k.permute(2, 0, 1, 3).contiguous().view(-1, len_k, d_k) # (n*b) x lk x dk
v = v.permute(2, 0, 1, 3).contiguous().view(-1, len_v, d_v) # (n*b) x lv x dv
mask = mask.repeat(n_head, 1, 1) # (n*b) x .. x ..
output, attn = self.attention(q, k, v, mask=mask)
output = output.view(n_head, sz_b, len_q, d_v)
output = output.permute(1, 2, 0, 3).contiguous().view(sz_b, len_q, -1) # b x lq x (n*dv)
output = self.dropout(self.fc(output))
output = self.layer_norm(output + residual)
return output, attn
class MultiHeadAttentionNew(nn.Module):
def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
super().__init__()
self.n_head = n_head
self.w_qs = nn.Linear(d_model, n_head * d_k)
self.w_ks = nn.Linear(d_model, n_head * d_k)
self.w_vs = nn.Linear(d_model, n_head * d_v)
nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v)))
self.fc = nn.Linear(n_head * d_v, d_model)
nn.init.xavier_normal_(self.fc.weight)
self.dropout = nn.Dropout(p=dropout)
self.layer_norm = nn.LayerNorm(d_model)
def forward(self, q, k, v, mask=None):
residual = q
q = rearrange(self.w_qs(q), 'b l (head k) -> head b l k', head=self.n_head)
k = rearrange(self.w_ks(k), 'b t (head k) -> head b t k', head=self.n_head)
v = rearrange(self.w_vs(v), 'b t (head v) -> head b t v', head=self.n_head)
attn = torch.einsum('hblk,hbtk->hblt', [q, k]) / np.sqrt(q.shape[-1])
if mask is not None:
attn = attn.masked_fill(mask[None], -np.inf)
attn = torch.softmax(attn, dim=3)
output = torch.einsum('hblt,hbtv->hblv', [attn, v])
output = rearrange(output, 'head b l v -> b l (head v)')
output = self.dropout(self.fc(output))
output = self.layer_norm(output + residual)
return output, attn
Benefits of new implementation
- we have one module, not two
- now code does not fail for None mask
- the amount of caveats in the original code that we removed is huge. Try erasing comments and deciphering what happens there
Self-attention GANs
SAGANs are currently SotA for image generation, and can be simplified using same tricks.
class Self_Attn_Old(nn.Module):
""" Self attention Layer"""
def __init__(self,in_dim,activation):
super(Self_Attn_Old,self).__init__()
self.chanel_in = in_dim
self.activation = activation
self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)
self.gamma = nn.Parameter(torch.zeros(1))
self.softmax = nn.Softmax(dim=-1) #
def forward(self, x):
"""
inputs :
x : input feature maps( B X C X W X H)
returns :
out : self attention value + input feature
attention: B X N X N (N is Width*Height)
"""
m_batchsize,C,width ,height = x.size()
proj_query = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)
proj_key = self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)
energy = torch.bmm(proj_query,proj_key) # transpose check
attention = self.softmax(energy) # BX (N) X (N)
proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N
out = torch.bmm(proj_value,attention.permute(0,2,1) )
out = out.view(m_batchsize,C,width,height)
out = self.gamma*out + x
return out,attention
class Self_Attn_New(nn.Module):
""" Self attention Layer"""
def __init__(self, in_dim):
super().__init__()
self.query_conv = nn.Conv2d(in_dim, out_channels=in_dim//8, kernel_size=1)
self.key_conv = nn.Conv2d(in_dim, out_channels=in_dim//8, kernel_size=1)
self.value_conv = nn.Conv2d(in_dim, out_channels=in_dim, kernel_size=1)
self.gamma = nn.Parameter(torch.zeros([1]))
def forward(self, x):
proj_query = rearrange(self.query_conv(x), 'b c h w -> b (h w) c')
proj_key = rearrange(self.key_conv(x), 'b c h w -> b c (h w)')
proj_value = rearrange(self.value_conv(x), 'b c h w -> b (h w) c')
energy = torch.bmm(proj_query, proj_key)
attention = F.softmax(energy, dim=2)
out = torch.bmm(attention, proj_value)
out = x + self.gamma * rearrange(out, 'b (h w) c -> b c h w',
**parse_shape(x, 'b c h w'))
return out, attention
Improving time sequence prediction
While this example was considered to be simplistic, I had to analyze surrounding code to understand what kind of input was expected. You can try yourself.
Additionally now the code works with any dtype, not only double; and new code supports using GPU.
class SequencePredictionOld(nn.Module):
def __init__(self):
super(SequencePredictionOld, self).__init__()
self.lstm1 = nn.LSTMCell(1, 51)
self.lstm2 = nn.LSTMCell(51, 51)
self.linear = nn.Linear(51, 1)
def forward(self, input, future = 0):
outputs = []
h_t = torch.zeros(input.size(0), 51, dtype=torch.double)
c_t = torch.zeros(input.size(0), 51, dtype=torch.double)
h_t2 = torch.zeros(input.size(0), 51, dtype=torch.double)
c_t2 = torch.zeros(input.size(0), 51, dtype=torch.double)
for i, input_t in enumerate(input.chunk(input.size(1), dim=1)):
h_t, c_t = self.lstm1(input_t, (h_t, c_t))
h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))
output = self.linear(h_t2)
outputs += [output]
for i in range(future):# if we should predict the future
h_t, c_t = self.lstm1(output, (h_t, c_t))
h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))
output = self.linear(h_t2)
outputs += [output]
outputs = torch.stack(outputs, 1).squeeze(2)
return outputs
class SequencePredictionNew(nn.Module):
def __init__(self):
super(SequencePredictionNew, self).__init__()
self.lstm1 = nn.LSTMCell(1, 51)
self.lstm2 = nn.LSTMCell(51, 51)
self.linear = nn.Linear(51, 1)
def forward(self, input, future=0):
b, t = input.shape
h_t, c_t, h_t2, c_t2 = torch.zeros(4, b, 51, dtype=self.linear.weight.dtype,
device=self.linear.weight.device)
outputs = []
for input_t in rearrange(input, 'b t -> t b ()'):
h_t, c_t = self.lstm1(input_t, (h_t, c_t))
h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))
output = self.linear(h_t2)
outputs += [output]
for i in range(future): # if we should predict the future
h_t, c_t = self.lstm1(output, (h_t, c_t))
h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))
output = self.linear(h_t2)
outputs += [output]
return rearrange(outputs, 't b () -> b t')
Transforming spacial transformer network (STN)
class SpacialTransformOld(nn.Module):
def __init__(self):
super(Net, self).__init__()
# Spatial transformer localization-network
self.localization = nn.Sequential(
nn.Conv2d(1, 8, kernel_size=7),
nn.MaxPool2d(2, stride=2),
nn.ReLU(True),
nn.Conv2d(8, 10, kernel_size=5),
nn.MaxPool2d(2, stride=2),
nn.ReLU(True)
)
# Regressor for the 3 * 2 affine matrix
self.fc_loc = nn.Sequential(
nn.Linear(10 * 3 * 3, 32),
nn.ReLU(True),
nn.Linear(32, 3 * 2)
)
# Initialize the weights/bias with identity transformation
self.fc_loc[2].weight.data.zero_()
self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))
# Spatial transformer network forward function
def stn(self, x):
xs = self.localization(x)
xs = xs.view(-1, 10 * 3 * 3)
theta = self.fc_loc(xs)
theta = theta.view(-1, 2, 3)
grid = F.affine_grid(theta, x.size())
x = F.grid_sample(x, grid)
return x
class SpacialTransformNew(nn.Module):
def __init__(self):
super(Net, self).__init__()
# Spatial transformer localization-network
linear = nn.Linear(32, 3 * 2)
# Initialize the weights/bias with identity transformation
linear.weight.data.zero_()
linear.bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))
self.compute_theta = nn.Sequential(
nn.Conv2d(1, 8, kernel_size=7),
nn.MaxPool2d(2, stride=2),
nn.ReLU(True),
nn.Conv2d(8, 10, kernel_size=5),
nn.MaxPool2d(2, stride=2),
nn.ReLU(True),
Rearrange('b c h w -> b (c h w)', h=3, w=3),
nn.Linear(10 * 3 * 3, 32),
nn.ReLU(True),
linear,
Rearrange('b (row col) -> b row col', row=2, col=3),
)
# Spatial transformer network forward function
def stn(self, x):
grid = F.affine_grid(self.compute_theta(x), x.size())
return F.grid_sample(x, grid)
- new code will give reasonable errors when passed image size is different from expected
- if batch size is divisible by 18, whatever you input in the old code, it'll fail no sooner than affine_grid.
Improving GLOW
That's a good old depth-to-space written manually!
Since GLOW is revertible, it will frequently rely on rearrange-like operations.
def unsqueeze2d_old(input, factor=2):
assert factor >= 1 and isinstance(factor, int)
factor2 = factor ** 2
if factor == 1:
return input
size = input.size()
B = size[0]
C = size[1]
H = size[2]
W = size[3]
assert C % (factor2) == 0, "{}".format(C)
x = input.view(B, C // factor2, factor, factor, H, W)
x = x.permute(0, 1, 4, 2, 5, 3).contiguous()
x = x.view(B, C // (factor2), H * factor, W * factor)
return x
def squeeze2d_old(input, factor=2):
assert factor >= 1 and isinstance(factor, int)
if factor == 1:
return input
size = input.size()
B = size[0]
C = size[1]
H = size[2]
W = size[3]
assert H % factor == 0 and W % factor == 0, "{}".format((H, W))
x = input.view(B, C, H // factor, factor, W // factor, factor)
x = x.permute(0, 1, 3, 5, 2, 4).contiguous()
x = x.view(B, C * factor * factor, H // factor, W // factor)
return x
def unsqueeze2d_new(input, factor=2):
return rearrange(input, 'b (c h2 w2) h w -> b c (h h2) (w w2)', h2=factor, w2=factor)
def squeeze2d_new(input, factor=2):
return rearrange(input, 'b c (h h2) (w w2) -> b (c h2 w2) h w', h2=factor, w2=factor)
- term
squeezeisn't very helpful: which dimension is squeezed? There istorch.squeeze, but it's very different. - in fact, we could skip creating functions completely - it is a single call to
einopsanyway
Detecting problems in YOLO detection
def YOLO_prediction_old(input, num_classes, num_anchors, anchors, stride_h, stride_w):
bs = input.size(0)
in_h = input.size(2)
in_w = input.size(3)
scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors]
prediction = input.view(bs, num_anchors,
5 + num_classes, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()
# Get outputs
x = torch.sigmoid(prediction[..., 0]) # Center x
y = torch.sigmoid(prediction[..., 1]) # Center y
w = prediction[..., 2] # Width
h = prediction[..., 3] # Height
conf = torch.sigmoid(prediction[..., 4]) # Conf
pred_cls = torch.sigmoid(prediction[..., 5:]) # Cls pred.
FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
# Calculate offsets for each grid
grid_x = torch.linspace(0, in_w - 1, in_w).repeat(in_w, 1).repeat(
bs * num_anchors, 1, 1).view(x.shape).type(FloatTensor)
grid_y = torch.linspace(0, in_h - 1, in_h).repeat(in_h, 1).t().repeat(
bs * num_anchors, 1, 1).view(y.shape).type(FloatTensor)
# Calculate anchor w, h
anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape)
anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape)
# Add offset and scale with anchors
pred_boxes = FloatTensor(prediction[..., :4].shape)
pred_boxes[..., 0] = x.data + grid_x
pred_boxes[..., 1] = y.data + grid_y
pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
pred_boxes[..., 3] = torch.exp(h.data) * anchor_h
# Results
_scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)
output = torch.cat((pred_boxes.view(bs, -1, 4) * _scale,
conf.view(bs, -1, 1), pred_cls.view(bs, -1, num_classes)), -1)
return output
def YOLO_prediction_new(input, num_classes, num_anchors, anchors, stride_h, stride_w):
raw_predictions = rearrange(input, 'b (anchor prediction) h w -> prediction b anchor h w',
anchor=num_anchors, prediction=5 + num_classes)
anchors = torch.FloatTensor(anchors).to(input.device)
anchor_sizes = rearrange(anchors, 'anchor dim -> dim () anchor () ()')
_, _, _, in_h, in_w = raw_predictions.shape
grid_h = rearrange(torch.arange(in_h).float(), 'h -> () () h ()').to(input.device)
grid_w = rearrange(torch.arange(in_w).float(), 'w -> () () () w').to(input.device)
predicted_bboxes = torch.zeros_like(raw_predictions)
predicted_bboxes[0] = (raw_predictions[0].sigmoid() + grid_w) * stride_w # center x
predicted_bboxes[1] = (raw_predictions[1].sigmoid() + grid_h) * stride_h # center y
predicted_bboxes[2:4] = (raw_predictions[2:4].exp()) * anchor_sizes # bbox width and height
predicted_bboxes[4] = raw_predictions[4].sigmoid() # confidence
predicted_bboxes[5:] = raw_predictions[5:].sigmoid() # class predictions
# merging all predicted bboxes for each image
return rearrange(predicted_bboxes, 'prediction b anchor h w -> b (anchor h w) prediction')
We changed and fixed a lot:
- new code won't fail if input is not on the first GPU
- old code has wrong grid_x and grid_y for non-square images
- new code doesn't use replication when broadcasting is sufficient
- old code strangely sometimes takes
.data, but this has no real effect, as some branches preserve gradient till the end- if gradients not needed, torch.no_grad should be used, so it's redundant
Simpler output for a bunch of pictures
Next time you need to output drawings of you generative models, you can use this trick
device = 'cpu'
plt.imshow(np.transpose(vutils.make_grid(fake_batch.to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
padded = F.pad(fake_batch[:64], [1, 1, 1, 1])
plt.imshow(rearrange(padded, '(b1 b2) c h w -> (b1 h) (b2 w) c', b1=8).cpu())
Instead of conclusion
Better code is a vague term; to be specific, code is expected to be:
- reliable: does what expected and does not fail. Explicitly fails for wrong inputs
- maintainable and modifiable
- reusable: understanding and modifying code should be easier than writing from scratch
- fast: in my measurements, proposed versions have speed similar to the original code
- readability counts, as a mean to achieve previous goals
Provided examples show how to improve these criteria for deep learning code. And einops helps a lot.
Links
Answer is: {x} (hover to see)
".format(x=tuple(x)), raw=True, ) arogozhnikov-einops-ad2c8d6/docs_src/ 0000775 0000000 0000000 00000000000 14752016746 0020035 5 ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/docs_src/1-einops-basics.ipynb 0000777 0000000 0000000 00000000000 14752016746 0031073 2../docs/1-einops-basics.ipynb ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/docs_src/2-einops-for-deep-learning.ipynb 0000777 0000000 0000000 00000000000 14752016746 0035161 2../docs/2-einops-for-deep-learning.ipynb ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/docs_src/3-einmix-layer.ipynb 0000777 0000000 0000000 00000000000 14752016746 0030607 2../docs/3-einmix-layer.ipynb ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/docs_src/4-pack-and-unpack.ipynb 0000777 0000000 0000000 00000000000 14752016746 0031475 2../docs/4-pack-and-unpack.ipynb ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/docs_src/CNAME 0000664 0000000 0000000 00000000014 14752016746 0020576 0 ustar 00root root 0000000 0000000 einops.rocks arogozhnikov-einops-ad2c8d6/docs_src/api/ 0000775 0000000 0000000 00000000000 14752016746 0020606 5 ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/docs_src/api/asnumpy.md 0000664 0000000 0000000 00000000044 14752016746 0022622 0 ustar 00root root 0000000 0000000 # einops.asnumpy ::: einops.asnumpy arogozhnikov-einops-ad2c8d6/docs_src/api/einsum.md 0000664 0000000 0000000 00000000042 14752016746 0022424 0 ustar 00root root 0000000 0000000 # einops.einsum ::: einops.einsum arogozhnikov-einops-ad2c8d6/docs_src/api/pack_unpack.md 0000664 0000000 0000000 00000000155 14752016746 0023410 0 ustar 00root root 0000000 0000000 # einops.pack and einops.unpack ## einops.pack ::: einops.pack## einops.unpack ::: einops.unpack arogozhnikov-einops-ad2c8d6/docs_src/api/parse_shape.md 0000664 0000000 0000000 00000000054 14752016746 0023421 0 ustar 00root root 0000000 0000000 # einops.parse_shape ::: einops.parse_shape arogozhnikov-einops-ad2c8d6/docs_src/api/rearrange.md 0000664 0000000 0000000 00000000050 14752016746 0023071 0 ustar 00root root 0000000 0000000 # einops.rearrange ::: einops.rearrange arogozhnikov-einops-ad2c8d6/docs_src/api/reduce.md 0000664 0000000 0000000 00000000042 14752016746 0022373 0 ustar 00root root 0000000 0000000 # einops.reduce ::: einops.reduce arogozhnikov-einops-ad2c8d6/docs_src/api/repeat.md 0000664 0000000 0000000 00000000042 14752016746 0022404 0 ustar 00root root 0000000 0000000 # einops.repeat ::: einops.repeat arogozhnikov-einops-ad2c8d6/docs_src/css/ 0000775 0000000 0000000 00000000000 14752016746 0020625 5 ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/docs_src/css/codehilite.css 0000664 0000000 0000000 00000010161 14752016746 0023447 0 ustar 00root root 0000000 0000000 /*, .codehilite pre*/ .codehilite code{ color:#3F3F3F; background-color:#F7F7F7; overflow-x: scroll; box-sizing: border-box; padding: 0.01em; border: 1px solid #CCC; } .codehilite .hll { background-color: #ffffcc } .codehilite .c { color: #999988; font-style: italic } /* Comment */ .codehilite .err { color: #a61717; background-color: #e3d2d2 } /* Error */ .codehilite .k { color: #000000; font-weight: bold } /* Keyword */ .codehilite .o { color: #000000; font-weight: bold } /* Operator */ .codehilite .cm { color: #999988; font-style: italic } /* Comment.Multiline */ .codehilite .cp { color: #999999; font-weight: bold; font-style: italic } /* Comment.Preproc */ .codehilite .c1 { color: #999988; font-style: italic } /* Comment.Single */ .codehilite .cs { color: #999999; font-weight: bold; font-style: italic } /* Comment.Special */ .codehilite .gd { color: #000000; background-color: #ffdddd } /* Generic.Deleted */ .codehilite .ge { color: #000000; font-style: italic } /* Generic.Emph */ .codehilite .gr { color: #aa0000 } /* Generic.Error */ .codehilite .gh { color: #999999 } /* Generic.Heading */ .codehilite .gi { color: #000000; background-color: #ddffdd } /* Generic.Inserted */ .codehilite .go { color: #888888 } /* Generic.Output */ .codehilite .gp { color: #555555 } /* Generic.Prompt */ .codehilite .gs { font-weight: bold } /* Generic.Strong */ .codehilite .gu { color: #aaaaaa } /* Generic.Subheading */ .codehilite .gt { color: #aa0000 } /* Generic.Traceback */ .codehilite .kc { color: #000000; font-weight: bold } /* Keyword.Constant */ .codehilite .kd { color: #000000; font-weight: bold } /* Keyword.Declaration */ .codehilite .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */ .codehilite .kp { color: #000000; font-weight: bold } /* Keyword.Pseudo */ .codehilite .kr { color: #000000; font-weight: bold } /* Keyword.Reserved */ .codehilite .kt { color: #445588; font-weight: bold } /* Keyword.Type */ .codehilite .m { color: #009999 } /* Literal.Number */ .codehilite .s { color: #d01040 } /* Literal.String */ .codehilite .na { color: #008080 } /* Name.Attribute */ .codehilite .nb { color: #0086B3 } /* Name.Builtin */ .codehilite .nc { color: #445588; font-weight: bold } /* Name.Class */ .codehilite .no { color: #008080 } /* Name.Constant */ .codehilite .nd { color: #3c5d5d; font-weight: bold } /* Name.Decorator */ .codehilite .ni { color: #800080 } /* Name.Entity */ .codehilite .ne { color: #990000; font-weight: bold } /* Name.Exception */ .codehilite .nf { color: #990000; font-weight: bold } /* Name.Function */ .codehilite .nl { color: #990000; font-weight: bold } /* Name.Label */ .codehilite .nn { color: #555555 } /* Name.Namespace */ .codehilite .nt { color: #000080 } /* Name.Tag */ .codehilite .nv { color: #008080 } /* Name.Variable */ .codehilite .ow { color: #000000; font-weight: bold } /* Operator.Word */ .codehilite .w { color: #bbbbbb } /* Text.Whitespace */ .codehilite .mf { color: #009999 } /* Literal.Number.Float */ .codehilite .mh { color: #009999 } /* Literal.Number.Hex */ .codehilite .mi { color: #009999 } /* Literal.Number.Integer */ .codehilite .mo { color: #009999 } /* Literal.Number.Oct */ .codehilite .sb { color: #d01040 } /* Literal.String.Backtick */ .codehilite .sc { color: #d01040 } /* Literal.String.Char */ .codehilite .sd { color: #d01040 } /* Literal.String.Doc */ .codehilite .s2 { color: #d01040 } /* Literal.String.Double */ .codehilite .se { color: #d01040 } /* Literal.String.Escape */ .codehilite .sh { color: #d01040 } /* Literal.String.Heredoc */ .codehilite .si { color: #d01040 } /* Literal.String.Interpol */ .codehilite .sx { color: #d01040 } /* Literal.String.Other */ .codehilite .sr { color: #009926 } /* Literal.String.Regex */ .codehilite .s1 { color: #d01040 } /* Literal.String.Single */ .codehilite .ss { color: #990073 } /* Literal.String.Symbol */ .codehilite .bp { color: #999999 } /* Name.Builtin.Pseudo */ .codehilite .vc { color: #008080 } /* Name.Variable.Class */ .codehilite .vg { color: #008080 } /* Name.Variable.Global */ .codehilite .vi { color: #008080 } /* Name.Variable.Instance */ .codehilite .il { color: #009999 } /* Literal.Number.Integer.Long */ arogozhnikov-einops-ad2c8d6/docs_src/css/mkdocs.css 0000664 0000000 0000000 00000001261 14752016746 0022617 0 ustar 00root root 0000000 0000000 div.highlight>pre { padding: 3px; } .md-typeset ul li, .md-typeset ol li { /*original padding makes lists huge*/ margin-bottom: 0.2em; } html { /* override 125% default, which looks terrible on smaller screens */ font-size: 100% !important; } .jupyter-wrapper { /* set css var used by markdown */ --jp-code-font-size: 90% !important; } .md-grid { /* change max width of page, looks much better on bigger screens */ max-width: 72rem !important; } div.highlight { /* this was introduced to fix size of documentation script. This applies to multiline code fragments in markdown */ font-size: 13px; background-color: #f2f2f2; } arogozhnikov-einops-ad2c8d6/docs_src/docs 0000777 0000000 0000000 00000000000 14752016746 0022050 2../docs/ ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/docs_src/pages/ 0000775 0000000 0000000 00000000000 14752016746 0021134 5 ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/docs_src/pages/projects.md 0000664 0000000 0000000 00000005555 14752016746 0023321 0 ustar 00root root 0000000 0000000 Einops tutorials cover multiple einops usages (and you'd better first follow tutorials), but it can also help to see einops in action. ## Selected projects Here are some open-source projects that can teach how to leverage einops for your problems - [@lucidrains](https://github.com/lucidrains) has a dramatic [collection of vision transformers](https://github.com/lucidrains/vit-pytorch) - there is a plenty of good examples how to use einops efficiently in your projects - lambda networks (non-conventional architecture) implemented by @lucidrains - nice demonstration how clearer code can be with einops, even compared to description in the paper - [implementation](https://github.com/lucidrains/lambda-networks) and [video](https://www.youtube.com/watch?v=3qxJ2WD8p4w) - capsule networks (aka capsnets) [implemented in einops](https://github.com/arogozhnikov/readable_capsnet) - blazingly fast, concise (3-10 times less code), and memory efficient capsule networks, written with einops - [NuX](https://github.com/Information-Fusion-Lab-Umass/NuX) — normalizing flows in Jax - different rearrangement patterns in normalizing flows have nice mapping to einops - For video recognition, look at [MotionFormer](https://github.com/facebookresearch/Motionformer) and [TimeSFormer](https://github.com/lucidrains/TimeSformer-pytorch) implementations - For protein folding, see [implementation](https://github.com/lucidrains/invariant-point-attention) of invariant point attention from AlphaFold 2 ## Community introductions to einops Tutorial in the AI summer about einops and einsum:
> You really have to try it out, you'll love it: https://github.com/arogozhnikov/einops
> Plus it supports NumPy, TensorFlow, PyTorch, and more. Aurélien Geron,
author of "Hands-On Machine Learning with Scikit-Learn and TensorFlow."
Former PM of YouTube video classification. [(ref)](https://twitter.com/aureliengeron/status/1382829421967515648) > This is your daily reminder to *never* trust pytorch's native reshape/view.
> Always use einops! > > (Just spend 1 h debugging code and it turned out tensor.view was shuffling the tensor in a weird way) Tom Lieberum, University of Amsterdam [(ref)](https://twitter.com/lieberum_t/status/1427282842250358787) > TIL einops can be faster than raw PyTorch 🤯 Zach Mueller, Novetta, author of "walk with fastai" [(ref)](https://twitter.com/TheZachMueller/status/1418003372494426113) > einops are also a joy to work with! Norman Casagrande, Deepmind [(ref)](https://twitter.com/nova77t/status/1419405150805008387) > — And btw I estimate that AI research suffers globally from a 5% loss of productivity because einops are not included in @PyTorch by default. > > — Might be true for research, but not likely to be true for engineering. I don't think a lot of people in the industry use PyTorch directly [...] > > — That’s why it’s 5% and not 30% > > — E-xa-ctly [Discussion thread](https://twitter.com/francoisfleuret/status/1409141186326106114) > After a while, it feels like einsum+einops is all you need ;) [...] Tim Rocktäschel, Facebook AI Research [(ref)](https://twitter.com/_rockt/status/1390049226193788930) > Yes, I am also using einops in that snippet! It's great!
> I wished I knew about it from the start when it was created Tim Dettmers, PhD Student at UoW and visiting researcher at Facebook AI [(ref)](https://twitter.com/Tim_Dettmers/status/1390027329351520256) > A little late to the party, but einops (https://github.com/arogozhnikov/einops) is a massive improvement to deep learning code readability. I love this! Daniel Havir, Apple [(ref)](https://twitter.com/danielhavir/status/1389070232853966849) > I recently discovered the beauty of torch.einsum and einops.rearrange
> and at this point I'm confused why I even bothered with other tensor operations in the first place. Robin M. Schmidt, AI/ML Resident at Apple [(ref)](https://twitter.com/robinschmidt_/status/1363709832788852736) > I end up using einops for ridiculously simple things,
> simply to be nice to my future self (because the code takes so little effort to read). Nasim Rahaman, MILA [(ref)](https://twitter.com/nasim_rahaman/status/1390027557546901504) > I love Einops for this kind of stuff, it makes code very readable,
> even if you are just doing a simple squeeze or expand_dims. [...] Cristian Garcia, ML Engineer @quansightai, [(ref)](https://twitter.com/cgarciae88/status/1331968395110211586) > i might be late to the party, but einsum and the einops package are unbelievably useful Samson Koelle, Stat PhD candidate at UoW, [(ref)](https://twitter.com/SeattleStatSam/status/1338673646898794496) > I really recommend einops for tensor shape manipulation Alex Mordvintsev, DeepDream creator, [(ref)](https://twitter.com/zzznah/status/1315297985585123328) > The einops Python package is worth checking out: > it provides a powerful declarative interface > for manipulating and reshaping multi-dimensional arrays https://github.com/arogozhnikov/einops Jake VanderPlas,
Google Research, core contributor to Altair, AstroML, scikit-learn, etc. [(ref)](https://twitter.com/jakevdp/status/1299012119761833989) > I can't believe after this many years of programming with NumPy/PyTorch/TensorFlow, I didn't know about ðšŽðš’ðš—ðšœðšžðš–. [...]
> — Arash Vahdat > > They smell like regexp to me -- concise, but I know it is going to take effort to understand or modify them in the future.
> — John Carmack > > I have a much easier time to read einsum than any equivalent combinations of matmul, reshape, broadcasting... you name it. > Regexps are ad-hoc, subtle and cryptic. > > Einstein summation is uniform, succinct with simple, clear semantics.
> Ein sum to rule them all ;)
> — Christian Szegedy > > The einops library, in particular, is magical. Best thing since baked bread and brewed beer.
> — @aertherks [Discussion thread](https://twitter.com/aertherks/status/1357054506656165889) arogozhnikov-einops-ad2c8d6/docs_src/pytorch-examples.html 0000777 0000000 0000000 00000000000 14752016746 0031557 2../docs/pytorch-examples.html ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/einops/ 0000775 0000000 0000000 00000000000 14752016746 0017533 5 ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/einops/__init__.py 0000664 0000000 0000000 00000000646 14752016746 0021652 0 ustar 00root root 0000000 0000000 # imports can use EinopsError class # ruff: noqa: E402 __author__ = "Alex Rogozhnikov" __version__ = "0.8.1" class EinopsError(RuntimeError): """Runtime error thrown by einops""" pass __all__ = ["rearrange", "reduce", "repeat", "einsum", "pack", "unpack", "parse_shape", "asnumpy", "EinopsError"] from .einops import rearrange, reduce, repeat, einsum, parse_shape, asnumpy from .packing import pack, unpack arogozhnikov-einops-ad2c8d6/einops/_backends.py 0000664 0000000 0000000 00000051441 14752016746 0022023 0 ustar 00root root 0000000 0000000 """ Backends in `einops` are organized to meet the following requirements - backends are not imported unless those are actually needed, because - backends may not be installed - importing all available backends will drive to significant memory footprint - backends may be present but installed with errors (but never used), importing may drive to crashes - backend should be either symbolic or imperative - this determines which methods (from_numpy/to_numpy or create_symbol/eval_symbol) should be defined - if backend can't provide symbols for shape dimensions, UnknownSize objects are used """ import sys __author__ = "Alex Rogozhnikov" _loaded_backends: dict = {} _type2backend: dict = {} _debug_importing = False def get_backend(tensor) -> "AbstractBackend": """ Takes a correct backend (e.g. numpy backend if tensor is numpy.ndarray) for a tensor. If needed, imports package and creates backend """ _type = type(tensor) _result = _type2backend.get(_type, None) if _result is not None: return _result for framework_name, backend in list(_loaded_backends.items()): if backend.is_appropriate_type(tensor): _type2backend[_type] = backend return backend # Find backend subclasses recursively backend_subclasses = [] backends = AbstractBackend.__subclasses__() while backends: backend = backends.pop() backends += backend.__subclasses__() backend_subclasses.append(backend) for BackendSubclass in backend_subclasses: if _debug_importing: print("Testing for subclass of ", BackendSubclass) if BackendSubclass.framework_name not in _loaded_backends: # check that module was already imported. Otherwise it can't be imported if BackendSubclass.framework_name in sys.modules: if _debug_importing: print("Imported backend for ", BackendSubclass.framework_name) backend = BackendSubclass() _loaded_backends[backend.framework_name] = backend if backend.is_appropriate_type(tensor): _type2backend[_type] = backend return backend raise RuntimeError("Tensor type unknown to einops {}".format(type(tensor))) class AbstractBackend: """Base backend class, major part of methods are only for debugging purposes.""" framework_name: str def is_appropriate_type(self, tensor): """helper method should recognize tensors it can handle""" raise NotImplementedError() def from_numpy(self, x): raise NotImplementedError("framework doesn't support imperative execution") def to_numpy(self, x): raise NotImplementedError("framework doesn't support imperative execution") def create_symbol(self, shape): raise NotImplementedError("framework doesn't support symbolic computations") def eval_symbol(self, symbol, symbol_value_pairs): # symbol-value pairs is list[tuple[symbol, value-tensor]] raise NotImplementedError("framework doesn't support symbolic computations") def arange(self, start, stop): # supplementary method used only in testing, so should implement CPU version raise NotImplementedError("framework doesn't implement arange") def shape(self, x): """shape should return a tuple with integers or "shape symbols" (which will evaluate to actual size)""" return x.shape def reshape(self, x, shape): return x.reshape(shape) def transpose(self, x, axes): return x.transpose(axes) def reduce(self, x, operation, axes): return getattr(x, operation)(axis=axes) def stack_on_zeroth_dimension(self, tensors: list): raise NotImplementedError() def add_axis(self, x, new_position): raise NotImplementedError() def add_axes(self, x, n_axes, pos2len): repeats = [1] * n_axes for axis_position, axis_length in pos2len.items(): x = self.add_axis(x, axis_position) repeats[axis_position] = axis_length return self.tile(x, tuple(repeats)) def tile(self, x, repeats): """repeats - same lengths as x.shape""" raise NotImplementedError() def concat(self, tensors, axis: int): """concatenates tensors along axis. Assume identical across tensors: devices, dtypes and shapes except selected axis.""" raise NotImplementedError() def is_float_type(self, x): # some backends (torch) can't compute average for non-floating types. # Decided to drop average for all backends if type is not floating raise NotImplementedError() def layers(self): raise NotImplementedError("backend does not provide layers") def __repr__(self): return "
\n\n
# """.strip() else: # other lines are not touched return line new_content = "\n".join([replace_with_video_tag(line) for line in original_text.splitlines()]) # save contents docs_index = Path(__file__).parent.parent.joinpath("docs_src", "index.md") assert docs_index.parent.exists() docs_index.write_bytes(new_content.encode("utf-8")) print("Converted README.md") arogozhnikov-einops-ad2c8d6/scripts/pytorch_examples_source/ 0000775 0000000 0000000 00000000000 14752016746 0024673 5 ustar 00root root 0000000 0000000 arogozhnikov-einops-ad2c8d6/scripts/pytorch_examples_source/Pytorch.ipynb 0000664 0000000 0000000 00000276444 14752016746 0027407 0 ustar 00root root 0000000 0000000 { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "
\n",
" \n",
" ![\"einops](\"http://arogozhnikov.github.io/images/einops/einops_logo_350x350.png\")
\n",
" \n",
" \n",
"
\n",
"\n",
"# Writing a better code with pytorch and einops\n",
"\n",
"\n",
"\n",
" \n",
" \n",
"\n", "\n", "## Rewriting building blocks of deep learning \n", "\n", "Now let's get to examples from real world.\n", "These code fragments taken from official tutorials and popular repositories.\n", "\n", "Learn how to improve code and how `einops` can help you.\n", "\n", "**Left**: as it was, **Right**: improved version" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "#right\n", "# start from importing some stuff\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import numpy as np\n", "import math\n", "\n", "from einops import rearrange, reduce, asnumpy, parse_shape\n", "from einops.layers.torch import Rearrange, Reduce" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "def initialize(model):\n", " for p in model.parameters():\n", " p.data[:] = torch.from_numpy(np.random.RandomState(sum(p.shape)).randn(*p.shape))\n", " return model" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Simple ConvNet" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "#left\n", "class Net(nn.Module):\n", " def __init__(self):\n", " super(Net, self).__init__()\n", " self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n", " self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n", " self.conv2_drop = nn.Dropout2d()\n", " self.fc1 = nn.Linear(320, 50)\n", " self.fc2 = nn.Linear(50, 10)\n", "\n", " def forward(self, x):\n", " x = F.relu(F.max_pool2d(self.conv1(x), 2))\n", " x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n", " x = x.view(-1, 320)\n", " x = F.relu(self.fc1(x))\n", " x = F.dropout(x, training=self.training)\n", " x = self.fc2(x)\n", " return F.log_softmax(x, dim=1)\n", "\n", "conv_net_old = Net()" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "#right\n", "conv_net_new = nn.Sequential(\n", " nn.Conv2d(1, 10, kernel_size=5),\n", " nn.MaxPool2d(kernel_size=2),\n", " nn.ReLU(),\n", " nn.Conv2d(10, 20, kernel_size=5),\n", " nn.MaxPool2d(kernel_size=2),\n", " nn.ReLU(),\n", " nn.Dropout2d(),\n", " Rearrange('b c h w -> b (c h w)'),\n", " nn.Linear(320, 50),\n", " nn.ReLU(),\n", " nn.Dropout(),\n", " nn.Linear(50, 10),\n", " nn.LogSoftmax(dim=1)\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Reasons to prefer new implementation:\n", "\n", "- in the original code (to the left) if input size is changed and batch size is divisible by 16 (that's usually so), we'll get something senseless after reshaping\n", " - new code will explicitly raise an error in this case\n", "- we won't forget to use dropout with flag self.training with new version\n", "- code is straightforward to read and analyze\n", "- sequential makes printing / saving / passing trivial. And there is no need in your code to load a model (which also has a number of benefits)\n", "- don't need logsoftmax? Now you can use `conv_net_new[:-1]`. One more reason to prefer `nn.Sequential`\n", "- ... and we could also add inplace for ReLU\n" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "torch.Size([4, 10])" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "conv_net_old(torch.zeros([16, 1, 20, 20])).shape\n", "# conv_net_new(torch.zeros([16, 1, 20, 20])).shape" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Super-resolution\n", "\n", "" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "#left\n", "class SuperResolutionNetOld(nn.Module):\n", " def __init__(self, upscale_factor):\n", " super(SuperResolutionNetOld, self).__init__()\n", "\n", " self.relu = nn.ReLU()\n", " self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))\n", " self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))\n", " self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))\n", " self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))\n", " self.pixel_shuffle = nn.PixelShuffle(upscale_factor)\n", "\n", " def forward(self, x):\n", " x = self.relu(self.conv1(x))\n", " x = self.relu(self.conv2(x))\n", " x = self.relu(self.conv3(x))\n", " x = self.pixel_shuffle(self.conv4(x))\n", " return x" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "#right\n", "def SuperResolutionNetNew(upscale_factor):\n", " return nn.Sequential(\n", " nn.Conv2d(1, 64, kernel_size=5, padding=2),\n", " nn.ReLU(inplace=True),\n", " nn.Conv2d(64, 64, kernel_size=3, padding=1),\n", " nn.ReLU(inplace=True),\n", " nn.Conv2d(64, 32, kernel_size=3, padding=1),\n", " nn.ReLU(inplace=True),\n", " nn.Conv2d(32, upscale_factor ** 2, kernel_size=3, padding=1),\n", " Rearrange('b (h2 w2) h w -> b (h h2) (w w2)', h2=upscale_factor, w2=upscale_factor),\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Here is the difference:\n", "\n", "- no need in special instruction pixel_shuffle (and result is transferrable between frameworks)\n", "- output doesn't contain a fake axis (and we could do the same for the input)\n", "- inplace ReLU used now, for high resolution pictures that becomes critical and saves us much memory\n", "- and all the benefits of nn.Sequential again" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "model1 = initialize(SuperResolutionNetOld(upscale_factor=3))\n", "model2 = initialize(SuperResolutionNetNew(upscale_factor=3))" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [ "assert torch.allclose(model1(torch.zeros(1, 1, 30, 30)), model2(torch.zeros(1, 1, 30, 30))[None])" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "## that's how this code was mentioned to use\n", "\n", "# from PIL import Image\n", "# img = Image.open(opt.input_image).convert('YCbCr')\n", "# y, cb, cr = img.split()\n", "\n", "# model = torch.load(opt.model)\n", "# img_to_tensor = ToTensor()\n", "# input = img_to_tensor(y).view(1, -1, y.size[1], y.size[0])\n", "\n", "# if opt.cuda:\n", "# model = model.cuda()\n", "# input = input.cuda()\n", "\n", "# out = model(input)\n", "# out = out.cpu()\n", "# out_img_y = out[0].detach().numpy()\n", "# out_img_y *= 255.0\n", "# out_img_y = out_img_y.clip(0, 255)\n", "# out_img_y = Image.fromarray(np.uint8(out_img_y[0]), mode='L')\n", "\n", "# out_img_cb = cb.resize(out_img_y.size, Image.BICUBIC)\n", "# out_img_cr = cr.resize(out_img_y.size, Image.BICUBIC)\n", "# out_img = Image.merge('YCbCr', [out_img_y, out_img_cb, out_img_cr]).convert('RGB')\n", "\n", "## Benefits\n", "\n", "# - no need to remembder the order of components in PIL.Image.size (as you see, it is actually different)\n", "# - code explicitly shows shapes passed in and out\n", "# - normalization to [0, 1] range and back is also explicit (it is needed to rememebed in original code that division by 255 is done by ToTensor)\n", "\n", "input_image = '../../logo/einops_logo_350x350.png'\n", "from PIL import Image\n", "import numpy as np\n", "\n", "from torchvision.transforms import ToTensor\n", "model = SuperResolutionNetOld(upscale_factor=2)\n", "\n", "img = Image.open(input_image).convert('YCbCr')\n", "y, cb, cr = img.split()\n", "\n", "img_to_tensor = ToTensor()\n", "input = img_to_tensor(y).view(1, -1, y.size[1], y.size[0])\n", "out = model(input)\n", "\n", "out_img_y = out[0].detach().numpy()\n", "out_img_y = np.clip(out_img_y[0] * 255, 0, 255)\n", "\n", "model = SuperResolutionNetNew(upscale_factor=2)\n", "\n", "img = Image.open(input_image).convert('YCbCr')\n", "y, cb, cr = img.split()\n", "\n", "# TODO numpy.asarray\n", "y = torch.from_numpy(np.array(y, dtype='float32') / 255)\n", "out = model(rearrange(y, 'h w -> () () h w'))\n", "\n", "out_img_y = asnumpy(rearrange(out, '() h w -> h w'))\n", "out_img_y = np.clip(out_img_y * 255, 0, 255)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Restyling Gram matrix for style transfer" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "\n", "Original code is already good - first line shows what kind of input is expected\n", "\n", "- einsum operation should be read like:\n", " - for each batch and for each pair of channels, we sum over h and w.\n", "- I've also changed normalization, because that's how Gram matrix is defined, otherwise we should call it normalized Gram matrix or alike" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [], "source": [ "#left\n", "def gram_matrix_old(y):\n", " (b, ch, h, w) = y.size()\n", " features = y.view(b, ch, w * h)\n", " features_t = features.transpose(1, 2)\n", " gram = features.bmm(features_t) / (ch * h * w)\n", " return gram" ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [], "source": [ "#right\n", "def gram_matrix_new(y):\n", " b, ch, h, w = y.shape\n", " return torch.einsum('bchw,bdhw->bcd', [y, y]) / (h * w)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "It would be great to use just `'b c1 h w,b c2 h w->b c1 c2'`, but einsum supports only one-letter axes" ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [], "source": [ "x = torch.randn([32, 128, 40, 40])" ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "7.58 ms ± 492 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n", "9.66 ms ± 258 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" ] } ], "source": [ "%timeit gram_matrix_old(x).sum()\n", "%timeit gram_matrix_new(x).sum()" ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [], "source": [ "assert torch.allclose(gram_matrix_old(x), gram_matrix_new(x) / 128)" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [], "source": [ "# x = x.to('cuda')\n", "# %timeit -n100 gram_matrix_old(x).sum(); torch.cuda.synchronize()\n", "# %timeit -n100 gram_matrix_new(x).sum(); torch.cuda.synchronize()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Recurrent model\n", "\n", "All we did here is just made information about shapes explicit to skip deciphering\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [], "source": [ "#left\n", "class RNNModelOld(nn.Module):\n", " \"\"\"Container module with an encoder, a recurrent module, and a decoder.\"\"\"\n", " def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):\n", " super(RNNModel, self).__init__()\n", " self.drop = nn.Dropout(dropout)\n", " self.encoder = nn.Embedding(ntoken, ninp)\n", " self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)\n", " self.decoder = nn.Linear(nhid, ntoken)\n", "\n", " def forward(self, input, hidden):\n", " emb = self.drop(self.encoder(input))\n", " output, hidden = self.rnn(emb, hidden)\n", " output = self.drop(output)\n", " decoded = self.decoder(output.view(output.size(0)*output.size(1), output.size(2)))\n", " return decoded.view(output.size(0), output.size(1), decoded.size(1)), hidden\n" ] }, { "cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [], "source": [ "#right\n", "class RNNModelNew(nn.Module):\n", " \"\"\"Container module with an encoder, a recurrent module, and a decoder.\"\"\"\n", " def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):\n", " super(RNNModel, self).__init__()\n", " self.drop = nn.Dropout(p=dropout)\n", " self.encoder = nn.Embedding(ntoken, ninp)\n", " self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)\n", " self.decoder = nn.Linear(nhid, ntoken)\n", "\n", " def forward(self, input, hidden):\n", " t, b = input.shape\n", " emb = self.drop(self.encoder(input))\n", " output, hidden = self.rnn(emb, hidden)\n", " output = rearrange(self.drop(output), 't b nhid -> (t b) nhid')\n", " decoded = rearrange(self.decoder(output), '(t b) token -> t b token', t=t, b=b)\n", " return decoded, hidden" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Channel shuffle (from shufflenet) \n", "" ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [], "source": [ "#left\n", "def channel_shuffle_old(x, groups):\n", " batchsize, num_channels, height, width = x.data.size()\n", "\n", " channels_per_group = num_channels // groups\n", "\n", " # reshape\n", " x = x.view(batchsize, groups,\n", " channels_per_group, height, width)\n", "\n", " # transpose\n", " # - contiguous() required if transpose() is used before view().\n", " # See https://github.com/pytorch/pytorch/issues/764\n", " x = torch.transpose(x, 1, 2).contiguous()\n", "\n", " # flatten\n", " x = x.view(batchsize, -1, height, width)\n", "\n", " return x" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [], "source": [ "#right\n", "def channel_shuffle_new(x, groups):\n", " return rearrange(x, 'b (c1 c2) h w -> b (c2 c1) h w', c1=groups)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "While progress is obvious, this is not the limit. As you'll see below, we don't even need to write these couple of lines." ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [], "source": [ "x = torch.zeros([32, 64, 100, 100])" ] }, { "cell_type": "code", "execution_count": 22, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "51.2 ms ± 2.18 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)\n", "49.9 ms ± 594 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" ] } ], "source": [ "%timeit -n100 channel_shuffle_old(x, 8); torch.cuda.synchronize()\n", "%timeit -n100 channel_shuffle_new(x, 8); torch.cuda.synchronize()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Shufflenet" ] }, { "cell_type": "code", "execution_count": 23, "metadata": {}, "outputs": [], "source": [ "def conv3x3(in_channels, out_channels, stride=1,\n", " padding=1, bias=True, groups=1):\n", " \"\"\"3x3 convolution with padding\n", " \"\"\"\n", " return nn.Conv2d(\n", " in_channels,\n", " out_channels,\n", " kernel_size=3,\n", " stride=stride,\n", " padding=padding,\n", " bias=bias,\n", " groups=groups)\n", "\n", "\n", "def conv1x1(in_channels, out_channels, groups=1):\n", " \"\"\"1x1 convolution with padding\n", " - Normal pointwise convolution When groups == 1\n", " - Grouped pointwise convolution when groups > 1\n", " \"\"\"\n", " return nn.Conv2d(\n", " in_channels,\n", " out_channels,\n", " kernel_size=1,\n", " groups=groups,\n", " stride=1)" ] }, { "cell_type": "code", "execution_count": 24, "metadata": {}, "outputs": [], "source": [ "#left\n", "from collections import OrderedDict\n", "\n", "def channel_shuffle(x, groups):\n", " batchsize, num_channels, height, width = x.data.size()\n", "\n", " channels_per_group = num_channels // groups\n", "\n", " # reshape\n", " x = x.view(batchsize, groups,\n", " channels_per_group, height, width)\n", "\n", " # transpose\n", " # - contiguous() required if transpose() is used before view().\n", " # See https://github.com/pytorch/pytorch/issues/764\n", " x = torch.transpose(x, 1, 2).contiguous()\n", "\n", " # flatten\n", " x = x.view(batchsize, -1, height, width)\n", "\n", " return x\n", "\n", "class ShuffleUnitOld(nn.Module):\n", " def __init__(self, in_channels, out_channels, groups=3,\n", " grouped_conv=True, combine='add'):\n", "\n", " super(ShuffleUnitOld, self).__init__()\n", "\n", " self.in_channels = in_channels\n", " self.out_channels = out_channels\n", " self.grouped_conv = grouped_conv\n", " self.combine = combine\n", " self.groups = groups\n", " self.bottleneck_channels = self.out_channels // 4\n", "\n", " # define the type of ShuffleUnit\n", " if self.combine == 'add':\n", " # ShuffleUnit Figure 2b\n", " self.depthwise_stride = 1\n", " self._combine_func = self._add\n", " elif self.combine == 'concat':\n", " # ShuffleUnit Figure 2c\n", " self.depthwise_stride = 2\n", " self._combine_func = self._concat\n", "\n", " # ensure output of concat has the same channels as\n", " # original output channels.\n", " self.out_channels -= self.in_channels\n", " else:\n", " raise ValueError(\"Cannot combine tensors with \\\"{}\\\"\" \\\n", " \"Only \\\"add\\\" and \\\"concat\\\" are\" \\\n", " \"supported\".format(self.combine))\n", "\n", " # Use a 1x1 grouped or non-grouped convolution to reduce input channels\n", " # to bottleneck channels, as in a ResNet bottleneck module.\n", " # NOTE: Do not use group convolution for the first conv1x1 in Stage 2.\n", " self.first_1x1_groups = self.groups if grouped_conv else 1\n", "\n", " self.g_conv_1x1_compress = self._make_grouped_conv1x1(\n", " self.in_channels,\n", " self.bottleneck_channels,\n", " self.first_1x1_groups,\n", " batch_norm=True,\n", " relu=True\n", " )\n", "\n", " # 3x3 depthwise convolution followed by batch normalization\n", " self.depthwise_conv3x3 = conv3x3(\n", " self.bottleneck_channels, self.bottleneck_channels,\n", " stride=self.depthwise_stride, groups=self.bottleneck_channels)\n", " self.bn_after_depthwise = nn.BatchNorm2d(self.bottleneck_channels)\n", "\n", " # Use 1x1 grouped convolution to expand from\n", " # bottleneck_channels to out_channels\n", " self.g_conv_1x1_expand = self._make_grouped_conv1x1(\n", " self.bottleneck_channels,\n", " self.out_channels,\n", " self.groups,\n", " batch_norm=True,\n", " relu=False\n", " )\n", "\n", "\n", " @staticmethod\n", " def _add(x, out):\n", " # residual connection\n", " return x + out\n", "\n", "\n", " @staticmethod\n", " def _concat(x, out):\n", " # concatenate along channel axis\n", " return torch.cat((x, out), 1)\n", "\n", "\n", " def _make_grouped_conv1x1(self, in_channels, out_channels, groups,\n", " batch_norm=True, relu=False):\n", "\n", " modules = OrderedDict()\n", " conv = conv1x1(in_channels, out_channels, groups=groups)\n", " modules['conv1x1'] = conv\n", "\n", " if batch_norm:\n", " modules['batch_norm'] = nn.BatchNorm2d(out_channels)\n", " if relu:\n", " modules['relu'] = nn.ReLU()\n", " if len(modules) > 1:\n", " return nn.Sequential(modules)\n", " else:\n", " return conv\n", "\n", "\n", " def forward(self, x):\n", " # save for combining later with output\n", " residual = x\n", " if self.combine == 'concat':\n", " residual = F.avg_pool2d(residual, kernel_size=3,\n", " stride=2, padding=1)\n", "\n", " out = self.g_conv_1x1_compress(x)\n", " out = channel_shuffle(out, self.groups)\n", " out = self.depthwise_conv3x3(out)\n", " out = self.bn_after_depthwise(out)\n", " out = self.g_conv_1x1_expand(out)\n", "\n", " out = self._combine_func(residual, out)\n", " return F.relu(out)" ] }, { "cell_type": "code", "execution_count": 25, "metadata": {}, "outputs": [], "source": [ "#right\n", "class ShuffleUnitNew(nn.Module):\n", " def __init__(self, in_channels, out_channels, groups=3,\n", " grouped_conv=True, combine='add'):\n", " super().__init__()\n", " first_1x1_groups = groups if grouped_conv else 1\n", " bottleneck_channels = out_channels // 4\n", " self.combine = combine\n", " if combine == 'add':\n", " # ShuffleUnit Figure 2b\n", " self.left = Rearrange('...->...') # identity\n", " depthwise_stride = 1\n", " else:\n", " # ShuffleUnit Figure 2c\n", " self.left = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)\n", " depthwise_stride = 2\n", " # ensure output of concat has the same channels as original output channels.\n", " out_channels -= in_channels\n", " assert out_channels > 0\n", "\n", " self.right = nn.Sequential(\n", " # Use a 1x1 grouped or non-grouped convolution to reduce input channels\n", " # to bottleneck channels, as in a ResNet bottleneck module.\n", " conv1x1(in_channels, bottleneck_channels, groups=first_1x1_groups),\n", " nn.BatchNorm2d(bottleneck_channels),\n", " nn.ReLU(inplace=True),\n", " # channel shuffle\n", " Rearrange('b (c1 c2) h w -> b (c2 c1) h w', c1=groups),\n", " # 3x3 depthwise convolution followed by batch\n", " conv3x3(bottleneck_channels, bottleneck_channels,\n", " stride=depthwise_stride, groups=bottleneck_channels),\n", " nn.BatchNorm2d(bottleneck_channels),\n", " # Use 1x1 grouped convolution to expand from\n", " # bottleneck_channels to out_channels\n", " conv1x1(bottleneck_channels, out_channels, groups=groups),\n", " nn.BatchNorm2d(out_channels),\n", " )\n", "\n", " def forward(self, x):\n", " if self.combine == 'add':\n", " combined = self.left(x) + self.right(x)\n", " else:\n", " combined = torch.cat([self.left(x), self.right(x)], dim=1)\n", " return F.relu(combined, inplace=True)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Rewriting the code helped to identify:\n", "\n", "- There is no sense in doing reshuffling and not using groups in the first convolution\n", " (indeed, in the paper it is not so). However, result is an equivalent model.\n", "- It is also strange that the first convolution may be not grouped, while the last convolution is always grouped\n", " (and that is different from the paper)\n", "\n", "Other comments:\n", "\n", "- There is an identity layer for pytorch introduced here\n", "- The last thing left is get rid of conv1x1 and conv3x3 in the code - those are not better than standard" ] }, { "cell_type": "code", "execution_count": 26, "metadata": {}, "outputs": [], "source": [ "model1 = ShuffleUnitOld(32, 32, groups=4, grouped_conv=True, combine='add')\n", "model2 = ShuffleUnitNew(32, 32, groups=4, grouped_conv=True, combine='add')" ] }, { "cell_type": "code", "execution_count": 27, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "True" ] }, "execution_count": 27, "metadata": {}, "output_type": "execute_result" } ], "source": [ "x = torch.randn(1, 32, 14, 14)\n", "initialize(model1)\n", "initialize(model2)\n", "torch.allclose(model1(x), model2(x))" ] }, { "cell_type": "code", "execution_count": 28, "metadata": {}, "outputs": [], "source": [ "import pickle\n", "dump1 = pickle.dumps(model1._combine_func)\n", "dump2 = pickle.dumps(model2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Simplifying ResNet" ] }, { "cell_type": "code", "execution_count": 29, "metadata": {}, "outputs": [], "source": [ "#left\n", "class ResNetOld(nn.Module):\n", "\n", " def __init__(self, block, layers, num_classes=1000):\n", " self.inplanes = 64\n", " super(ResNetOld, self).__init__()\n", " self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,\n", " bias=False)\n", " self.bn1 = nn.BatchNorm2d(64)\n", " self.relu = nn.ReLU(inplace=True)\n", " self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n", " self.layer1 = self._make_layer(block, 64, layers[0])\n", " self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n", " self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n", " self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n", " self.avgpool = nn.AvgPool2d(7, stride=1)\n", "\n", " self.fc = nn.Linear(512 * block.expansion, num_classes)\n", "\n", " for m in self.modules():\n", " if isinstance(m, nn.Conv2d):\n", " n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n", " m.weight.data.normal_(0, math.sqrt(2. / n))\n", " elif isinstance(m, nn.BatchNorm2d):\n", " m.weight.data.fill_(1)\n", " m.bias.data.zero_()\n", "\n", " def _make_layer(self, block, planes, blocks, stride=1):\n", " downsample = None\n", " if stride != 1 or self.inplanes != planes * block.expansion:\n", " downsample = nn.Sequential(\n", " nn.Conv2d(self.inplanes, planes * block.expansion,\n", " kernel_size=1, stride=stride, bias=False),\n", " nn.BatchNorm2d(planes * block.expansion),\n", " )\n", "\n", " layers = []\n", " layers.append(block(self.inplanes, planes, stride, downsample))\n", " self.inplanes = planes * block.expansion\n", " for i in range(1, blocks):\n", " layers.append(block(self.inplanes, planes))\n", "\n", " return nn.Sequential(*layers)\n", "\n", " def forward(self, x):\n", " x = self.conv1(x)\n", " x = self.bn1(x)\n", " x = self.relu(x)\n", " x = self.maxpool(x)\n", "\n", " x = self.layer1(x)\n", " x = self.layer2(x)\n", " x = self.layer3(x)\n", " x = self.layer4(x)\n", " x = self.avgpool(x)\n", " x = x.view(x.size(0), -1)\n", " x = self.fc(x)\n", "\n", " return x" ] }, { "cell_type": "code", "execution_count": 30, "metadata": {}, "outputs": [], "source": [ "#right\n", "def make_layer(inplanes, planes, block, n_blocks, stride=1):\n", " downsample = None\n", " if stride != 1 or inplanes != planes * block.expansion:\n", " # output size won't match input, so adjust residual\n", " downsample = nn.Sequential(\n", " nn.Conv2d(inplanes, planes * block.expansion,\n", " kernel_size=1, stride=stride, bias=False),\n", " nn.BatchNorm2d(planes * block.expansion),\n", " )\n", " return nn.Sequential(\n", " block(inplanes, planes, stride, downsample),\n", " *[block(planes * block.expansion, planes) for _ in range(1, n_blocks)]\n", " )\n", "\n", "\n", "def ResNetNew(block, layers, num_classes=1000):\n", " e = block.expansion\n", "\n", " resnet = nn.Sequential(\n", " Rearrange('b c h w -> b c h w', c=3, h=224, w=224),\n", " nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),\n", " nn.BatchNorm2d(64),\n", " nn.ReLU(inplace=True),\n", " nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n", " make_layer(64, 64, block, layers[0], stride=1),\n", " make_layer(64 * e, 128, block, layers[1], stride=2),\n", " make_layer(128 * e, 256, block, layers[2], stride=2),\n", " make_layer(256 * e, 512, block, layers[3], stride=2),\n", " # combined AvgPool and view in one averaging operation\n", " Reduce('b c h w -> b c', 'mean'),\n", " nn.Linear(512 * e, num_classes),\n", " )\n", "\n", " # initialization\n", " for m in resnet.modules():\n", " if isinstance(m, nn.Conv2d):\n", " n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n", " m.weight.data.normal_(0, math.sqrt(2. / n))\n", " elif isinstance(m, nn.BatchNorm2d):\n", " m.weight.data.fill_(1)\n", " m.bias.data.zero_()\n", " return resnet" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Changes:\n", "\n", "- explicit check for input shape\n", "- no views and simple sequential structure, output is just nn.Sequential, so can always be saved/passed/etc\n", "- no need in AvgPool and additional views, this place is much clearer now\n", "- `make_layer` doesn't use internal state (that's quite faulty place)" ] }, { "cell_type": "code", "execution_count": 31, "metadata": {}, "outputs": [], "source": [ "from torchvision.models.resnet import BasicBlock, Bottleneck, ResNet" ] }, { "cell_type": "code", "execution_count": 32, "metadata": {}, "outputs": [], "source": [ "x = torch.randn(2, 3, 224, 224)\n", "with torch.no_grad():\n", " model_old = ResNetOld(BasicBlock, layers=[2, 2, 2, 3])\n", " model_new = ResNetNew(BasicBlock, layers=[2, 2, 2, 3])\n", " initialize(model_old)\n", " initialize(model_new)\n", " assert torch.allclose(model_old(x), model_new(x), atol=1e-3)" ] }, { "cell_type": "code", "execution_count": 33, "metadata": {}, "outputs": [], "source": [ "# with torch.no_grad():\n", "# x = torch.randn([2, 512, 7, 7])\n", "# torch.allclose(nn.AvgPool2d(7)(x), reduce(x, 'b c h w -> b c', 'mean'), atol=1e-8)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Improving RNN language modelling" ] }, { "cell_type": "code", "execution_count": 34, "metadata": {}, "outputs": [], "source": [ "#left\n", "class RNNOld(nn.Module):\n", " def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, bidirectional, dropout):\n", " super().__init__()\n", "\n", " self.embedding = nn.Embedding(vocab_size, embedding_dim)\n", " self.rnn = nn.LSTM(embedding_dim, hidden_dim, num_layers=n_layers,\n", " bidirectional=bidirectional, dropout=dropout)\n", " self.fc = nn.Linear(hidden_dim*2, output_dim)\n", " self.dropout = nn.Dropout(dropout)\n", "\n", " def forward(self, x):\n", " #x = [sent len, batch size]\n", "\n", " embedded = self.dropout(self.embedding(x))\n", "\n", " #embedded = [sent len, batch size, emb dim]\n", "\n", " output, (hidden, cell) = self.rnn(embedded)\n", "\n", " #output = [sent len, batch size, hid dim * num directions]\n", " #hidden = [num layers * num directions, batch size, hid dim]\n", " #cell = [num layers * num directions, batch size, hid dim]\n", "\n", " #concat the final forward (hidden[-2,:,:]) and backward (hidden[-1,:,:]) hidden layers\n", " #and apply dropout\n", "\n", " hidden = self.dropout(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1))\n", "\n", " #hidden = [batch size, hid dim * num directions]\n", "\n", " return self.fc(hidden.squeeze(0))\n" ] }, { "cell_type": "code", "execution_count": 35, "metadata": {}, "outputs": [], "source": [ "#right\n", "class RNNNew(nn.Module):\n", " def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, bidirectional, dropout):\n", " super().__init__()\n", "\n", " self.embedding = nn.Embedding(vocab_size, embedding_dim)\n", " self.rnn = nn.LSTM(embedding_dim, hidden_dim, num_layers=n_layers,\n", " bidirectional=bidirectional, dropout=dropout)\n", " self.dropout = nn.Dropout(dropout)\n", " self.directions = 2 if bidirectional else 1\n", " self.fc = nn.Linear(hidden_dim * self.directions, output_dim)\n", "\n", " def forward(self, x):\n", " #x = [sent len, batch size]\n", " embedded = self.dropout(self.embedding(x))\n", "\n", " #embedded = [sent len, batch size, emb dim]\n", " output, (hidden, cell) = self.rnn(embedded)\n", "\n", " hidden = rearrange(hidden, '(layer dir) b c -> layer b (dir c)',\n", " dir=self.directions)\n", " # take the final layer's hidden\n", " return self.fc(self.dropout(hidden[-1]))\n" ] }, { "cell_type": "code", "execution_count": 36, "metadata": {}, "outputs": [], "source": [ "model_old = initialize(RNNOld(10, 10, 10, output_dim=15, n_layers=2, bidirectional=True, dropout=0.1)).eval()\n", "model_new = initialize(RNNNew(10, 10, 10, output_dim=15, n_layers=2, bidirectional=True, dropout=0.1)).eval()\n", "\n", "x = torch.randint(0, 10, size=[23, 10]).long()\n", "\n", "assert torch.allclose(model_old(x), model_new(x))" ] }, { "cell_type": "code", "execution_count": 37, "metadata": {}, "outputs": [], "source": [ "# this code fails\n", "# model_old = initialize(RNNOld(10, 10, 10, output_dim=15, n_layers=1, bidirectional=False, dropout=0.1)).eval()\n", "# model_old(x).shape" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "- original code misbehaves for non-bidirectional models\n", "- ... and fails when bidirectional = False, and there is only one layer\n", "- modification of the code shows both how hidden is structured and how it is modified" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Writing FastText faster\n", "\n", "" ] }, { "cell_type": "code", "execution_count": 38, "metadata": {}, "outputs": [], "source": [ "#left\n", "class FastTextOld(nn.Module):\n", " def __init__(self, vocab_size, embedding_dim, output_dim):\n", " super().__init__()\n", "\n", " self.embedding = nn.Embedding(vocab_size, embedding_dim)\n", " self.fc = nn.Linear(embedding_dim, output_dim)\n", "\n", " def forward(self, x):\n", "\n", " #x = [sent len, batch size]\n", "\n", " embedded = self.embedding(x)\n", "\n", " #embedded = [sent len, batch size, emb dim]\n", "\n", " embedded = embedded.permute(1, 0, 2)\n", "\n", " #embedded = [batch size, sent len, emb dim]\n", "\n", " pooled = F.avg_pool2d(embedded, (embedded.shape[1], 1)).squeeze(1)\n", "\n", " #pooled = [batch size, embedding_dim]\n", "\n", " return self.fc(pooled)" ] }, { "cell_type": "code", "execution_count": 39, "metadata": {}, "outputs": [], "source": [ "#right\n", "def FastTextNew(vocab_size, embedding_dim, output_dim):\n", " return nn.Sequential(\n", " Rearrange('t b -> t b'),\n", " nn.Embedding(vocab_size, embedding_dim),\n", " Reduce('t b c -> b c', 'mean'),\n", " nn.Linear(embedding_dim, output_dim),\n", " Rearrange('b c -> b c'),\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Some comments on new code:\n", "\n", "- first and last operations do nothing and can be removed\n", " - but were added to explicitly show expected input and output\n", "- this also gives you a flexibility of changing interface by editing a single line. Should you need to accept inputs as (batch, time), \n", " you just change first line to `Rearrange('b t -> t b'),`" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# CNNs for text classification" ] }, { "cell_type": "code", "execution_count": 40, "metadata": {}, "outputs": [], "source": [ "#left\n", "class CNNOld(nn.Module):\n", " def __init__(self, vocab_size, embedding_dim, n_filters, filter_sizes, output_dim, dropout):\n", " super().__init__()\n", " self.embedding = nn.Embedding(vocab_size, embedding_dim)\n", " self.conv_0 = nn.Conv2d(in_channels=1, out_channels=n_filters, kernel_size=(filter_sizes[0],embedding_dim))\n", " self.conv_1 = nn.Conv2d(in_channels=1, out_channels=n_filters, kernel_size=(filter_sizes[1],embedding_dim))\n", " self.conv_2 = nn.Conv2d(in_channels=1, out_channels=n_filters, kernel_size=(filter_sizes[2],embedding_dim))\n", " self.fc = nn.Linear(len(filter_sizes)*n_filters, output_dim)\n", " self.dropout = nn.Dropout(dropout)\n", "\n", " def forward(self, x):\n", "\n", " #x = [sent len, batch size]\n", "\n", " x = x.permute(1, 0)\n", "\n", " #x = [batch size, sent len]\n", "\n", " embedded = self.embedding(x)\n", "\n", " #embedded = [batch size, sent len, emb dim]\n", "\n", " embedded = embedded.unsqueeze(1)\n", "\n", " #embedded = [batch size, 1, sent len, emb dim]\n", "\n", " conved_0 = F.relu(self.conv_0(embedded).squeeze(3))\n", " conved_1 = F.relu(self.conv_1(embedded).squeeze(3))\n", " conved_2 = F.relu(self.conv_2(embedded).squeeze(3))\n", "\n", " #conv_n = [batch size, n_filters, sent len - filter_sizes[n]]\n", "\n", " pooled_0 = F.max_pool1d(conved_0, conved_0.shape[2]).squeeze(2)\n", " pooled_1 = F.max_pool1d(conved_1, conved_1.shape[2]).squeeze(2)\n", " pooled_2 = F.max_pool1d(conved_2, conved_2.shape[2]).squeeze(2)\n", "\n", " #pooled_n = [batch size, n_filters]\n", "\n", " cat = self.dropout(torch.cat((pooled_0, pooled_1, pooled_2), dim=1))\n", "\n", " #cat = [batch size, n_filters * len(filter_sizes)]\n", "\n", " return self.fc(cat)" ] }, { "cell_type": "code", "execution_count": 41, "metadata": {}, "outputs": [], "source": [ "#right\n", "class CNNNew(nn.Module):\n", " def __init__(self, vocab_size, embedding_dim, n_filters, filter_sizes, output_dim, dropout):\n", " super().__init__()\n", " self.embedding = nn.Embedding(vocab_size, embedding_dim)\n", " self.convs = nn.ModuleList([\n", " nn.Conv1d(embedding_dim, n_filters, kernel_size=size) for size in filter_sizes\n", " ])\n", " self.fc = nn.Linear(len(filter_sizes) * n_filters, output_dim)\n", " self.dropout = nn.Dropout(dropout)\n", "\n", " def forward(self, x):\n", " x = rearrange(x, 't b -> t b')\n", " emb = rearrange(self.embedding(x), 't b c -> b c t')\n", " pooled = [reduce(conv(emb), 'b c t -> b c', 'max') for conv in self.convs]\n", " concatenated = rearrange(pooled, 'filter b c -> b (filter c)')\n", " return self.fc(self.dropout(F.relu(concatenated)))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "- Original code misuses Conv2d, while Conv1d is the right choice\n", "- Fixed code can work with any number of filter_sizes (and won't fail)\n", "- First line in new code does nothing, but was added for simplicity" ] }, { "cell_type": "code", "execution_count": 42, "metadata": {}, "outputs": [], "source": [ "# old_model = initialize(CNNOld(32, 32, 32, [1, 2, 4], 32, dropout=0.1)).eval()\n", "# new_model = initialize(CNNNew(32, 32, 32, [1, 2, 4], 32, dropout=0.1)).eval()\n", "\n", "# x = torch.zeros([10, 20]).long()\n", "# assert torch.allclose(old_model(x), new_model(x), atol=1e-3)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Highway convolutions\n", "\n", "- Highway convolutions are common in TTS systems. Code below makes splitting a bit more explicit.\n", "- Splitting policy may eventually turn out to be important if input had previously groups over channel axes (group convolutions or bidirectional LSTMs/GRUs)\n", "- Same applies to GLU and gated units in general" ] }, { "cell_type": "code", "execution_count": 43, "metadata": {}, "outputs": [], "source": [ "#left\n", "class HighwayConv1dOld(nn.Conv1d):\n", " def forward(self, inputs):\n", " L = super(HighwayConv1dOld, self).forward(inputs)\n", " H1, H2 = torch.chunk(L, 2, 1) # chunk at the feature dim\n", " torch.sigmoid_(H1)\n", " return H1 * H2 + (1.0 - H1) * inputs" ] }, { "cell_type": "code", "execution_count": 44, "metadata": {}, "outputs": [], "source": [ "#right\n", "class HighwayConv1dNew(nn.Conv1d):\n", " def forward(self, inputs):\n", " L = super().forward(inputs)\n", " H1, H2 = rearrange(L, 'b (split c) t -> split b c t', split=2)\n", " torch.sigmoid_(H1)\n", " return H1 * H2 + (1.0 - H1) * inputs" ] }, { "cell_type": "code", "execution_count": 45, "metadata": {}, "outputs": [], "source": [ "hc1 = HighwayConv1dOld(10, 20, kernel_size=3, padding=1)\n", "hc2 = HighwayConv1dNew(10, 20, kernel_size=3, padding=1)\n", "initialize(hc1)\n", "initialize(hc2)\n", "fw1 = hc1(torch.zeros(1, 10, 100))\n", "fw2 = hc2(torch.zeros(1, 10, 100))\n", "assert torch.allclose(fw1, fw2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Tacotron's CBHG module\n", "\n", "" ] }, { "cell_type": "code", "execution_count": 46, "metadata": {}, "outputs": [], "source": [ "#right\n", "class CBHG_Old(nn.Module):\n", " \"\"\"CBHG module: a recurrent neural network composed of:\n", " - 1-d convolution banks\n", " - Highway networks + residual connections\n", " - Bidirectional gated recurrent units\n", " \"\"\"\n", "\n", " def __init__(self, in_dim, K=16, projections=[128, 128]):\n", " super(CBHG, self).__init__()\n", " self.in_dim = in_dim\n", " self.relu = nn.ReLU()\n", " self.conv1d_banks = nn.ModuleList(\n", " [BatchNormConv1d(in_dim, in_dim, kernel_size=k, stride=1,\n", " padding=k // 2, activation=self.relu)\n", " for k in range(1, K + 1)])\n", " self.max_pool1d = nn.MaxPool1d(kernel_size=2, stride=1, padding=1)\n", "\n", " in_sizes = [K * in_dim] + projections[:-1]\n", " activations = [self.relu] * (len(projections) - 1) + [None]\n", " self.conv1d_projections = nn.ModuleList(\n", " [BatchNormConv1d(in_size, out_size, kernel_size=3, stride=1,\n", " padding=1, activation=ac)\n", " for (in_size, out_size, ac) in zip(\n", " in_sizes, projections, activations)])\n", "\n", " self.pre_highway = nn.Linear(projections[-1], in_dim, bias=False)\n", " self.highways = nn.ModuleList(\n", " [Highway(in_dim, in_dim) for _ in range(4)])\n", "\n", " self.gru = nn.GRU(\n", " in_dim, in_dim, 1, batch_first=True, bidirectional=True)" ] }, { "cell_type": "code", "execution_count": 47, "metadata": {}, "outputs": [], "source": [ "#left\n", "def forward_old(self, inputs):\n", " # (B, T_in, in_dim)\n", " x = inputs\n", "\n", " # Needed to perform conv1d on time-axis\n", " # (B, in_dim, T_in)\n", " if x.size(-1) == self.in_dim:\n", " x = x.transpose(1, 2)\n", "\n", " T = x.size(-1)\n", "\n", " # (B, in_dim*K, T_in)\n", " # Concat conv1d bank outputs\n", " x = torch.cat([conv1d(x)[:, :, :T] for conv1d in self.conv1d_banks], dim=1)\n", " assert x.size(1) == self.in_dim * len(self.conv1d_banks)\n", " x = self.max_pool1d(x)[:, :, :T]\n", "\n", " for conv1d in self.conv1d_projections:\n", " x = conv1d(x)\n", "\n", " # (B, T_in, in_dim)\n", " # Back to the original shape\n", " x = x.transpose(1, 2)\n", "\n", " if x.size(-1) != self.in_dim:\n", " x = self.pre_highway(x)\n", "\n", " # Residual connection\n", " x += inputs\n", " for highway in self.highways:\n", " x = highway(x)\n", "\n", " # (B, T_in, in_dim*2)\n", " outputs, _ = self.gru(x)\n", "\n", " return outputs" ] }, { "cell_type": "code", "execution_count": 48, "metadata": {}, "outputs": [], "source": [ "#right\n", "def forward_new(self, inputs, input_lengths=None):\n", " x = rearrange(inputs, 'b t c -> b c t')\n", " _, _, T = x.shape\n", " # Concat conv1d bank outputs\n", " x = rearrange([conv1d(x)[:, :, :T] for conv1d in self.conv1d_banks],\n", " 'bank b c t -> b (bank c) t', c=self.in_dim)\n", " x = self.max_pool1d(x)[:, :, :T]\n", "\n", " for conv1d in self.conv1d_projections:\n", " x = conv1d(x)\n", " x = rearrange(x, 'b c t -> b t c')\n", " if x.size(-1) != self.in_dim:\n", " x = self.pre_highway(x)\n", "\n", " # Residual connection\n", " x += inputs\n", " for highway in self.highways:\n", " x = highway(x)\n", "\n", " # (B, T_in, in_dim*2)\n", " outputs, _ = self.gru(self.highways(x))\n", "\n", " return outputs" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "There is still a large room for improvements, but in this example only forward function was changed" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Simple attention\n", "Good news: there is no more need to guess order of dimensions. Neither for inputs nor for outputs" ] }, { "cell_type": "code", "execution_count": 49, "metadata": {}, "outputs": [], "source": [ "#left\n", "class Attention(nn.Module):\n", " def __init__(self):\n", " super(Attention, self).__init__()\n", "\n", " def forward(self, K, V, Q):\n", " A = torch.bmm(K.transpose(1,2), Q) / np.sqrt(Q.shape[1])\n", " A = F.softmax(A, 1)\n", " R = torch.bmm(V, A)\n", " return torch.cat((R, Q), dim=1)" ] }, { "cell_type": "code", "execution_count": 50, "metadata": {}, "outputs": [], "source": [ "#right\n", "def attention(K, V, Q):\n", " _, n_channels, _ = K.shape\n", " A = torch.einsum('bct,bcl->btl', [K, Q])\n", " A = F.softmax(A * n_channels ** (-0.5), 1)\n", " R = torch.einsum('bct,btl->bcl', [V, A])\n", " return torch.cat((R, Q), dim=1)" ] }, { "cell_type": "code", "execution_count": 51, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "336 µs ± 7.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n", "336 µs ± 7.48 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" ] } ], "source": [ "args = dict(\n", " K=torch.zeros(32, 128, 40).cuda(),\n", " V=torch.zeros(32, 128, 40).cuda(),\n", " Q=torch.zeros(32, 128, 30).cuda(),\n", ")\n", "\n", "%timeit -n100 result_old = Attention()(**args); torch.cuda.synchronize()\n", "%timeit -n100 result_new = attention(**args); torch.cuda.synchronize()" ] }, { "cell_type": "code", "execution_count": 52, "metadata": {}, "outputs": [], "source": [ "result_old = Attention()(**args); torch.cuda.synchronize()\n", "result_new = attention(**args); torch.cuda.synchronize()\n", "assert torch.allclose(result_old, result_new)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Transformer's attention needs more attention" ] }, { "cell_type": "code", "execution_count": 53, "metadata": {}, "outputs": [], "source": [ "#left\n", "class ScaledDotProductAttention(nn.Module):\n", " ''' Scaled Dot-Product Attention '''\n", "\n", " def __init__(self, temperature, attn_dropout=0.1):\n", " super().__init__()\n", " self.temperature = temperature\n", " self.dropout = nn.Dropout(attn_dropout)\n", " self.softmax = nn.Softmax(dim=2)\n", "\n", " def forward(self, q, k, v, mask=None):\n", "\n", " attn = torch.bmm(q, k.transpose(1, 2))\n", " attn = attn / self.temperature\n", "\n", " if mask is not None:\n", " attn = attn.masked_fill(mask, -np.inf)\n", "\n", " attn = self.softmax(attn)\n", " attn = self.dropout(attn)\n", " output = torch.bmm(attn, v)\n", "\n", " return output, attn\n", "\n", "\n", "\n", "class MultiHeadAttentionOld(nn.Module):\n", " ''' Multi-Head Attention module '''\n", "\n", " def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):\n", " super().__init__()\n", "\n", " self.n_head = n_head\n", " self.d_k = d_k\n", " self.d_v = d_v\n", "\n", " self.w_qs = nn.Linear(d_model, n_head * d_k)\n", " self.w_ks = nn.Linear(d_model, n_head * d_k)\n", " self.w_vs = nn.Linear(d_model, n_head * d_v)\n", " nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))\n", " nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))\n", " nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v)))\n", "\n", " self.attention = ScaledDotProductAttention(temperature=np.power(d_k, 0.5))\n", " self.layer_norm = nn.LayerNorm(d_model)\n", "\n", " self.fc = nn.Linear(n_head * d_v, d_model)\n", " nn.init.xavier_normal_(self.fc.weight)\n", "\n", " self.dropout = nn.Dropout(dropout)\n", "\n", "\n", " def forward(self, q, k, v, mask=None):\n", "\n", " d_k, d_v, n_head = self.d_k, self.d_v, self.n_head\n", "\n", " sz_b, len_q, _ = q.size()\n", " sz_b, len_k, _ = k.size()\n", " sz_b, len_v, _ = v.size()\n", "\n", " residual = q\n", "\n", " q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)\n", " k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)\n", " v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)\n", "\n", " q = q.permute(2, 0, 1, 3).contiguous().view(-1, len_q, d_k) # (n*b) x lq x dk\n", " k = k.permute(2, 0, 1, 3).contiguous().view(-1, len_k, d_k) # (n*b) x lk x dk\n", " v = v.permute(2, 0, 1, 3).contiguous().view(-1, len_v, d_v) # (n*b) x lv x dv\n", "\n", " mask = mask.repeat(n_head, 1, 1) # (n*b) x .. x ..\n", " output, attn = self.attention(q, k, v, mask=mask)\n", "\n", " output = output.view(n_head, sz_b, len_q, d_v)\n", " output = output.permute(1, 2, 0, 3).contiguous().view(sz_b, len_q, -1) # b x lq x (n*dv)\n", "\n", " output = self.dropout(self.fc(output))\n", " output = self.layer_norm(output + residual)\n", "\n", " return output, attn\n" ] }, { "cell_type": "code", "execution_count": 54, "metadata": {}, "outputs": [], "source": [ "#right\n", "class MultiHeadAttentionNew(nn.Module):\n", " def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):\n", " super().__init__()\n", " self.n_head = n_head\n", "\n", " self.w_qs = nn.Linear(d_model, n_head * d_k)\n", " self.w_ks = nn.Linear(d_model, n_head * d_k)\n", " self.w_vs = nn.Linear(d_model, n_head * d_v)\n", "\n", " nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))\n", " nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))\n", " nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v)))\n", "\n", " self.fc = nn.Linear(n_head * d_v, d_model)\n", " nn.init.xavier_normal_(self.fc.weight)\n", " self.dropout = nn.Dropout(p=dropout)\n", " self.layer_norm = nn.LayerNorm(d_model)\n", "\n", " def forward(self, q, k, v, mask=None):\n", " residual = q\n", " q = rearrange(self.w_qs(q), 'b l (head k) -> head b l k', head=self.n_head)\n", " k = rearrange(self.w_ks(k), 'b t (head k) -> head b t k', head=self.n_head)\n", " v = rearrange(self.w_vs(v), 'b t (head v) -> head b t v', head=self.n_head)\n", " attn = torch.einsum('hblk,hbtk->hblt', [q, k]) / np.sqrt(q.shape[-1])\n", " if mask is not None:\n", " attn = attn.masked_fill(mask[None], -np.inf)\n", " attn = torch.softmax(attn, dim=3)\n", " output = torch.einsum('hblt,hbtv->hblv', [attn, v])\n", " output = rearrange(output, 'head b l v -> b l (head v)')\n", " output = self.dropout(self.fc(output))\n", " output = self.layer_norm(output + residual)\n", " return output, attn\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Benefits of new implementation\n", "\n", "- we have one module, not two\n", "- now code does not fail for None mask\n", "- the amount of caveats in the original code that we removed is huge. \n", " Try erasing comments and deciphering what happens there" ] }, { "cell_type": "code", "execution_count": 55, "metadata": {}, "outputs": [], "source": [ "# Poor implementation of torch.einsum, so code below doesn't work\n", "class MultiHeadAttentionHard(nn.Module):\n", " def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):\n", " super().__init__()\n", " self.w_qs = nn.Parameter(torch.randn(d_model, n_head, d_k) * np.sqrt(2.0 / (d_model + d_k)))\n", " self.w_ks = nn.Parameter(torch.randn(d_model, n_head, d_k) * np.sqrt(2.0 / (d_model + d_k)))\n", " self.w_vs = nn.Parameter(torch.randn(d_model, n_head, d_v) * np.sqrt(2.0 / (d_model + d_v)))\n", " self.w_fc = nn.Parameter(torch.randn(d_model, n_head, d_v) * np.sqrt(2.0 / (d_model + n_head * d_v)))\n", "\n", " self.dropout = nn.Dropout(p=dropout)\n", " self.layer_norm = nn.LayerNorm(d_model)\n", "\n", " def forward(self, q, k, v, mask=None):\n", " attn = torch.einsum('bld,dhc,bte,ehc->hblt', [q, self.w_qs, k, self.w_ks])\n", " if mask is not None:\n", " attn = attn.masked_fill(mask[None], -np.inf)\n", " attn = torch.softmax(attn, dim=3)\n", " output = torch.einsum('hblt,bte,ehv,dhv->hbd', [attn, v, self.w_vs, self.w_fc])\n", " output = self.dropout(output)\n", " output = self.layer_norm(output + q)\n", " return output, attn" ] }, { "cell_type": "code", "execution_count": 56, "metadata": {}, "outputs": [], "source": [ "n_heads = 8\n", "d_k = 32\n", "d_v = 64\n", "d_model = 100\n", "t = 51\n", "l = 53\n", "batch = 30\n", "\n", "layer1 = initialize(MultiHeadAttentionOld(n_heads, d_k=d_k, d_v=d_v, d_model=d_model)).eval().cuda()\n", "layer2 = initialize(MultiHeadAttentionNew(n_heads, d_k=d_k, d_v=d_v, d_model=d_model)).eval().cuda()\n", "\n", "args = dict(\n", " q=torch.randn(batch, l, d_model),\n", " k=torch.randn(batch, t, d_model) * 0.1,\n", " v=torch.randn(batch, t, d_model),\n", " mask=torch.randn(batch, l, t) > 0,\n", ")\n", "args = {k:v.cuda() for k, v in args.items()}\n", "\n", "o1, a1 = layer1(**args)\n", "o2, a2 = layer2(**args)" ] }, { "cell_type": "code", "execution_count": 57, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(torch.Size([240, 53, 51]), torch.Size([8, 30, 53, 51]))" ] }, "execution_count": 57, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a1.shape, a2.shape" ] }, { "cell_type": "code", "execution_count": 58, "metadata": {}, "outputs": [], "source": [ "assert torch.allclose(o1, o2)" ] }, { "cell_type": "code", "execution_count": 59, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "4.82 ms ± 73.9 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)\n", "4.6 ms ± 116 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)\n" ] } ], "source": [ "%timeit -n 200 layer1(**args); torch.cuda.synchronize()\n", "%timeit -n 200 layer2(**args); torch.cuda.synchronize()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Self-attention GANs\n", "\n", "SAGANs are currently SotA for image generation, and can be simplified using same tricks.\n", "\n", "\n", "" ] }, { "cell_type": "code", "execution_count": 60, "metadata": {}, "outputs": [], "source": [ "#left\n", "class Self_Attn_Old(nn.Module):\n", " \"\"\" Self attention Layer\"\"\"\n", " def __init__(self,in_dim,activation):\n", " super(Self_Attn_Old,self).__init__()\n", " self.chanel_in = in_dim\n", " self.activation = activation\n", "\n", " self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n", " self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n", " self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)\n", " self.gamma = nn.Parameter(torch.zeros(1))\n", "\n", " self.softmax = nn.Softmax(dim=-1) #\n", "\n", " def forward(self, x):\n", " \"\"\"\n", " inputs :\n", " x : input feature maps( B X C X W X H)\n", " returns :\n", " out : self attention value + input feature\n", " attention: B X N X N (N is Width*Height)\n", " \"\"\"\n", "\n", " m_batchsize,C,width ,height = x.size()\n", " proj_query = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)\n", " proj_key = self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)\n", " energy = torch.bmm(proj_query,proj_key) # transpose check\n", " attention = self.softmax(energy) # BX (N) X (N)\n", " proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N\n", "\n", " out = torch.bmm(proj_value,attention.permute(0,2,1) )\n", " out = out.view(m_batchsize,C,width,height)\n", "\n", " out = self.gamma*out + x\n", " return out,attention" ] }, { "cell_type": "code", "execution_count": 61, "metadata": {}, "outputs": [], "source": [ "#right\n", "class Self_Attn_New(nn.Module):\n", " \"\"\" Self attention Layer\"\"\"\n", " def __init__(self, in_dim):\n", " super().__init__()\n", " self.query_conv = nn.Conv2d(in_dim, out_channels=in_dim//8, kernel_size=1)\n", " self.key_conv = nn.Conv2d(in_dim, out_channels=in_dim//8, kernel_size=1)\n", " self.value_conv = nn.Conv2d(in_dim, out_channels=in_dim, kernel_size=1)\n", " self.gamma = nn.Parameter(torch.zeros([1]))\n", "\n", " def forward(self, x):\n", " proj_query = rearrange(self.query_conv(x), 'b c h w -> b (h w) c')\n", " proj_key = rearrange(self.key_conv(x), 'b c h w -> b c (h w)')\n", " proj_value = rearrange(self.value_conv(x), 'b c h w -> b (h w) c')\n", " energy = torch.bmm(proj_query, proj_key)\n", " attention = F.softmax(energy, dim=2)\n", " out = torch.bmm(attention, proj_value)\n", " out = x + self.gamma * rearrange(out, 'b (h w) c -> b c h w',\n", " **parse_shape(x, 'b c h w'))\n", " return out, attention" ] }, { "cell_type": "code", "execution_count": 62, "metadata": {}, "outputs": [], "source": [ "model_old = initialize(Self_Attn_Old(128, None))\n", "model_new = initialize(Self_Attn_New(128))" ] }, { "cell_type": "code", "execution_count": 63, "metadata": {}, "outputs": [], "source": [ "x = torch.randn(2, 128, 30, 30)\n", "assert torch.allclose(model_old(x)[0], model_new(x)[0], atol=1e-4)\n", "# returned attention is transposed\n", "assert torch.allclose(model_old(x)[1], model_new(x)[1], atol=1e-4)" ] }, { "cell_type": "code", "execution_count": 64, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "15.9 ms ± 990 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n", "14.5 ms ± 81.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" ] } ], "source": [ "%timeit model_old(x)[0].sum().item()\n", "%timeit model_new(x)[0].sum().item()\n", "# surprise - I had slow down here due to the order of softmax, not einops" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Improving time sequence prediction\n", "\n", "\n", "\n", "While this example was considered to be simplistic, I had to analyze surrounding code to understand what kind of input was expected.\n", "You can try yourself. \n", "\n", "Additionally now the code works with any dtype, not only double; and new code supports using GPU." ] }, { "cell_type": "code", "execution_count": 65, "metadata": {}, "outputs": [], "source": [ "#left\n", "class SequencePredictionOld(nn.Module):\n", " def __init__(self):\n", " super(SequencePredictionOld, self).__init__()\n", " self.lstm1 = nn.LSTMCell(1, 51)\n", " self.lstm2 = nn.LSTMCell(51, 51)\n", " self.linear = nn.Linear(51, 1)\n", "\n", " def forward(self, input, future = 0):\n", " outputs = []\n", " h_t = torch.zeros(input.size(0), 51, dtype=torch.double)\n", " c_t = torch.zeros(input.size(0), 51, dtype=torch.double)\n", " h_t2 = torch.zeros(input.size(0), 51, dtype=torch.double)\n", " c_t2 = torch.zeros(input.size(0), 51, dtype=torch.double)\n", "\n", " for i, input_t in enumerate(input.chunk(input.size(1), dim=1)):\n", " h_t, c_t = self.lstm1(input_t, (h_t, c_t))\n", " h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))\n", " output = self.linear(h_t2)\n", " outputs += [output]\n", "\n", " for i in range(future):# if we should predict the future\n", " h_t, c_t = self.lstm1(output, (h_t, c_t))\n", " h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))\n", " output = self.linear(h_t2)\n", " outputs += [output]\n", " outputs = torch.stack(outputs, 1).squeeze(2)\n", " return outputs" ] }, { "cell_type": "code", "execution_count": 66, "metadata": {}, "outputs": [], "source": [ "#right\n", "class SequencePredictionNew(nn.Module):\n", " def __init__(self):\n", " super(SequencePredictionNew, self).__init__()\n", " self.lstm1 = nn.LSTMCell(1, 51)\n", " self.lstm2 = nn.LSTMCell(51, 51)\n", " self.linear = nn.Linear(51, 1)\n", "\n", " def forward(self, input, future=0):\n", " b, t = input.shape\n", " h_t, c_t, h_t2, c_t2 = torch.zeros(4, b, 51, dtype=self.linear.weight.dtype,\n", " device=self.linear.weight.device)\n", "\n", " outputs = []\n", " for input_t in rearrange(input, 'b t -> t b ()'):\n", " h_t, c_t = self.lstm1(input_t, (h_t, c_t))\n", " h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))\n", " output = self.linear(h_t2)\n", " outputs += [output]\n", "\n", " for i in range(future): # if we should predict the future\n", " h_t, c_t = self.lstm1(output, (h_t, c_t))\n", " h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))\n", " output = self.linear(h_t2)\n", " outputs += [output]\n", " return rearrange(outputs, 't b () -> b t')" ] }, { "cell_type": "code", "execution_count": 67, "metadata": {}, "outputs": [], "source": [ "seq_old = SequencePredictionOld().double()\n", "seq_new = SequencePredictionNew().double()\n", "initialize(seq_old)\n", "initialize(seq_new)\n", "x = torch.randn([10, 10], dtype=torch.double)" ] }, { "cell_type": "code", "execution_count": 68, "metadata": {}, "outputs": [], "source": [ "result_old = seq_old(x)\n", "result_new = seq_new(x)\n", "assert torch.allclose(result_old, result_new)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Transforming spacial transformer network (STN)\n", "\n", "\n", "" ] }, { "cell_type": "code", "execution_count": 69, "metadata": {}, "outputs": [], "source": [ "#left\n", "class SpacialTransformOld(nn.Module):\n", " def __init__(self):\n", " super(Net, self).__init__()\n", "\n", " # Spatial transformer localization-network\n", " self.localization = nn.Sequential(\n", " nn.Conv2d(1, 8, kernel_size=7),\n", " nn.MaxPool2d(2, stride=2),\n", " nn.ReLU(True),\n", " nn.Conv2d(8, 10, kernel_size=5),\n", " nn.MaxPool2d(2, stride=2),\n", " nn.ReLU(True)\n", " )\n", "\n", " # Regressor for the 3 * 2 affine matrix\n", " self.fc_loc = nn.Sequential(\n", " nn.Linear(10 * 3 * 3, 32),\n", " nn.ReLU(True),\n", " nn.Linear(32, 3 * 2)\n", " )\n", "\n", " # Initialize the weights/bias with identity transformation\n", " self.fc_loc[2].weight.data.zero_()\n", " self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))\n", "\n", " # Spatial transformer network forward function\n", " def stn(self, x):\n", " xs = self.localization(x)\n", " xs = xs.view(-1, 10 * 3 * 3)\n", " theta = self.fc_loc(xs)\n", " theta = theta.view(-1, 2, 3)\n", "\n", " grid = F.affine_grid(theta, x.size())\n", " x = F.grid_sample(x, grid)\n", "\n", " return x\n" ] }, { "cell_type": "code", "execution_count": 70, "metadata": {}, "outputs": [], "source": [ "#right\n", "class SpacialTransformNew(nn.Module):\n", " def __init__(self):\n", " super(Net, self).__init__()\n", " # Spatial transformer localization-network\n", " linear = nn.Linear(32, 3 * 2)\n", " # Initialize the weights/bias with identity transformation\n", " linear.weight.data.zero_()\n", " linear.bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))\n", "\n", " self.compute_theta = nn.Sequential(\n", " nn.Conv2d(1, 8, kernel_size=7),\n", " nn.MaxPool2d(2, stride=2),\n", " nn.ReLU(True),\n", " nn.Conv2d(8, 10, kernel_size=5),\n", " nn.MaxPool2d(2, stride=2),\n", " nn.ReLU(True),\n", " Rearrange('b c h w -> b (c h w)', h=3, w=3),\n", " nn.Linear(10 * 3 * 3, 32),\n", " nn.ReLU(True),\n", " linear,\n", " Rearrange('b (row col) -> b row col', row=2, col=3),\n", " )\n", "\n", " # Spatial transformer network forward function\n", " def stn(self, x):\n", " grid = F.affine_grid(self.compute_theta(x), x.size())\n", " return F.grid_sample(x, grid)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "- new code will give reasonable errors when passed image size is different from expected\n", "- if batch size is divisible by 18, whatever you input in the old code, it'll fail no sooner than affine_grid." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Improving GLOW\n", "\n", "That's a good old depth-to-space written manually!\n", "\n", "Since GLOW is revertible, it will frequently rely on `rearrange`-like operations.\n", "\n", "" ] }, { "cell_type": "code", "execution_count": 71, "metadata": {}, "outputs": [], "source": [ "#left\n", "def unsqueeze2d_old(input, factor=2):\n", " assert factor >= 1 and isinstance(factor, int)\n", " factor2 = factor ** 2\n", " if factor == 1:\n", " return input\n", " size = input.size()\n", " B = size[0]\n", " C = size[1]\n", " H = size[2]\n", " W = size[3]\n", " assert C % (factor2) == 0, \"{}\".format(C)\n", " x = input.view(B, C // factor2, factor, factor, H, W)\n", " x = x.permute(0, 1, 4, 2, 5, 3).contiguous()\n", " x = x.view(B, C // (factor2), H * factor, W * factor)\n", " return x\n", "\n", "def squeeze2d_old(input, factor=2):\n", " assert factor >= 1 and isinstance(factor, int)\n", " if factor == 1:\n", " return input\n", " size = input.size()\n", " B = size[0]\n", " C = size[1]\n", " H = size[2]\n", " W = size[3]\n", " assert H % factor == 0 and W % factor == 0, \"{}\".format((H, W))\n", " x = input.view(B, C, H // factor, factor, W // factor, factor)\n", " x = x.permute(0, 1, 3, 5, 2, 4).contiguous()\n", " x = x.view(B, C * factor * factor, H // factor, W // factor)\n", " return x" ] }, { "cell_type": "code", "execution_count": 72, "metadata": {}, "outputs": [], "source": [ "#right\n", "def unsqueeze2d_new(input, factor=2):\n", " return rearrange(input, 'b (c h2 w2) h w -> b c (h h2) (w w2)', h2=factor, w2=factor)\n", "\n", "def squeeze2d_new(input, factor=2):\n", " return rearrange(input, 'b c (h h2) (w w2) -> b (c h2 w2) h w', h2=factor, w2=factor)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "- term `squeeze` isn't very helpful: which dimension is squeezed? There is `torch.squeeze`, but it's very different.\n", "- in fact, we could skip creating functions completely - it is a single call to `einops` anyway" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Detecting problems in YOLO detection\n", "\n", "" ] }, { "cell_type": "code", "execution_count": 73, "metadata": {}, "outputs": [], "source": [ "#left\n", "def YOLO_prediction_old(input, num_classes, num_anchors, anchors, stride_h, stride_w):\n", " bs = input.size(0)\n", " in_h = input.size(2)\n", " in_w = input.size(3)\n", " scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors]\n", "\n", " prediction = input.view(bs, num_anchors,\n", " 5 + num_classes, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()\n", " # Get outputs\n", " x = torch.sigmoid(prediction[..., 0]) # Center x\n", " y = torch.sigmoid(prediction[..., 1]) # Center y\n", " w = prediction[..., 2] # Width\n", " h = prediction[..., 3] # Height\n", " conf = torch.sigmoid(prediction[..., 4]) # Conf\n", " pred_cls = torch.sigmoid(prediction[..., 5:]) # Cls pred.\n", "\n", " FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor\n", " LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor\n", " # Calculate offsets for each grid\n", " grid_x = torch.linspace(0, in_w - 1, in_w).repeat(in_w, 1).repeat(\n", " bs * num_anchors, 1, 1).view(x.shape).type(FloatTensor)\n", " grid_y = torch.linspace(0, in_h - 1, in_h).repeat(in_h, 1).t().repeat(\n", " bs * num_anchors, 1, 1).view(y.shape).type(FloatTensor)\n", " # Calculate anchor w, h\n", " anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))\n", " anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))\n", " anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape)\n", " anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape)\n", " # Add offset and scale with anchors\n", " pred_boxes = FloatTensor(prediction[..., :4].shape)\n", " pred_boxes[..., 0] = x.data + grid_x\n", " pred_boxes[..., 1] = y.data + grid_y\n", " pred_boxes[..., 2] = torch.exp(w.data) * anchor_w\n", " pred_boxes[..., 3] = torch.exp(h.data) * anchor_h\n", " # Results\n", " _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)\n", " output = torch.cat((pred_boxes.view(bs, -1, 4) * _scale,\n", " conf.view(bs, -1, 1), pred_cls.view(bs, -1, num_classes)), -1)\n", " return output" ] }, { "cell_type": "code", "execution_count": 74, "metadata": {}, "outputs": [], "source": [ "#right\n", "def YOLO_prediction_new(input, num_classes, num_anchors, anchors, stride_h, stride_w):\n", " raw_predictions = rearrange(input, 'b (anchor prediction) h w -> prediction b anchor h w',\n", " anchor=num_anchors, prediction=5 + num_classes)\n", " anchors = torch.FloatTensor(anchors).to(input.device)\n", " anchor_sizes = rearrange(anchors, 'anchor dim -> dim () anchor () ()')\n", "\n", " _, _, _, in_h, in_w = raw_predictions.shape\n", " grid_h = rearrange(torch.arange(in_h).float(), 'h -> () () h ()').to(input.device)\n", " grid_w = rearrange(torch.arange(in_w).float(), 'w -> () () () w').to(input.device)\n", "\n", " predicted_bboxes = torch.zeros_like(raw_predictions)\n", " predicted_bboxes[0] = (raw_predictions[0].sigmoid() + grid_w) * stride_w # center x\n", " predicted_bboxes[1] = (raw_predictions[1].sigmoid() + grid_h) * stride_h # center y\n", " predicted_bboxes[2:4] = (raw_predictions[2:4].exp()) * anchor_sizes # bbox width and height\n", " predicted_bboxes[4] = raw_predictions[4].sigmoid() # confidence\n", " predicted_bboxes[5:] = raw_predictions[5:].sigmoid() # class predictions\n", " # merging all predicted bboxes for each image\n", " return rearrange(predicted_bboxes, 'prediction b anchor h w -> b (anchor h w) prediction')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We changed and fixed a lot:\n", "\n", "- new code won't fail if input is not on the first GPU\n", "- old code has wrong grid_x and grid_y for non-square images\n", "- new code doesn't use replication when broadcasting is sufficient\n", "- old code strangely sometimes takes `.data`, but this has no real effect, as some branches preserve gradient till the end\n", " - if gradients not needed, torch.no_grad should be used, so it's redundant" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Simpler output for a bunch of pictures\n", "\n", "Next time you need to output drawings of you generative models, you can use this trick\n", "\n", "" ] }, { "cell_type": "code", "execution_count": 75, "metadata": {}, "outputs": [], "source": [ "fake_batch = torch.rand([100, 3, 1, 1]) + torch.zeros([100, 3, 32, 32])\n", "from matplotlib import pyplot as plt\n", "import torchvision.utils as vutils" ] }, { "cell_type": "code", "execution_count": 76, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "
{}
".format(highlight(trimmed_source, PythonLexer(), HtmlFormatter()))
if source.startswith("#right"):
content += "{}
{}
".format(markdown.markdown(cell["source"]))
else:
raise RuntimeError("not expected type of cell" + cell["cell_type"])
styles = HtmlFormatter().get_style_defs(".highlight")
styles += """
body {
padding: 50px 10px;
}
.leftright-wrapper {
text-align: center;
overflow-x: auto;
}
.leftright-cells {
display: inline-flex;
text-align: left;
}
.leftright-cells > div {
padding: 0px 10px;
min-width: 350px;
}
.markdown-cell{
max-width: 700px;
margin: 0px auto;
}
h1 {
text-align: center;
padding: 10px 0px 0px;
}
"""
meta_tags = """
"""
github_ribbon = """
"""
result = f"""
{meta_tags}