Note
Click here to download the full example code
Compares implementations of Einsum#
This example compares different equations for function numpy.einsum.
It compares numpy implementation to a custom implementation,
onnxruntime implementation and opt-einsum optimisation.
If available, tensorflow and pytorch are included as well.
The custom implementation does not do any transpose.
It uses parallelisation and SIMD optimization when the summation
happens on the last axis of both matrices. It only implements
matrix multiplication. We also measure the improvment made with
function einsum.
Available optimisation#
The code shows which optimisation is used for the custom implementation, AVX or SSE and the number of available processors, equal to the default number of used threads to parallelize.
import numpy
import pandas
import matplotlib.pyplot as plt
from onnxruntime import InferenceSession
from skl2onnx.common.data_types import FloatTensorType
from skl2onnx.algebra.onnx_ops import OnnxEinsum
from cpyquickhelper.numbers import measure_time
from tqdm import tqdm
from opt_einsum import contract
from mlprodict.testing.experimental_c_impl.experimental_c import (
custom_einsum_float, code_optimisation)
from mlprodict.testing.einsum.einsum_fct import _einsum
print(code_optimisation())
Out:
AVX-omp=8
Einsum: common code#
try:
from tensorflow import einsum as tf_einsum, convert_to_tensor
except ImportError:
tf_einsum = None
try:
from torch import einsum as torch_einsum, from_numpy
except ImportError:
torch_einsum = None
def build_ort_einsum(equation, op_version=14): # opset=13, 14, ...
node = OnnxEinsum('x', 'y', equation=equation,
op_version=op_version,
output_names=['z'])
onx = node.to_onnx(inputs=[('x', FloatTensorType()),
('y', FloatTensorType())],
target_opset=op_version)
sess = InferenceSession(onx.SerializeToString())
return lambda x, y: sess.run(None, {'x': x, 'y': y})
def build_ort_decomposed(equation, op_version=14): # opset=13, 14, ...
cache = _einsum(equation, numpy.float32, opset=op_version,
optimize=True, verbose=True, runtime="python")
if not hasattr(cache, 'onnx_'):
cache.build()
sess = InferenceSession(cache.onnx_.SerializeToString())
return lambda x, y: sess.run(None, {'X0': x, 'X1': y})
def loop_einsum_eq(fct, equation, xs, ys):
for x, y in zip(xs, ys):
fct(equation, x, y)
def loop_einsum_eq_th(fct, equation, xs, ys):
for x, y in zip(xs, ys):
fct(equation, x, y, nthread=-1)
def loop_einsum(fct, xs, ys):
for x, y in zip(xs, ys):
fct(x, y)
def custom_einsum_float_tr(eq, x, y):
if eq == "bshn,bthn->bnts":
x = x.transpose((0, 1, 3, 2))
y = y.transpose((0, 1, 3, 2))
return custom_einsum_float("bsnh,btnh->bnts", x, y, nthread=-1)
if eq == "bhsn,bhtn->bnts":
x = x.transpose((0, 2, 3, 1))
y = y.transpose((0, 2, 3, 1))
return custom_einsum_float("bsnh,btnh->bnts", x, y, nthread=-1)
return custom_einsum_float(eq, x, y, nthread=-1)
def benchmark_equation(equation):
# equations
ort_einsum = build_ort_einsum(equation)
ort_einsum_decomposed = build_ort_decomposed(equation)
res = []
for dim in tqdm([8, 16, 32, 64, 100, 128, 200,
256, 500, 512]):
xs = [numpy.random.rand(2, dim, 12, 64).astype(numpy.float32)
for _ in range(5)]
ys = [numpy.random.rand(2, dim, 12, 64).astype(numpy.float32)
for _ in range(5)]
# numpy
ctx = dict(equation=equation, xs=xs, ys=ys, einsum=numpy.einsum,
loop_einsum=loop_einsum, loop_einsum_eq=loop_einsum_eq,
loop_einsum_eq_th=loop_einsum_eq_th)
obs = measure_time(
"loop_einsum_eq(einsum, equation, xs, ys)",
div_by_number=True, context=ctx, repeat=5, number=1)
obs['dim'] = dim
obs['fct'] = 'numpy.einsum'
res.append(obs)
# opt-einsum
ctx['einsum'] = contract
obs = measure_time(
"loop_einsum_eq(einsum, equation, xs, ys)",
div_by_number=True, context=ctx, repeat=5, number=1)
obs['dim'] = dim
obs['fct'] = 'opt-einsum'
res.append(obs)
# onnxruntime
ctx['einsum'] = ort_einsum
obs = measure_time(
"loop_einsum(einsum, xs, ys)",
div_by_number=True, context=ctx, repeat=5, number=1)
obs['dim'] = dim
obs['fct'] = 'ort_einsum'
res.append(obs)
# onnxruntime decomposed
ctx['einsum'] = ort_einsum_decomposed
obs = measure_time(
"loop_einsum(einsum, xs, ys)",
div_by_number=True, context=ctx, repeat=5, number=1)
obs['dim'] = dim
obs['fct'] = 'ort_dec'
res.append(obs)
# custom implementation
ctx['einsum'] = custom_einsum_float
obs = measure_time(
"loop_einsum_eq_th(einsum, equation, xs, ys)",
div_by_number=True, context=ctx, repeat=5, number=1)
obs['dim'] = dim
obs['fct'] = 'c_einsum'
res.append(obs)
# transpose + custom implementation
ctx['einsum'] = custom_einsum_float_tr
obs = measure_time(
"loop_einsum_eq(einsum, equation, xs, ys)",
div_by_number=True, context=ctx, repeat=5, number=1)
obs['dim'] = dim
obs['fct'] = 'c_einsum_tr'
res.append(obs)
if tf_einsum is not None:
# tensorflow
ctx['einsum'] = tf_einsum
ctx['xs'] = [convert_to_tensor(x) for x in xs]
ctx['ys'] = [convert_to_tensor(y) for y in ys]
obs = measure_time(
"loop_einsum_eq(einsum, equation, xs, ys)",
div_by_number=True, context=ctx, repeat=5, number=1)
obs['dim'] = dim
obs['fct'] = 'tf_einsum'
res.append(obs)
if torch_einsum is not None:
# torch
ctx['einsum'] = torch_einsum
ctx['xs'] = [from_numpy(x) for x in xs]
ctx['ys'] = [from_numpy(y) for y in ys]
obs = measure_time(
"loop_einsum_eq(einsum, equation, xs, ys)",
div_by_number=True, context=ctx, repeat=5, number=1)
obs['dim'] = dim
obs['fct'] = 'torch_einsum'
res.append(obs)
# Dataframes
df = pandas.DataFrame(res)
piv = df.pivot('dim', 'fct', 'average')
rs = piv.copy()
rs['c_einsum'] = rs['numpy.einsum'] / rs['c_einsum']
rs['ort_einsum'] = rs['numpy.einsum'] / rs['ort_einsum']
rs['ort_dec'] = rs['numpy.einsum'] / rs['ort_dec']
rs['opt-einsum'] = rs['numpy.einsum'] / rs['opt-einsum']
if 'c_einsum_tr' in rs.columns:
rs['c_einsum_tr'] = rs['numpy.einsum'] / rs['c_einsum_tr']
if 'tf_einsum' in rs.columns:
rs['tf_einsum'] = rs['numpy.einsum'] / rs['tf_einsum']
if 'torch_einsum' in rs.columns:
rs['torch_einsum'] = rs['numpy.einsum'] / rs['torch_einsum']
rs['numpy.einsum'] = 1.
# Graphs.
fig, ax = plt.subplots(1, 2, figsize=(14, 5))
piv.plot(logx=True, logy=True, ax=ax[0],
title="Einsum benchmark\n%s -- (2, N, 12, 64)"
" lower better" % equation)
ax[0].legend(prop={"size": 9})
rs.plot(logx=True, logy=True, ax=ax[1],
title="Einsum Speedup, baseline=numpy\n%s -- (2, N, 12, 64)"
" higher better" % equation)
ax[1].plot([min(rs.index), max(rs.index)], [0.5, 0.5], 'g--')
ax[1].plot([min(rs.index), max(rs.index)], [2., 2.], 'g--')
ax[1].legend(prop={"size": 9})
return df, rs, ax
First equation: bsnh,btnh->bnts#
The decomposition of this equation without einsum function gives the following.
dfs = []
equation = "bsnh,btnh->bnts"
df, piv, ax = benchmark_equation(equation)
df.pivot("fct", "dim", "average")
dfs.append(df)
Out:
0%| | 0/121 [00:00<?, ?it/s]
0.032 rtbest='bsnh,btnh->bnts': 0%| | 0/121 [00:00<?, ?it/s]
0.032 rtbest='bsnh,btnh->bnts': 1%| | 1/121 [00:00<00:17, 6.80it/s]
0.032 rtbest='bsnh,btnh->bnts': 2%|2 | 3/121 [00:00<00:09, 12.16it/s]
0.032 rtbest='bnsh,btsh->bstn': 2%|2 | 3/121 [00:00<00:09, 12.16it/s]
0.032 rtbest='bnsh,btsh->bstn': 4%|4 | 5/121 [00:00<00:08, 14.10it/s]
0.032 rtbest='bnsh,btsh->bstn': 6%|5 | 7/121 [00:00<00:07, 15.09it/s]
0.032 rtbest='bnsh,btsh->bstn': 7%|7 | 9/121 [00:00<00:07, 15.68it/s]
0.032 rtbest='bnsh,btsh->bstn': 9%|9 | 11/121 [00:00<00:06, 16.10it/s]
0.032 rtbest='bnsh,btsh->bstn': 11%|# | 13/121 [00:00<00:06, 16.01it/s]
0.032 rtbest='bnsh,btsh->bstn': 12%|#2 | 15/121 [00:00<00:06, 16.21it/s]
0.032 rtbest='bnsh,btsh->bstn': 14%|#4 | 17/121 [00:01<00:06, 16.39it/s]
0.032 rtbest='bhts,bnts->btnh': 14%|#4 | 17/121 [00:01<00:06, 16.39it/s]
0.032 rtbest='bhts,bnts->btnh': 16%|#5 | 19/121 [00:01<00:06, 16.48it/s]
0.032 rtbest='bhts,bnts->btnh': 17%|#7 | 21/121 [00:01<00:06, 16.55it/s]
0.032 rtbest='bhts,bnts->btnh': 19%|#9 | 23/121 [00:01<00:05, 16.64it/s]
0.032 rtbest='bhts,bnts->btnh': 21%|## | 25/121 [00:01<00:05, 16.65it/s]
0.032 rtbest='bhts,bnts->btnh': 22%|##2 | 27/121 [00:01<00:05, 16.34it/s]
0.032 rtbest='bhts,bnts->btnh': 24%|##3 | 29/121 [00:01<00:05, 16.45it/s]
0.032 rtbest='bhts,bnts->btnh': 26%|##5 | 31/121 [00:01<00:05, 16.51it/s]
0.032 rtbest='nshb,nthb->nhts': 26%|##5 | 31/121 [00:02<00:05, 16.51it/s]
0.032 rtbest='nshb,nthb->nhts': 27%|##7 | 33/121 [00:02<00:05, 16.55it/s]
0.032 rtbest='snhb,sthb->shtn': 27%|##7 | 33/121 [00:02<00:05, 16.55it/s]
0.032 rtbest='tnhb,tshb->thsn': 27%|##7 | 33/121 [00:02<00:05, 16.55it/s]
0.032 rtbest='tnhb,tshb->thsn': 29%|##8 | 35/121 [00:02<00:05, 16.54it/s]
0.032 rtbest='sthb,snhb->shnt': 29%|##8 | 35/121 [00:02<00:05, 16.54it/s]
0.032 rtbest='sthb,snhb->shnt': 31%|### | 37/121 [00:02<00:05, 16.60it/s]
0.032 rtbest='sthb,snhb->shnt': 32%|###2 | 39/121 [00:02<00:04, 16.63it/s]
0.032 rtbest='sthb,snhb->shnt': 34%|###3 | 41/121 [00:02<00:04, 16.33it/s]
0.032 rtbest='sthb,snhb->shnt': 36%|###5 | 43/121 [00:02<00:04, 16.44it/s]
0.032 rtbest='sthb,snhb->shnt': 37%|###7 | 45/121 [00:02<00:04, 16.51it/s]
0.032 rtbest='sthb,snhb->shnt': 39%|###8 | 47/121 [00:02<00:04, 16.59it/s]
0.032 rtbest='tnsb,thsb->tshn': 39%|###8 | 47/121 [00:03<00:04, 16.59it/s]
0.032 rtbest='tnsb,thsb->tshn': 40%|#### | 49/121 [00:03<00:04, 16.61it/s]
0.032 rtbest='tnsb,thsb->tshn': 42%|####2 | 51/121 [00:03<00:04, 16.69it/s]
0.032 rtbest='hnbt,hsbt->hbsn': 42%|####2 | 51/121 [00:03<00:04, 16.69it/s]
0.032 rtbest='hnbt,hsbt->hbsn': 44%|####3 | 53/121 [00:03<00:04, 16.76it/s]
0.032 rtbest='hnbt,hsbt->hbsn': 45%|####5 | 55/121 [00:03<00:04, 16.43it/s]
0.032 rtbest='hnbt,hsbt->hbsn': 47%|####7 | 57/121 [00:03<00:03, 16.52it/s]
0.032 rtbest='hnbt,hsbt->hbsn': 49%|####8 | 59/121 [00:03<00:03, 16.60it/s]
0.032 rtbest='hnbt,hsbt->hbsn': 50%|##### | 61/121 [00:03<00:03, 16.63it/s]
0.032 rtbest='hnbt,hsbt->hbsn': 52%|#####2 | 63/121 [00:03<00:03, 16.68it/s]
0.032 rtbest='hnbt,hsbt->hbsn': 54%|#####3 | 65/121 [00:04<00:03, 16.73it/s]
0.032 rtbest='hnbt,hsbt->hbsn': 55%|#####5 | 67/121 [00:04<00:03, 16.76it/s]
0.032 rtbest='ntbs,nhbs->nbht': 55%|#####5 | 67/121 [00:04<00:03, 16.76it/s]
0.032 rtbest='ntbs,nhbs->nbht': 57%|#####7 | 69/121 [00:04<00:03, 16.35it/s]
0.032 rtbest='ntbs,nhbs->nbht': 59%|#####8 | 71/121 [00:04<00:03, 16.45it/s]
0.032 rtbest='ntbs,nhbs->nbht': 60%|###### | 73/121 [00:04<00:02, 16.54it/s]
0.032 rtbest='ntbs,nhbs->nbht': 62%|######1 | 75/121 [00:04<00:02, 16.61it/s]
0.032 rtbest='ntbs,nhbs->nbht': 64%|######3 | 77/121 [00:04<00:02, 16.68it/s]
0.032 rtbest='ntbs,nhbs->nbht': 65%|######5 | 79/121 [00:04<00:02, 16.70it/s]
0.032 rtbest='ntbs,nhbs->nbht': 67%|######6 | 81/121 [00:04<00:02, 16.72it/s]
0.032 rtbest='ntbs,nhbs->nbht': 69%|######8 | 83/121 [00:05<00:02, 16.35it/s]
0.032 rtbest='ntbs,nhbs->nbht': 70%|####### | 85/121 [00:05<00:02, 16.41it/s]
0.032 rtbest='ntbs,nhbs->nbht': 72%|#######1 | 87/121 [00:05<00:02, 16.50it/s]
0.032 rtbest='ntbs,nhbs->nbht': 74%|#######3 | 89/121 [00:05<00:01, 16.58it/s]
0.032 rtbest='ntbs,nhbs->nbht': 75%|#######5 | 91/121 [00:05<00:01, 16.66it/s]
0.032 rtbest='ntbs,nhbs->nbht': 77%|#######6 | 93/121 [00:05<00:01, 16.68it/s]
0.032 rtbest='ntbs,nhbs->nbht': 79%|#######8 | 95/121 [00:05<00:01, 16.71it/s]
0.032 rtbest='ntbs,nhbs->nbht': 80%|######## | 97/121 [00:05<00:01, 16.34it/s]
0.032 rtbest='ntbs,nhbs->nbht': 82%|########1 | 99/121 [00:06<00:01, 16.43it/s]
0.032 rtbest='ntbs,nhbs->nbht': 83%|########3 | 101/121 [00:06<00:01, 16.50it/s]
0.032 rtbest='ntbs,nhbs->nbht': 85%|########5 | 103/121 [00:06<00:01, 16.58it/s]
0.032 rtbest='ntbs,nhbs->nbht': 87%|########6 | 105/121 [00:06<00:00, 16.64it/s]
0.032 rtbest='ntbs,nhbs->nbht': 88%|########8 | 107/121 [00:06<00:00, 16.66it/s]
0.032 rtbest='ntbs,nhbs->nbht': 90%|######### | 109/121 [00:06<00:00, 16.69it/s]
0.032 rtbest='ntbs,nhbs->nbht': 92%|#########1| 111/121 [00:06<00:00, 16.33it/s]
0.032 rtbest='ntbs,nhbs->nbht': 93%|#########3| 113/121 [00:06<00:00, 16.40it/s]
0.032 rtbest='ntbs,nhbs->nbht': 95%|#########5| 115/121 [00:07<00:00, 16.50it/s]
0.032 rtbest='ntbs,nhbs->nbht': 97%|#########6| 117/121 [00:07<00:00, 16.57it/s]
0.032 rtbest='ntbs,nhbs->nbht': 98%|#########8| 119/121 [00:07<00:00, 16.60it/s]
0.032 rtbest='ntbs,nhbs->nbht': 100%|##########| 121/121 [00:07<00:00, 16.62it/s]
0.032 rtbest='ntbs,nhbs->nbht': 100%|##########| 121/121 [00:07<00:00, 16.38it/s]
0%| | 0/10 [00:00<?, ?it/s]
10%|# | 1/10 [00:01<00:12, 1.35s/it]
20%|## | 2/10 [00:01<00:06, 1.16it/s]
30%|### | 3/10 [00:03<00:07, 1.08s/it]
40%|#### | 4/10 [00:05<00:08, 1.39s/it]
50%|##### | 5/10 [00:08<00:09, 1.97s/it]
60%|###### | 6/10 [00:12<00:10, 2.67s/it]
70%|####### | 7/10 [00:20<00:13, 4.51s/it]
80%|######## | 8/10 [00:33<00:14, 7.16s/it]
90%|######### | 9/10 [01:22<00:20, 20.33s/it]
100%|##########| 10/10 [02:18<00:00, 31.28s/it]
100%|##########| 10/10 [02:18<00:00, 13.83s/it]
Second equation: bshn,bthn->bnts#
The summation does not happen on the last axis but on the previous one. Is it worth transposing before doing the summation… The decomposition of this equation without einsum function gives the following.
equation = "bshn,bthn->bnts"
df, piv, ax = benchmark_equation(equation)
df.pivot("fct", "dim", "average")
dfs.append(df)
Out:
0%| | 0/121 [00:00<?, ?it/s]
0.031 rtbest='bshn,bthn->bnts': 0%| | 0/121 [00:00<?, ?it/s]
0.031 rtbest='bshn,bthn->bnts': 0%| | 0/121 [00:00<?, ?it/s]
0.031 rtbest='bshn,bthn->bnts': 2%|1 | 2/121 [00:00<00:07, 16.28it/s]
0.031 rtbest='bthn,bshn->bnst': 2%|1 | 2/121 [00:00<00:07, 16.28it/s]
0.031 rtbest='bthn,bshn->bnst': 3%|3 | 4/121 [00:00<00:07, 16.69it/s]
0.031 rtbest='bthn,bshn->bnst': 5%|4 | 6/121 [00:00<00:06, 16.87it/s]
0.031 rtbest='bsnh,btnh->bhts': 5%|4 | 6/121 [00:00<00:06, 16.87it/s]
0.031 rtbest='bsnh,btnh->bhts': 7%|6 | 8/121 [00:00<00:06, 16.94it/s]
0.031 rtbest='bnsh,btsh->bhtn': 7%|6 | 8/121 [00:00<00:06, 16.94it/s]
0.031 rtbest='bnsh,btsh->bhtn': 8%|8 | 10/121 [00:00<00:06, 17.02it/s]
0.031 rtbest='bnsh,btsh->bhtn': 10%|9 | 12/121 [00:00<00:06, 17.21it/s]
0.031 rtbest='bnsh,btsh->bhtn': 12%|#1 | 14/121 [00:00<00:06, 16.93it/s]
0.031 rtbest='bnsh,btsh->bhtn': 13%|#3 | 16/121 [00:00<00:06, 17.01it/s]
0.031 rtbest='bnsh,btsh->bhtn': 15%|#4 | 18/121 [00:01<00:06, 17.06it/s]
0.031 rtbest='bnsh,btsh->bhtn': 17%|#6 | 20/121 [00:01<00:05, 17.09it/s]
0.031 rtbest='bnsh,btsh->bhtn': 18%|#8 | 22/121 [00:01<00:05, 17.13it/s]
0.031 rtbest='bnsh,btsh->bhtn': 20%|#9 | 24/121 [00:01<00:05, 17.20it/s]
0.031 rtbest='bnsh,btsh->bhtn': 21%|##1 | 26/121 [00:01<00:05, 17.16it/s]
0.031 rtbest='bnsh,btsh->bhtn': 23%|##3 | 28/121 [00:01<00:05, 16.86it/s]
0.031 rtbest='bnsh,btsh->bhtn': 25%|##4 | 30/121 [00:01<00:05, 16.90it/s]
0.031 rtbest='bnsh,btsh->bhtn': 26%|##6 | 32/121 [00:01<00:05, 16.89it/s]
0.031 rtbest='bnsh,btsh->bhtn': 28%|##8 | 34/121 [00:02<00:05, 16.84it/s]
0.031 rtbest='bnsh,btsh->bhtn': 30%|##9 | 36/121 [00:02<00:05, 16.83it/s]
0.031 rtbest='bnsh,btsh->bhtn': 31%|###1 | 38/121 [00:02<00:04, 16.85it/s]
0.031 rtbest='bnsh,btsh->bhtn': 33%|###3 | 40/121 [00:02<00:04, 16.89it/s]
0.031 rtbest='bnsh,btsh->bhtn': 35%|###4 | 42/121 [00:02<00:04, 16.63it/s]
0.031 rtbest='bnsh,btsh->bhtn': 36%|###6 | 44/121 [00:02<00:04, 16.69it/s]
0.031 rtbest='bnsh,btsh->bhtn': 38%|###8 | 46/121 [00:02<00:04, 16.75it/s]
0.031 rtbest='bnsh,btsh->bhtn': 40%|###9 | 48/121 [00:02<00:04, 16.79it/s]
0.031 rtbest='bnsh,btsh->bhtn': 41%|####1 | 50/121 [00:02<00:04, 16.78it/s]
0.031 rtbest='bnsh,btsh->bhtn': 43%|####2 | 52/121 [00:03<00:04, 16.74it/s]
0.031 rtbest='bnsh,btsh->bhtn': 45%|####4 | 54/121 [00:03<00:03, 16.81it/s]
0.031 rtbest='bnsh,btsh->bhtn': 46%|####6 | 56/121 [00:03<00:03, 16.52it/s]
0.031 rtbest='bnsh,btsh->bhtn': 48%|####7 | 58/121 [00:03<00:03, 16.57it/s]
0.031 rtbest='bnsh,btsh->bhtn': 50%|####9 | 60/121 [00:03<00:03, 16.59it/s]
0.031 rtbest='bnsh,btsh->bhtn': 51%|#####1 | 62/121 [00:03<00:03, 16.63it/s]
0.031 rtbest='bnsh,btsh->bhtn': 53%|#####2 | 64/121 [00:03<00:03, 16.63it/s]
0.031 rtbest='bnsh,btsh->bhtn': 55%|#####4 | 66/121 [00:03<00:03, 16.66it/s]
0.031 rtbest='bnsh,btsh->bhtn': 56%|#####6 | 68/121 [00:04<00:03, 16.69it/s]
0.031 rtbest='bnsh,btsh->bhtn': 58%|#####7 | 70/121 [00:04<00:03, 16.72it/s]
0.031 rtbest='bnsh,btsh->bhtn': 60%|#####9 | 72/121 [00:04<00:02, 16.39it/s]
0.031 rtbest='bnsh,btsh->bhtn': 61%|######1 | 74/121 [00:04<00:02, 16.52it/s]
0.031 rtbest='bnsh,btsh->bhtn': 63%|######2 | 76/121 [00:04<00:02, 16.65it/s]
0.031 rtbest='bnsh,btsh->bhtn': 64%|######4 | 78/121 [00:04<00:02, 16.79it/s]
0.031 rtbest='bnsh,btsh->bhtn': 66%|######6 | 80/121 [00:04<00:02, 16.90it/s]
0.031 rtbest='bnsh,btsh->bhtn': 68%|######7 | 82/121 [00:04<00:02, 16.89it/s]
0.031 rtbest='bnsh,btsh->bhtn': 69%|######9 | 84/121 [00:04<00:02, 16.93it/s]
0.031 rtbest='bnsh,btsh->bhtn': 71%|#######1 | 86/121 [00:05<00:02, 16.51it/s]
0.031 rtbest='bnsh,btsh->bhtn': 73%|#######2 | 88/121 [00:05<00:01, 16.53it/s]
0.031 rtbest='bnsh,btsh->bhtn': 74%|#######4 | 90/121 [00:05<00:01, 16.56it/s]
0.031 rtbest='bnsh,btsh->bhtn': 76%|#######6 | 92/121 [00:05<00:01, 16.66it/s]
0.031 rtbest='bnsh,btsh->bhtn': 78%|#######7 | 94/121 [00:05<00:01, 16.77it/s]
0.031 rtbest='bnsh,btsh->bhtn': 79%|#######9 | 96/121 [00:05<00:01, 16.76it/s]
0.031 rtbest='bnsh,btsh->bhtn': 81%|######## | 98/121 [00:05<00:01, 16.84it/s]
0.031 rtbest='bnsh,btsh->bhtn': 83%|########2 | 100/121 [00:05<00:01, 16.57it/s]
0.031 rtbest='bnsh,btsh->bhtn': 84%|########4 | 102/121 [00:06<00:01, 16.72it/s]
0.031 rtbest='bnsh,btsh->bhtn': 86%|########5 | 104/121 [00:06<00:01, 16.82it/s]
0.031 rtbest='bnsh,btsh->bhtn': 88%|########7 | 106/121 [00:06<00:00, 16.83it/s]
0.031 rtbest='bnsh,btsh->bhtn': 89%|########9 | 108/121 [00:06<00:00, 16.84it/s]
0.031 rtbest='bnsh,btsh->bhtn': 91%|######### | 110/121 [00:06<00:00, 16.79it/s]
0.031 rtbest='bnsh,btsh->bhtn': 93%|#########2| 112/121 [00:06<00:00, 16.76it/s]
0.031 rtbest='bnsh,btsh->bhtn': 94%|#########4| 114/121 [00:06<00:00, 16.42it/s]
0.031 rtbest='bnsh,btsh->bhtn': 96%|#########5| 116/121 [00:06<00:00, 16.52it/s]
0.031 rtbest='bnsh,btsh->bhtn': 98%|#########7| 118/121 [00:07<00:00, 16.66it/s]
0.031 rtbest='bnsh,btsh->bhtn': 99%|#########9| 120/121 [00:07<00:00, 16.65it/s]
0.031 rtbest='bnsh,btsh->bhtn': 100%|##########| 121/121 [00:07<00:00, 16.76it/s]
0%| | 0/10 [00:00<?, ?it/s]
10%|# | 1/10 [00:00<00:00, 9.90it/s]
20%|## | 2/10 [00:00<00:03, 2.04it/s]
30%|### | 3/10 [00:02<00:06, 1.07it/s]
40%|#### | 4/10 [00:04<00:08, 1.46s/it]
50%|##### | 5/10 [00:08<00:11, 2.31s/it]
60%|###### | 6/10 [00:14<00:13, 3.44s/it]
70%|####### | 7/10 [00:27<00:20, 6.73s/it]
80%|######## | 8/10 [00:51<00:24, 12.17s/it]
90%|######### | 9/10 [02:49<00:45, 45.31s/it]
100%|##########| 10/10 [05:19<00:00, 77.54s/it]
100%|##########| 10/10 [05:19<00:00, 31.92s/it]
Third equation: bhsn,bhtn->bnts#
The summation does not happen on the last axis but on the second one. It is worth transposing before multiplying. The decomposition of this equation without einsum function gives the following.
equation = "bhsn,bhtn->bnts"
df, piv, ax = benchmark_equation(equation)
df.pivot("fct", "dim", "average")
dfs.append(df)
Out:
0%| | 0/121 [00:00<?, ?it/s]
0.032 rtbest='bhsn,bhtn->bnts': 0%| | 0/121 [00:00<?, ?it/s]
0.032 rtbest='bhsn,bhtn->bnts': 2%|1 | 2/121 [00:00<00:07, 16.03it/s]
0.032 rtbest='bhsn,bhtn->bnts': 3%|3 | 4/121 [00:00<00:07, 16.52it/s]
0.032 rtbest='bhsn,bhtn->bnts': 5%|4 | 6/121 [00:00<00:06, 16.65it/s]
0.032 rtbest='bhsn,bhtn->bnts': 7%|6 | 8/121 [00:00<00:06, 16.74it/s]
0.032 rtbest='bhsn,bhtn->bnts': 8%|8 | 10/121 [00:00<00:06, 16.77it/s]
0.032 rtbest='btnh,btsh->bhsn': 8%|8 | 10/121 [00:00<00:06, 16.77it/s]
0.032 rtbest='btnh,btsh->bhsn': 10%|9 | 12/121 [00:00<00:06, 16.84it/s]
0.032 rtbest='bnhs,bnts->bsth': 10%|9 | 12/121 [00:00<00:06, 16.84it/s]
0.032 rtbest='bnhs,bnts->bsth': 12%|#1 | 14/121 [00:00<00:06, 16.49it/s]
0.032 rtbest='bnhs,bnts->bsth': 13%|#3 | 16/121 [00:00<00:06, 16.59it/s]
0.032 rtbest='bnhs,bnts->bsth': 15%|#4 | 18/121 [00:01<00:06, 16.63it/s]
0.032 rtbest='bnhs,bnts->bsth': 17%|#6 | 20/121 [00:01<00:06, 16.66it/s]
0.032 rtbest='bnhs,bnts->bsth': 18%|#8 | 22/121 [00:01<00:05, 16.71it/s]
0.032 rtbest='bnhs,bnts->bsth': 20%|#9 | 24/121 [00:01<00:05, 16.73it/s]
0.032 rtbest='bnhs,bnts->bsth': 21%|##1 | 26/121 [00:01<00:05, 16.74it/s]
0.032 rtbest='hbns,hbts->hstn': 21%|##1 | 26/121 [00:01<00:05, 16.74it/s]
0.032 rtbest='hbns,hbts->hstn': 23%|##3 | 28/121 [00:01<00:05, 16.44it/s]
0.032 rtbest='hbns,hbts->hstn': 25%|##4 | 30/121 [00:01<00:05, 16.51it/s]
0.032 rtbest='hbns,hbts->hstn': 26%|##6 | 32/121 [00:01<00:05, 16.51it/s]
0.032 rtbest='hbns,hbts->hstn': 28%|##8 | 34/121 [00:02<00:05, 16.49it/s]
0.032 rtbest='hbns,hbts->hstn': 30%|##9 | 36/121 [00:02<00:05, 16.46it/s]
0.032 rtbest='hbns,hbts->hstn': 31%|###1 | 38/121 [00:02<00:05, 16.49it/s]
0.032 rtbest='hbns,hbts->hstn': 33%|###3 | 40/121 [00:02<00:04, 16.53it/s]
0.032 rtbest='hbns,hbts->hstn': 35%|###4 | 42/121 [00:02<00:04, 16.56it/s]
0.032 rtbest='hbns,hbts->hstn': 36%|###6 | 44/121 [00:02<00:04, 16.26it/s]
0.032 rtbest='hbns,hbts->hstn': 38%|###8 | 46/121 [00:02<00:04, 16.35it/s]
0.032 rtbest='hbns,hbts->hstn': 40%|###9 | 48/121 [00:02<00:04, 16.38it/s]
0.032 rtbest='hbns,hbts->hstn': 41%|####1 | 50/121 [00:03<00:04, 16.38it/s]
0.032 rtbest='hbns,hbts->hstn': 43%|####2 | 52/121 [00:03<00:04, 16.42it/s]
0.032 rtbest='hbns,hbts->hstn': 45%|####4 | 54/121 [00:03<00:04, 16.45it/s]
0.032 rtbest='hbns,hbts->hstn': 46%|####6 | 56/121 [00:03<00:04, 15.89it/s]
0.032 rtbest='hbns,hbts->hstn': 48%|####7 | 58/121 [00:03<00:03, 15.76it/s]
0.032 rtbest='hbns,hbts->hstn': 50%|####9 | 60/121 [00:03<00:03, 15.91it/s]
0.032 rtbest='hbns,hbts->hstn': 51%|#####1 | 62/121 [00:03<00:03, 16.03it/s]
0.032 rtbest='hbns,hbts->hstn': 53%|#####2 | 64/121 [00:03<00:03, 16.15it/s]
0.032 rtbest='hbns,hbts->hstn': 55%|#####4 | 66/121 [00:04<00:03, 16.20it/s]
0.032 rtbest='hbns,hbts->hstn': 56%|#####6 | 68/121 [00:04<00:03, 16.28it/s]
0.032 rtbest='hbns,hbts->hstn': 58%|#####7 | 70/121 [00:04<00:03, 16.35it/s]
0.032 rtbest='hbns,hbts->hstn': 60%|#####9 | 72/121 [00:04<00:03, 16.07it/s]
0.032 rtbest='hbns,hbts->hstn': 61%|######1 | 74/121 [00:04<00:02, 16.20it/s]
0.032 rtbest='hbns,hbts->hstn': 63%|######2 | 76/121 [00:04<00:02, 16.33it/s]
0.032 rtbest='hbns,hbts->hstn': 64%|######4 | 78/121 [00:04<00:02, 16.42it/s]
0.032 rtbest='hbns,hbts->hstn': 66%|######6 | 80/121 [00:04<00:02, 16.51it/s]
0.032 rtbest='hbns,hbts->hstn': 68%|######7 | 82/121 [00:04<00:02, 16.52it/s]
0.032 rtbest='hbns,hbts->hstn': 69%|######9 | 84/121 [00:05<00:02, 16.55it/s]
0.032 rtbest='hbns,hbts->hstn': 71%|#######1 | 86/121 [00:05<00:02, 16.21it/s]
0.032 rtbest='hbns,hbts->hstn': 73%|#######2 | 88/121 [00:05<00:02, 16.24it/s]
0.032 rtbest='hbns,hbts->hstn': 74%|#######4 | 90/121 [00:05<00:01, 16.26it/s]
0.032 rtbest='hbns,hbts->hstn': 76%|#######6 | 92/121 [00:05<00:01, 16.34it/s]
0.032 rtbest='hbns,hbts->hstn': 78%|#######7 | 94/121 [00:05<00:01, 16.43it/s]
0.032 rtbest='hbns,hbts->hstn': 79%|#######9 | 96/121 [00:05<00:01, 16.42it/s]
0.032 rtbest='hbns,hbts->hstn': 81%|######## | 98/121 [00:05<00:01, 16.48it/s]
0.032 rtbest='hbns,hbts->hstn': 83%|########2 | 100/121 [00:06<00:01, 16.57it/s]
0.032 rtbest='hbns,hbts->hstn': 84%|########4 | 102/121 [00:06<00:01, 16.28it/s]
0.032 rtbest='hbns,hbts->hstn': 86%|########5 | 104/121 [00:06<00:01, 16.39it/s]
0.032 rtbest='hbns,hbts->hstn': 88%|########7 | 106/121 [00:06<00:00, 16.40it/s]
0.032 rtbest='hbns,hbts->hstn': 89%|########9 | 108/121 [00:06<00:00, 16.43it/s]
0.032 rtbest='hbns,hbts->hstn': 91%|######### | 110/121 [00:06<00:00, 16.44it/s]
0.032 rtbest='hbns,hbts->hstn': 93%|#########2| 112/121 [00:06<00:00, 16.41it/s]
0.032 rtbest='hbns,hbts->hstn': 94%|#########4| 114/121 [00:06<00:00, 16.42it/s]
0.032 rtbest='hbns,hbts->hstn': 96%|#########5| 116/121 [00:07<00:00, 16.16it/s]
0.032 rtbest='hbns,hbts->hstn': 98%|#########7| 118/121 [00:07<00:00, 16.27it/s]
0.032 rtbest='hbns,hbts->hstn': 99%|#########9| 120/121 [00:07<00:00, 16.28it/s]
0.032 rtbest='hbns,hbts->hstn': 100%|##########| 121/121 [00:07<00:00, 16.39it/s]
0%| | 0/10 [00:00<?, ?it/s]
10%|# | 1/10 [00:00<00:06, 1.36it/s]
20%|## | 2/10 [00:01<00:06, 1.28it/s]
30%|### | 3/10 [00:02<00:07, 1.07s/it]
40%|#### | 4/10 [00:04<00:07, 1.31s/it]
50%|##### | 5/10 [00:06<00:07, 1.49s/it]
60%|###### | 6/10 [00:08<00:07, 1.84s/it]
70%|####### | 7/10 [00:12<00:07, 2.37s/it]
80%|######## | 8/10 [00:16<00:06, 3.02s/it]
90%|######### | 9/10 [00:27<00:05, 5.31s/it]
100%|##########| 10/10 [00:37<00:00, 6.92s/it]
100%|##########| 10/10 [00:37<00:00, 3.77s/it]
Conclusion#
pytorch seems quite efficient on these examples. The custom implementation was a way to investigate the implementation of einsum and find some ways to optimize it.
merged = pandas.concat(dfs)
name = "einsum"
merged.to_csv("plot_%s.csv" % name, index=False)
merged.to_excel("plot_%s.xlsx" % name, index=False)
plt.savefig("plot_%s.png" % name)
plt.show()
Total running time of the script: ( 8 minutes 46.672 seconds)