{"cells": [{"cell_type": "markdown", "metadata": {}, "source": ["# Einsum decomposition\n", "\n", "This notebook shows a way to decompose [einsum](https://numpy.org/doc/stable/reference/generated/numpy.einsum.html) into a subset of operations (expand_dims, squeeze, transpose, extended matrix multiplication)."]}, {"cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 2, "metadata": {}, "output_type": "execute_result"}], "source": ["from jyquickhelper import add_notebook_menu\n", "add_notebook_menu()"]}, {"cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": ["%load_ext mlprodict"]}, {"cell_type": "markdown", "metadata": {}, "source": ["## Operator explanation with equation `bac,cd,def=ebc`\n", "\n", "The operator einsum takes an equation and some inputs. Every letter involved in the equation is a loop. Let's see on one example."]}, {"cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [{"data": {"text/plain": ["array([[[ 8866198., 9864696.],\n", " [12090270., 13152928.]],\n", "\n", " [[ 8883886., 9884376.],\n", " [12114390., 13179168.]]], dtype=float32)"]}, "execution_count": 4, "metadata": {}, "output_type": "execute_result"}], "source": ["import numpy\n", "\n", "m1 = numpy.arange(0, 8).astype(numpy.float32).reshape((2, 2, 2)) + 10\n", "m2 = numpy.arange(0, 4).astype(numpy.float32).reshape((2, 2)) + 100\n", "m3 = numpy.arange(0, 8).astype(numpy.float32).reshape((2, 2, 2)) + 1000\n", "\n", "equation = \"bac,cd,def->ebc\"\n", "truth = numpy.einsum(equation, m1, m2, m3)\n", "truth"]}, {"cell_type": "markdown", "metadata": {}, "source": ["This summation is equalent to:"]}, {"cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [{"data": {"text/plain": ["array([[[ 8866198., 9864696.],\n", " [12090270., 13152928.]],\n", "\n", " [[ 8883886., 9884376.],\n", " [12114390., 13179168.]]])"]}, "execution_count": 5, "metadata": {}, "output_type": "execute_result"}], "source": ["res = numpy.zeros((2, 2, 2))\n", "for a in range(0, 2):\n", " for b in range(0, 2):\n", " for c in range(0, 2):\n", " for d in range(0, 2):\n", " for e in range(0, 2):\n", " for f in range(0, 2):\n", " res[e, b, c] += m1[b, a, c] * m2[c, d] * m3[d, e, f]\n", "res"]}, {"cell_type": "markdown", "metadata": {}, "source": ["Theoritically, this summation is in this case has a cost of $O(N^6)$. However this simple computation is usually much longer than using matrix multiplications along the path. $O(N^4)$ is the cost of the heaviest matrix multiplication in this case). But to do that, the equation needs to be decomposed into a sequence of matrix multiplications."]}, {"cell_type": "markdown", "metadata": {}, "source": ["### Decomposition of `bac,cd,def=ebc`"]}, {"cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": ["import numpy\n", "from mlprodict.testing.einsum import (\n", " decompose_einsum_equation, apply_einsum_sequence)"]}, {"cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": ["m1 = numpy.arange(0, 8).astype(numpy.float32).reshape((2, 2, 2)) + 10\n", "m2 = numpy.arange(0, 4).astype(numpy.float32).reshape((2, 2)) + 100\n", "m3 = numpy.arange(0, 8).astype(numpy.float32).reshape((2, 2, 2)) + 1000"]}, {"cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": ["seq = decompose_einsum_equation(\"bac,cd,def->ebc\")"]}, {"cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 9, "metadata": {}, "output_type": "execute_result"}], "source": ["from jyquickhelper import RenderJsDot\n", "RenderJsDot(seq.to_dot(size=7))"]}, {"cell_type": "markdown", "metadata": {}, "source": ["Then the result can be obtained as follows:"]}, {"cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [{"data": {"text/plain": ["array([[[ 8866198., 9864696.],\n", " [12090270., 13152928.]],\n", "\n", " [[ 8883886., 9884376.],\n", " [12114390., 13179168.]]], dtype=float32)"]}, "execution_count": 10, "metadata": {}, "output_type": "execute_result"}], "source": ["apply_einsum_sequence(seq, m1, m2, m3)"]}, {"cell_type": "markdown", "metadata": {}, "source": ["### operator matmul\n", "\n", "This operator can be used to represent either a multiplication, either a matrix multiplication but it applies only on arrays with the same number of dimensions. It can be broken into multiplication of matrix multiplication."]}, {"cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 11, "metadata": {}, "output_type": "execute_result"}], "source": ["seq_clean = decompose_einsum_equation(\"bac,cd,def->ebc\", strategy='numpy', clean=True)\n", "RenderJsDot(seq_clean.to_dot(size=7))"]}, {"cell_type": "markdown", "metadata": {}, "source": ["Operator *transpose_mm* is a regular transposition, it takes two inputs but only tranposes the first input before returning it. Operator *batch_dot* is a matrix multiplication. It is left that way on purpose as it may be implemented with function dot or gemm. The operator distinguishes between 3 kind of axes: batch axes, kept axes, sum(mation) axes. It then reshapes both input matrices with 3D tensors, batch axis, row axis, column axis to use function [numpy.dot](https://numpy.org/doc/stable/reference/generated/numpy.dot.html)."]}, {"cell_type": "markdown", "metadata": {}, "source": ["### ONNX\n", "\n", "The previous graph can be converted into ONNX."]}, {"cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 12, "metadata": {}, "output_type": "execute_result"}], "source": ["onx = seq_clean.to_onnx(\"Y\", \"X1\", \"X2\", \"X3\", dtype=numpy.float32)\n", "# with open(\"einsum.onnx\", \"wb\") as f:\n", "# f.write(onx.SerializeToString())\n", "%onnxview onx "]}, {"cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [{"data": {"text/plain": ["array([[[ 8866198., 9864696.],\n", " [12090270., 13152928.]],\n", "\n", " [[ 8883886., 9884376.],\n", " [12114390., 13179168.]]], dtype=float32)"]}, "execution_count": 13, "metadata": {}, "output_type": "execute_result"}], "source": ["from onnxruntime import InferenceSession\n", "sess = InferenceSession(onx.SerializeToString())\n", "sess.run(None, {'X1': m1.astype(numpy.float32), \n", " 'X2': m2.astype(numpy.float32), \n", " 'X3': m3.astype(numpy.float32)})[0]"]}, {"cell_type": "markdown", "metadata": {}, "source": ["### onnxruntime"]}, {"cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [], "source": ["import onnx\n", "from onnx import helper, numpy_helper\n", "from onnxruntime import InferenceSession\n", "\n", "\n", "def make_model1(equation):\n", " model = helper.make_model(\n", " opset_imports=[helper.make_operatorsetid('', 13)],\n", " graph=helper.make_graph(\n", " name='einsum_test',\n", " inputs=[helper.make_tensor_value_info(\"X\", onnx.TensorProto.FLOAT, None),\n", " helper.make_tensor_value_info(\"Y\", onnx.TensorProto.FLOAT, None),\n", " helper.make_tensor_value_info(\"Z\", onnx.TensorProto.FLOAT, None)],\n", " outputs=[helper.make_tensor_value_info(\"A\", onnx.TensorProto.FLOAT, None)],\n", " nodes=[\n", " helper.make_node(\"Einsum\", [\"X\", \"Y\", \"Z\"], [\"A\"], equation=equation)\n", " ]\n", " )\n", " )\n", " return model\n", "\n", "\n", "model = make_model1(\"bac,cd,def->ebc\")\n", "sess = InferenceSession(model.SerializeToString())"]}, {"cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [{"data": {"text/plain": ["array([[[ 8866198., 9864696.],\n", " [12090270., 13152928.]],\n", "\n", " [[ 8883886., 9884376.],\n", " [12114390., 13179168.]]], dtype=float32)"]}, "execution_count": 15, "metadata": {}, "output_type": "execute_result"}], "source": ["sess.run(None, {'X': m1.astype(numpy.float32), \n", " 'Y': m2.astype(numpy.float32), \n", " 'Z': m3.astype(numpy.float32)})[0]"]}, {"cell_type": "markdown", "metadata": {}, "source": ["### Benchmark\n", "\n", "It clearly shows the summation done with the basic algorithm is the slowest."]}, {"cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["C:\\xavierdupre\\__home_\\github_fork\\scikit-learn\\sklearn\\experimental\\enable_hist_gradient_boosting.py:16: UserWarning: Since version 1.0, it is not needed to import enable_hist_gradient_boosting anymore. HistGradientBoostingClassifier and HistGradientBoostingRegressor are now stable and can be normally imported from sklearn.ensemble.\n", " warnings.warn(\n", "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 14/14 [00:20<00:00, 1.47s/it]\n"]}, {"data": {"text/html": ["\n", "\n", "
\n", " \n", " \n", " average \n", " deviation \n", " min_exec \n", " max_exec \n", " repeat \n", " number \n", " total \n", " name \n", " N \n", " \n", " \n", " \n", " \n", " 82 \n", " 0.065132 \n", " 0.001338 \n", " 0.063801 \n", " 0.068927 \n", " 10 \n", " 10 \n", " 0.651318 \n", " custom_einsum \n", " 60 \n", " \n", " \n", " 83 \n", " 0.051615 \n", " 0.001206 \n", " 0.049987 \n", " 0.053465 \n", " 10 \n", " 10 \n", " 0.516154 \n", " dec-matmul \n", " 60 \n", " \n", " \n", " 84 \n", " 0.062689 \n", " 0.003658 \n", " 0.058949 \n", " 0.073073 \n", " 10 \n", " 10 \n", " 0.626888 \n", " dec-batch_dot \n", " 60 \n", " \n", " \n", " 85 \n", " 0.009917 \n", " 0.000274 \n", " 0.009737 \n", " 0.010686 \n", " 10 \n", " 10 \n", " 0.099166 \n", " ort-einsum \n", " 60 \n", " \n", " \n", " 86 \n", " 0.015518 \n", " 0.001107 \n", " 0.014413 \n", " 0.018179 \n", " 10 \n", " 10 \n", " 0.155178 \n", " ort-matmul \n", " 60 \n", " \n", " \n", "
\n", "
"]}, "metadata": {"needs_background": "light"}, "output_type": "display_data"}], "source": ["import matplotlib.pyplot as plt\n", "\n", "piv = df.pivot(\"N\", \"name\", \"average\")\n", "piv2 = piv.copy()\n", "np = piv[\"numpy.einsum\"]\n", "for c in piv2.columns:\n", " piv2[c] /= np\n", " \n", "fig, ax = plt.subplots(1, 2, figsize=(14, 6))\n", "piv.plot(logy=True, logx=True, ax=ax[0])\n", "ax[0].set_title(\"Benchmark einsum function\\nbac,cd,def->ebc\")\n", "piv2.plot(logy=True, logx=True, ax=ax[1])\n", "ax[1].set_title(\"Benchmark einsum function\\n(ratio, baseline=numpy)\");"]}, {"cell_type": "markdown", "metadata": {}, "source": ["Version `dec-matmul` is an implementation based on the decomposition of a simplified einsum into a sequence of transpose, reshape, (batch_)dot or mul operations. This decomposition is converted into ONNX and executed with *onnxruntime*, version `ort-matmul`. Both versions are faster than the numpy optimized version."]}, {"cell_type": "markdown", "metadata": {}, "source": ["## Another example with `bsnh,btnh=bnts`\n", "\n", "Another case, more frequent in deep learning."]}, {"cell_type": "markdown", "metadata": {}, "source": ["### Decomposition of `bsnh,btnh=bnts`"]}, {"cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 18, "metadata": {}, "output_type": "execute_result"}], "source": ["seq2 = decompose_einsum_equation(\"bsnh,btnh->bnts\", strategy='numpy', clean=True)\n", "RenderJsDot(seq2.to_dot(size=7))"]}, {"cell_type": "markdown", "metadata": {}, "source": ["### ONNX version"]}, {"cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 19, "metadata": {}, "output_type": "execute_result"}], "source": ["onx2 = seq2.to_onnx(\"Y\", \"X1\", \"X2\", dtype=numpy.float32)\n", "%onnxview onx2 "]}, {"cell_type": "markdown", "metadata": {}, "source": ["### Benchmark"]}, {"cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 7/7 [00:13<00:00, 1.93s/it]\n"]}, {"data": {"text/html": ["\n", "\n", "
\n", " \n", " \n", " average \n", " deviation \n", " min_exec \n", " max_exec \n", " repeat \n", " number \n", " total \n", " name \n", " N \n", " \n", " \n", " \n", " \n", " 37 \n", " 0.229418 \n", " 0.020792 \n", " 0.217997 \n", " 0.291032 \n", " 10 \n", " 10 \n", " 2.294175 \n", " custom_einsum \n", " 40 \n", " \n", " \n", " 38 \n", " 0.160575 \n", " 0.005435 \n", " 0.150772 \n", " 0.167411 \n", " 10 \n", " 10 \n", " 1.605746 \n", " dec-matmul \n", " 40 \n", " \n", " \n", " 39 \n", " 0.112844 \n", " 0.011305 \n", " 0.102173 \n", " 0.141890 \n", " 10 \n", " 10 \n", " 1.128436 \n", " dec-batch_dot \n", " 40 \n", " \n", " \n", " 40 \n", " 0.051181 \n", " 0.003533 \n", " 0.047244 \n", " 0.057054 \n", " 10 \n", " 10 \n", " 0.511815 \n", " ort-einsum \n", " 40 \n", " \n", " \n", " 41 \n", " 0.078827 \n", " 0.008735 \n", " 0.067893 \n", " 0.099156 \n", " 10 \n", " 10 \n", " 0.788271 \n", " ort-matmul \n", " 40 \n", " \n", " \n", "
\n", "
"]}, "metadata": {"needs_background": "light"}, "output_type": "display_data"}], "source": ["piv = df.pivot(\"N\", \"name\", \"average\")\n", "piv2 = piv.copy()\n", "np = piv[\"numpy.einsum\"]\n", "for c in piv2.columns:\n", " piv2[c] /= np\n", " \n", "fig, ax = plt.subplots(1, 2, figsize=(14, 6))\n", "piv.plot(logy=True, logx=True, ax=ax[0])\n", "ax[0].set_title(\"Benchmark einsum function\\nbsnh,btnh->bnts\")\n", "piv2.plot(logy=True, logx=True, ax=ax[1])\n", "ax[1].set_title(\"Benchmark einsum function\\n(ratio, baseline=numpy)\");"]}, {"cell_type": "markdown", "metadata": {}, "source": ["### Permutation\n", "\n", "Einsum's algorithm started by aligning all matrices involved in the computation to the same dimension in the same order. But which order is the best, that's the question."]}, {"cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [{"data": {"text/plain": ["['b', 'h', 'n', 's', 't']"]}, "execution_count": 22, "metadata": {}, "output_type": "execute_result"}], "source": ["equation = \"bsnh,btnh->bnts\"\n", "letters = list(sorted(set([c for c in equation if \"a\" <= c < \"z\"])))\n", "letters"]}, {"cell_type": "code", "execution_count": 22, "metadata": {"scrolled": false}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 120/120 [00:11<00:00, 10.23it/s]\n"]}, {"data": {"text/html": ["\n", "\n", "
\n", " \n", " \n", " average \n", " deviation \n", " min_exec \n", " max_exec \n", " repeat \n", " number \n", " total \n", " name \n", " N \n", " eq \n", " transpose \n", " kind \n", " gemm \n", " \n", " \n", " \n", " \n", " 715 \n", " 0.006162 \n", " 0.000038 \n", " 0.006128 \n", " 0.006216 \n", " 3 \n", " 4 \n", " 0.018485 \n", " custom_einsum \n", " 15 \n", " thns,tbns->tnbh \n", " 3.0 \n", " NN \n", " NaN \n", " \n", " \n", " 716 \n", " 0.002343 \n", " 0.000046 \n", " 0.002294 \n", " 0.002405 \n", " 3 \n", " 4 \n", " 0.007029 \n", " dec-matmul \n", " 15 \n", " thns,tbns->tnbh \n", " 3.0 \n", " NN \n", " NaN \n", " \n", " \n", " 717 \n", " 0.001645 \n", " 0.000035 \n", " 0.001610 \n", " 0.001694 \n", " 3 \n", " 4 \n", " 0.004934 \n", " dec-batch_dot \n", " 15 \n", " thns,tbns->tnbh \n", " 3.0 \n", " NN \n", " NaN \n", " \n", " \n", " 718 \n", " 0.000833 \n", " 0.000015 \n", " 0.000820 \n", " 0.000853 \n", " 3 \n", " 4 \n", " 0.002498 \n", " ort-einsum \n", " 15 \n", " thns,tbns->tnbh \n", " 4.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", " 719 \n", " 0.001251 \n", " 0.000012 \n", " 0.001238 \n", " 0.001268 \n", " 3 \n", " 4 \n", " 0.003753 \n", " ort-matmul \n", " 15 \n", " thns,tbns->tnbh \n", " 4.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", "
\n", "
\n", "\n", "
\n", " \n", " \n", " average \n", " deviation \n", " min_exec \n", " max_exec \n", " repeat \n", " number \n", " total \n", " name \n", " N \n", " eq \n", " transpose \n", " kind \n", " gemm \n", " \n", " \n", " \n", " \n", " 0 \n", " 0.000758 \n", " 0.000015 \n", " 0.000738 \n", " 0.000771 \n", " 3 \n", " 4 \n", " 0.002275 \n", " ort-matmul \n", " 15 \n", " hsnt,hbnt->hnbs \n", " 4.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", " 1 \n", " 0.000770 \n", " 0.000023 \n", " 0.000739 \n", " 0.000793 \n", " 3 \n", " 4 \n", " 0.002310 \n", " ort-matmul \n", " 15 \n", " hnts,hbts->htbn \n", " 4.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", " 2 \n", " 0.000778 \n", " 0.000020 \n", " 0.000758 \n", " 0.000806 \n", " 3 \n", " 4 \n", " 0.002334 \n", " ort-matmul \n", " 15 \n", " bnst,bhst->bshn \n", " 4.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", " 3 \n", " 0.000783 \n", " 0.000021 \n", " 0.000760 \n", " 0.000812 \n", " 3 \n", " 4 \n", " 0.002350 \n", " ort-matmul \n", " 15 \n", " bnht,bsht->bhsn \n", " 4.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", " 4 \n", " 0.000784 \n", " 0.000011 \n", " 0.000774 \n", " 0.000799 \n", " 3 \n", " 4 \n", " 0.002351 \n", " ort-matmul \n", " 15 \n", " hnst,hbst->hsbn \n", " 4.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", "
\n", "
\n", "\n", "
\n", " \n", " \n", " average \n", " deviation \n", " min_exec \n", " max_exec \n", " repeat \n", " number \n", " total \n", " name \n", " N \n", " eq \n", " transpose \n", " kind \n", " gemm \n", " \n", " \n", " \n", " \n", " 715 \n", " 0.011529 \n", " 0.000882 \n", " 0.010456 \n", " 0.012617 \n", " 3 \n", " 4 \n", " 0.034587 \n", " custom_einsum \n", " 15 \n", " sbnt,shnt->snhb \n", " 3.0 \n", " NN \n", " NaN \n", " \n", " \n", " 716 \n", " 0.011548 \n", " 0.000422 \n", " 0.010967 \n", " 0.011953 \n", " 3 \n", " 4 \n", " 0.034644 \n", " custom_einsum \n", " 15 \n", " htsb,hnsb->hsnt \n", " 3.0 \n", " NN \n", " NaN \n", " \n", " \n", " 717 \n", " 0.013971 \n", " 0.001984 \n", " 0.012279 \n", " 0.016754 \n", " 3 \n", " 4 \n", " 0.041912 \n", " custom_einsum \n", " 15 \n", " nbsh,ntsh->nstb \n", " 3.0 \n", " NN \n", " NaN \n", " \n", " \n", " 718 \n", " 0.014765 \n", " 0.001483 \n", " 0.013366 \n", " 0.016818 \n", " 3 \n", " 4 \n", " 0.044295 \n", " numpy.einsum \n", " 15 \n", " bnsh,btsh->bstn \n", " NaN \n", " NaN \n", " NaN \n", " \n", " \n", " 719 \n", " 0.015813 \n", " 0.002921 \n", " 0.012546 \n", " 0.019636 \n", " 3 \n", " 4 \n", " 0.047438 \n", " numpy.einsum \n", " 15 \n", " nbsh,ntsh->nstb \n", " NaN \n", " NaN \n", " NaN \n", " \n", " \n", "
\n", "
"]}, "metadata": {"needs_background": "light"}, "output_type": "display_data"}], "source": ["piv = df.pivot(\"eq\", \"name\", \"average\").sort_values(\"numpy.einsum\")\n", " \n", "fig, ax = plt.subplots(1, 1, figsize=(14, 6))\n", "piv.plot(logy=True, logx=True, ax=ax)\n", "ax.set_title(\"Benchmark einsum function - bsnh,btnh->bnts\");"]}, {"cell_type": "code", "execution_count": 26, "metadata": {}, "outputs": [{"data": {"text/plain": ["({3.0, 4.0}, {0.0}, {'NN'})"]}, "execution_count": 27, "metadata": {}, "output_type": "execute_result"}], "source": ["set(df['transpose'].dropna()), set(df['gemm'].dropna()), set(df['kind'].dropna())"]}, {"cell_type": "markdown", "metadata": {}, "source": ["## Decomposition of `bsnh,ctnh=nts`"]}, {"cell_type": "code", "execution_count": 27, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 28, "metadata": {}, "output_type": "execute_result"}], "source": ["seq3 = decompose_einsum_equation(\"bsnh,ctnh->nts\", strategy='numpy', clean=True)\n", "RenderJsDot(seq3.to_dot(size=7))"]}, {"cell_type": "code", "execution_count": 28, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 29, "metadata": {}, "output_type": "execute_result"}], "source": ["onx3 = seq3.to_onnx(\"Y\", \"X1\", \"X2\", dtype=numpy.float32)\n", "%onnxview onx3"]}, {"cell_type": "markdown", "metadata": {}, "source": ["### Benchmark size"]}, {"cell_type": "code", "execution_count": 29, "metadata": {}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 7/7 [00:39<00:00, 5.71s/it]\n"]}, {"data": {"text/html": ["\n", "\n", "
\n", " \n", " \n", " average \n", " deviation \n", " min_exec \n", " max_exec \n", " repeat \n", " number \n", " total \n", " name \n", " N \n", " \n", " \n", " \n", " \n", " 37 \n", " 0.043389 \n", " 0.016879 \n", " 0.030195 \n", " 0.077480 \n", " 10 \n", " 10 \n", " 0.433885 \n", " custom_einsum \n", " 40 \n", " \n", " \n", " 38 \n", " 0.015310 \n", " 0.000222 \n", " 0.014909 \n", " 0.015622 \n", " 10 \n", " 10 \n", " 0.153098 \n", " dec-matmul \n", " 40 \n", " \n", " \n", " 39 \n", " 0.013508 \n", " 0.000425 \n", " 0.013148 \n", " 0.014576 \n", " 10 \n", " 10 \n", " 0.135085 \n", " dec-batch_dot \n", " 40 \n", " \n", " \n", " 40 \n", " 0.032725 \n", " 0.000266 \n", " 0.032409 \n", " 0.033212 \n", " 10 \n", " 10 \n", " 0.327254 \n", " ort-einsum \n", " 40 \n", " \n", " \n", " 41 \n", " 0.057384 \n", " 0.002703 \n", " 0.053734 \n", " 0.062845 \n", " 10 \n", " 10 \n", " 0.573841 \n", " ort-matmul \n", " 40 \n", " \n", " \n", "
\n", "
"]}, "metadata": {"needs_background": "light"}, "output_type": "display_data"}], "source": ["piv = df.pivot(\"N\", \"name\", \"average\")\n", "piv2 = piv.copy()\n", "np = piv[\"numpy.einsum\"]\n", "for c in piv2.columns:\n", " piv2[c] /= np\n", " \n", "fig, ax = plt.subplots(1, 2, figsize=(14, 6))\n", "piv.plot(logy=True, logx=True, ax=ax[0])\n", "ax[0].set_title(\"Benchmark einsum function\\nbsnh,ctnh->nts\")\n", "piv2.plot(logy=True, logx=True, ax=ax[1])\n", "ax[1].set_title(\"Benchmark einsum function\\n(ratio, baseline=numpy)\");"]}, {"cell_type": "markdown", "metadata": {}, "source": ["### Benchmark permutation"]}, {"cell_type": "code", "execution_count": 31, "metadata": {}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 120/120 [00:06<00:00, 17.41it/s]\n"]}], "source": ["df = benchmark_perm(\"bsnh,ctnh->nts\", number=2, repeat=3, N=10)"]}, {"cell_type": "code", "execution_count": 32, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", "\n", "
\n", " \n", " \n", " average \n", " deviation \n", " min_exec \n", " max_exec \n", " repeat \n", " number \n", " total \n", " name \n", " N \n", " eq \n", " transpose \n", " kind \n", " gemm \n", " \n", " \n", " \n", " \n", " 0 \n", " 0.000125 \n", " 0.000008 \n", " 0.000118 \n", " 0.000136 \n", " 3 \n", " 2 \n", " 0.000374 \n", " ort-matmul \n", " 10 \n", " bnst,chst->shn \n", " 4.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", " 1 \n", " 0.000126 \n", " 0.000007 \n", " 0.000119 \n", " 0.000136 \n", " 3 \n", " 2 \n", " 0.000377 \n", " ort-matmul \n", " 10 \n", " bhst,cnst->snh \n", " 4.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", " 2 \n", " 0.000141 \n", " 0.000006 \n", " 0.000136 \n", " 0.000150 \n", " 3 \n", " 2 \n", " 0.000422 \n", " ort-matmul \n", " 10 \n", " hbst,cnst->snb \n", " 5.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", " 3 \n", " 0.000141 \n", " 0.000007 \n", " 0.000135 \n", " 0.000151 \n", " 3 \n", " 2 \n", " 0.000423 \n", " ort-matmul \n", " 10 \n", " nbst,chst->shb \n", " 5.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", " 4 \n", " 0.000144 \n", " 0.000007 \n", " 0.000138 \n", " 0.000154 \n", " 3 \n", " 2 \n", " 0.000432 \n", " ort-matmul \n", " 10 \n", " btns,chns->nht \n", " 5.0 \n", " NaN \n", " 0.0 \n", " \n", " \n", "
\n", "
"]}, "metadata": {"needs_background": "light"}, "output_type": "display_data"}], "source": ["piv = df.pivot(\"eq\", \"name\", \"average\").sort_values(\"numpy.einsum\")\n", " \n", "fig, ax = plt.subplots(1, 1, figsize=(14, 6))\n", "piv.plot(logy=True, logx=True, ax=ax)\n", "ax.set_title(\"Benchmark einsum function\");"]}, {"cell_type": "markdown", "metadata": {}, "source": ["### Best permutation"]}, {"cell_type": "markdown", "metadata": {}, "source": ["One of the best permutation is `bnst,chst->shn`."]}, {"cell_type": "code", "execution_count": 35, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 36, "metadata": {}, "output_type": "execute_result"}], "source": ["seq4 = decompose_einsum_equation(\"bnst,chst->shn\", strategy='numpy', clean=True)\n", "RenderJsDot(seq4.to_dot(size=7))"]}, {"cell_type": "code", "execution_count": 36, "metadata": {}, "outputs": [{"data": {"text/html": ["\n", ""], "text/plain": [""]}, "execution_count": 37, "metadata": {}, "output_type": "execute_result"}], "source": ["onx4 = seq4.to_onnx(\"Y\", \"X1\", \"X2\", dtype=numpy.float32)\n", "%onnxview onx4"]}, {"cell_type": "code", "execution_count": 37, "metadata": {}, "outputs": [], "source": []}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"codemirror_mode": {"name": "ipython", "version": 3}, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.7"}}, "nbformat": 4, "nbformat_minor": 4}