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