Note

Click here to download the full example code

Compares implementations of ReduceSum#

This example compares the numpy.sum from numpy, to onnxruntime implementation. If available, tensorflow and pytorch are included as well.

Available optimisation#

The code shows which parallelisation optimisation could be used, AVX or SSE and the number of available processors.

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 OnnxReduceSumApi11
from cpyquickhelper.numbers import measure_time
from tqdm import tqdm
from mlprodict.testing.experimental_c_impl.experimental_c import (
    code_optimisation, custom_reducesum_rk_float)
print(code_optimisation())

Out:

AVX-omp=8

ReduceSum implementations#

try:
    from tensorflow.math import reduce_sum as tf_reduce_sum
    from tensorflow import convert_to_tensor
except ImportError:
    tf_reduce_sum = None
try:
    from torch import sum as torch_sum, from_numpy
except ImportError:
    torch_sum = None


def build_ort_reducesum(axes, op_version=14):  # opset=13, 14, ...
    node = OnnxReduceSumApi11('x', axes=axes, op_version=op_version,
                              output_names=['z'])
    onx = node.to_onnx(inputs=[('x', FloatTensorType())],
                       target_opset=op_version)
    sess = InferenceSession(onx.SerializeToString())
    return lambda x, y: sess.run(None, {'x': x})


def loop_fct(fct, xs, ys):
    for x, y in zip(xs, ys):
        fct(x, y)


def benchmark_op(axes, repeat=5, number=5, name="ReduceSum", shape_fct=None,
                 custom_impl=False):
    if shape_fct is None:
        def shape_fct(dim):
            return (3, dim, 1, 128, 64)
    ort_fct = build_ort_reducesum(axes)
    res = []
    for dim in tqdm([8, 16, 32, 64, 100, 128, 200,
                     256, 400, 512, 1024]):
        shape = shape_fct(dim)
        n_arrays = 10 if dim < 512 else 4
        xs = [numpy.random.rand(*shape).astype(numpy.float32)
              for _ in range(n_arrays)]
        ys = [numpy.array(axes, dtype=numpy.int64)
              for _ in range(n_arrays)]
        info = dict(axes=axes, shape=shape)

        # numpy
        ctx = dict(
            xs=xs, ys=ys,
            fct=lambda x, y: numpy.sum(x, *y),
            loop_fct=loop_fct)
        obs = measure_time(
            "loop_fct(fct, xs, ys)",
            div_by_number=True, context=ctx, repeat=repeat, number=number)
        obs['dim'] = dim
        obs['fct'] = 'numpy'
        obs.update(info)
        res.append(obs)

        # onnxruntime
        ctx['fct'] = ort_fct
        obs = measure_time(
            "loop_fct(fct, xs, ys)",
            div_by_number=True, context=ctx, repeat=repeat, number=number)
        obs['dim'] = dim
        obs['fct'] = 'ort'
        obs.update(info)
        res.append(obs)

        if custom_impl:
            if axes != (0, ):
                raise RuntimeError(
                    "Unexpected axes=%r." % axes)
            ctx['fct'] = lambda x, y: custom_reducesum_rk_float(x)
            ctx['xs'] = [x.reshape((x.shape[0], -1)).copy() for x in xs]
            obs = measure_time(
                "loop_fct(fct, xs, ys)",
                div_by_number=True, context=ctx, repeat=repeat, number=number)
            obs['dim'] = dim
            obs['fct'] = 'custom'
            obs.update(info)
            res.append(obs)

        if tf_reduce_sum is not None:
            # tensorflow
            ctx['fct'] = tf_reduce_sum
            ctx['xs'] = [convert_to_tensor(x) for x in xs]
            ctx['ys'] = ys
            obs = measure_time(
                "loop_fct(fct, xs, ys)",
                div_by_number=True, context=ctx, repeat=repeat, number=number)
            obs['dim'] = dim
            obs['fct'] = 'tf'
            obs.update(info)
            res.append(obs)

        if torch_sum is not None:
            def torch_sum1(x, y):
                return torch_sum(x, y[0])

            def torch_sum2(x, y):
                return torch_sum(torch_sum(x, y[1]), y[0])

            # torch
            ctx['fct'] = torch_sum1 if len(axes) == 1 else torch_sum2
            ctx['xs'] = [from_numpy(x) for x in xs]
            ctx['ys'] = ys  # [from_numpy(y) for y in ys]
            obs = measure_time(
                "loop_fct(fct, xs, ys)",
                div_by_number=True, context=ctx, repeat=repeat, number=number)
            obs['dim'] = dim
            obs['fct'] = 'torch'
            obs.update(info)
            res.append(obs)

    # Dataframes
    shape_name = str(shape).replace(str(dim), "N")
    df = pandas.DataFrame(res)
    df.columns = [_.replace('dim', 'N') for _ in df.columns]
    piv = df.pivot('N', 'fct', 'average')

    rs = piv.copy()
    for c in ['ort', 'torch', 'tf', 'tf_copy']:
        if c in rs.columns:
            rs[c] = rs['numpy'] / rs[c]
    rs['numpy'] = 1.

    # Graphs.
    fig, ax = plt.subplots(1, 2, figsize=(12, 4))
    piv.plot(logx=True, logy=True, ax=ax[0],
             title="%s benchmark\n%r - %r"
                   " lower better" % (name, shape_name, axes))
    ax[0].legend(prop={"size": 9})
    rs.plot(logx=True, logy=True, ax=ax[1],
            title="%s Speedup, baseline=numpy\n%r - %r"
                  " higher better" % (name, shape_name, axes))
    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


dfs = []

Reduction on a particular case KR#

Consecutive axis not reduced and consecutive reduced axis are merged. KR means kept axis - reduced axis

(8, 24, 48, N), axis=(3, )#

axes = (3, )
df, piv, ax = benchmark_op(axes, shape_fct=lambda dim: (8, 24, 48, dim))
dfs.append(df)
df.pivot("fct", "N", "average")
ReduceSum benchmark '(8, 24, 48, N)' - (3,) lower better, ReduceSum Speedup, baseline=numpy '(8, 24, 48, N)' - (3,) higher better

Out:

  0%|          | 0/11 [00:00<?, ?it/s]
  9%|9         | 1/11 [00:02<00:27,  2.76s/it]
 18%|#8        | 2/11 [00:04<00:20,  2.31s/it]
 27%|##7       | 3/11 [00:06<00:18,  2.25s/it]
 36%|###6      | 4/11 [00:09<00:16,  2.35s/it]
 45%|####5     | 5/11 [00:12<00:16,  2.69s/it]
 55%|#####4    | 6/11 [00:16<00:14,  2.95s/it]
 64%|######3   | 7/11 [00:20<00:13,  3.27s/it]
 73%|#######2  | 8/11 [00:24<00:10,  3.55s/it]
 82%|########1 | 9/11 [00:29<00:07,  3.94s/it]
 91%|######### | 10/11 [00:31<00:03,  3.47s/it]
100%|##########| 11/11 [00:35<00:00,  3.56s/it]
100%|##########| 11/11 [00:35<00:00,  3.20s/it]
N 8 16 32 64 100 128 200 256 400 512 1024
fct
numpy 0.004929 0.005906 0.007664 0.009708 0.012695 0.015623 0.019775 0.024812 0.038165 0.019914 0.038882
ort 0.001316 0.001505 0.002065 0.003613 0.005311 0.005671 0.008493 0.010143 0.012658 0.006492 0.013934
torch 0.094325 0.070283 0.073618 0.079018 0.101657 0.100752 0.105172 0.099990 0.092052 0.045838 0.048306


Reduction on a particular case RK#

Consecutive axis not reduced and consecutive reduced axis are merged. RK means reduced axis - kept axis

(8, 24, 48, N), axis=(0, )#

axes = (0, )
df, piv, ax = benchmark_op(axes, shape_fct=lambda dim: (8, 24, 48, dim),
                           custom_impl=True)
dfs.append(df)
df.pivot("fct", "N", "average")
ReduceSum benchmark '(8, 24, 48, N)' - (0,) lower better, ReduceSum Speedup, baseline=numpy '(8, 24, 48, N)' - (0,) higher better

Out:

  0%|          | 0/11 [00:00<?, ?it/s]
  9%|9         | 1/11 [00:01<00:18,  1.84s/it]
 18%|#8        | 2/11 [00:04<00:18,  2.08s/it]
 27%|##7       | 3/11 [00:09<00:27,  3.42s/it]
 36%|###6      | 4/11 [00:13<00:26,  3.83s/it]
 45%|####5     | 5/11 [00:19<00:26,  4.45s/it]
 55%|#####4    | 6/11 [00:26<00:26,  5.35s/it]
 64%|######3   | 7/11 [00:34<00:25,  6.37s/it]
 73%|#######2  | 8/11 [00:44<00:21,  7.33s/it]
 82%|########1 | 9/11 [00:55<00:16,  8.48s/it]
 91%|######### | 10/11 [00:59<00:07,  7.29s/it]
100%|##########| 11/11 [01:06<00:00,  7.22s/it]
100%|##########| 11/11 [01:06<00:00,  6.07s/it]
N 8 16 32 64 100 128 200 256 400 512 1024
fct
custom 0.035236 0.060521 0.049322 0.072903 0.065538 0.071541 0.074178 0.082577 0.082933 0.039643 0.042977
numpy 0.001472 0.002760 0.004691 0.009286 0.013820 0.019101 0.031963 0.040465 0.061406 0.031054 0.061026
ort 0.002256 0.002634 0.004467 0.008639 0.012126 0.015895 0.022726 0.028648 0.043041 0.021753 0.042620
torch 0.033399 0.021862 0.138008 0.079144 0.117278 0.161489 0.184393 0.191790 0.203607 0.067169 0.084775


Reduction on a particular case KRK#

Consecutive axis not reduced and consecutive reduced axis are merged. KRK means kept axis - reduced axis - kept axis,

(8, 24, 48, N), axis=(1, 2)#

axes = (1, 2)
df, piv, ax = benchmark_op(axes, shape_fct=lambda dim: (8, 24, 48, dim))
dfs.append(df)
df.pivot("fct", "N", "average")
ReduceSum benchmark '(8, 24, 48, N)' - (1, 2) lower better, ReduceSum Speedup, baseline=numpy '(8, 24, 48, N)' - (1, 2) higher better

Out:

  0%|          | 0/11 [00:00<?, ?it/s]
  9%|9         | 1/11 [00:02<00:25,  2.54s/it]
 18%|#8        | 2/11 [00:04<00:17,  1.92s/it]
 27%|##7       | 3/11 [00:06<00:15,  1.95s/it]
 36%|###6      | 4/11 [00:09<00:16,  2.37s/it]
 45%|####5     | 5/11 [00:12<00:17,  2.93s/it]
 55%|#####4    | 6/11 [00:17<00:17,  3.43s/it]
 64%|######3   | 7/11 [00:27<00:22,  5.55s/it]
 73%|#######2  | 8/11 [00:38<00:21,  7.23s/it]
 82%|########1 | 9/11 [00:51<00:18,  9.03s/it]
 91%|######### | 10/11 [00:56<00:08,  8.02s/it]
100%|##########| 11/11 [01:05<00:00,  8.35s/it]
100%|##########| 11/11 [01:05<00:00,  5.99s/it]
N 8 16 32 64 100 128 200 256 400 512 1024
fct
numpy 0.003597 0.006729 0.012626 0.024718 0.038416 0.048322 0.072567 0.092406 0.147141 0.076664 0.154081
ort 0.009008 0.002004 0.003502 0.006426 0.008180 0.009815 0.014964 0.018595 0.028174 0.014898 0.032897
torch 0.087827 0.048687 0.059385 0.081032 0.098034 0.102456 0.283950 0.291501 0.295572 0.113852 0.128267


(8, 24 * 48, N), axis=1#

axes = (1, )
df, piv, ax = benchmark_op(axes, shape_fct=lambda dim: (8, 24 * 48, dim))
dfs.append(df)
df.pivot("fct", "N", "average")
ReduceSum benchmark '(8, 1152, N)' - (1,) lower better, ReduceSum Speedup, baseline=numpy '(8, 1152, N)' - (1,) higher better

Out:

  0%|          | 0/11 [00:00<?, ?it/s]
  9%|9         | 1/11 [00:02<00:25,  2.54s/it]
 18%|#8        | 2/11 [00:03<00:16,  1.85s/it]
 27%|##7       | 3/11 [00:05<00:14,  1.80s/it]
 36%|###6      | 4/11 [00:07<00:12,  1.78s/it]
 45%|####5     | 5/11 [00:10<00:13,  2.27s/it]
 55%|#####4    | 6/11 [00:13<00:13,  2.61s/it]
 64%|######3   | 7/11 [00:17<00:12,  3.08s/it]
 73%|#######2  | 8/11 [00:22<00:10,  3.58s/it]
 82%|########1 | 9/11 [00:28<00:08,  4.26s/it]
 91%|######### | 10/11 [00:31<00:03,  3.85s/it]
100%|##########| 11/11 [00:36<00:00,  4.24s/it]
100%|##########| 11/11 [00:36<00:00,  3.30s/it]
N 8 16 32 64 100 128 200 256 400 512 1024
fct
numpy 0.005178 0.006004 0.007908 0.011514 0.014139 0.016274 0.022036 0.026456 0.037860 0.018881 0.034943
ort 0.007758 0.001697 0.002974 0.005531 0.007926 0.009489 0.014751 0.018047 0.028224 0.014550 0.032866
torch 0.087599 0.044776 0.054460 0.045345 0.091173 0.088943 0.100494 0.111157 0.115269 0.059262 0.088189


(2, 8, 12, 24, 2, N), axis=(2, 3)#

axes = (2, 3)
df, piv, ax = benchmark_op(axes, shape_fct=lambda dim: (2, 8, 12, 24, 2, dim))
dfs.append(df)
df.pivot("fct", "N", "average")
ReduceSum benchmark '(2, 8, 12, 24, 2, N)' - (2, 3) lower better, ReduceSum Speedup, baseline=numpy '(2, 8, 12, 24, 2, N)' - (2, 3) higher better

Out:

  0%|          | 0/11 [00:00<?, ?it/s]
  9%|9         | 1/11 [00:01<00:15,  1.55s/it]
 18%|#8        | 2/11 [00:02<00:12,  1.38s/it]
 27%|##7       | 3/11 [00:05<00:14,  1.76s/it]
 36%|###6      | 4/11 [00:08<00:16,  2.40s/it]
 45%|####5     | 5/11 [00:16<00:25,  4.31s/it]
 55%|#####4    | 6/11 [00:24<00:28,  5.71s/it]
 64%|######3   | 7/11 [00:34<00:28,  7.12s/it]
 73%|#######2  | 8/11 [00:44<00:24,  8.17s/it]
 82%|########1 | 9/11 [00:57<00:18,  9.42s/it]
 91%|######### | 10/11 [01:02<00:08,  8.31s/it]
100%|##########| 11/11 [01:12<00:00,  8.59s/it]
100%|##########| 11/11 [01:12<00:00,  6.56s/it]
N 8 16 32 64 100 128 200 256 400 512 1024
fct
numpy 0.003652 0.006716 0.012702 0.025129 0.039174 0.049306 0.075470 0.095763 0.151691 0.076731 0.158002
ort 0.001194 0.001783 0.004115 0.005554 0.007871 0.009620 0.014232 0.018830 0.026383 0.016496 0.033553
torch 0.056021 0.039926 0.067373 0.096344 0.248595 0.262609 0.287274 0.270624 0.260206 0.115724 0.127295


Reduction on a particular case RKRK#

(8, 24, 48, N), axis=(0, 2)#

axes = (0, 2)
df, piv, ax = benchmark_op(axes, shape_fct=lambda dim: (8, 24, 48, dim))
dfs.append(df)
df.pivot("fct", "N", "average")
ReduceSum benchmark '(8, 24, 48, N)' - (0, 2) lower better, ReduceSum Speedup, baseline=numpy '(8, 24, 48, N)' - (0, 2) higher better

Out:

  0%|          | 0/11 [00:00<?, ?it/s]
  9%|9         | 1/11 [00:01<00:14,  1.43s/it]
 18%|#8        | 2/11 [00:03<00:14,  1.65s/it]
 27%|##7       | 3/11 [00:06<00:17,  2.17s/it]
 36%|###6      | 4/11 [00:09<00:17,  2.53s/it]
 45%|####5     | 5/11 [00:13<00:18,  3.04s/it]
 55%|#####4    | 6/11 [00:17<00:18,  3.61s/it]
 64%|######3   | 7/11 [00:27<00:22,  5.61s/it]
 73%|#######2  | 8/11 [00:39<00:22,  7.52s/it]
 82%|########1 | 9/11 [00:53<00:19,  9.52s/it]
 91%|######### | 10/11 [01:00<00:08,  8.88s/it]
100%|##########| 11/11 [01:16<00:00, 11.06s/it]
100%|##########| 11/11 [01:16<00:00,  6.95s/it]
N 8 16 32 64 100 128 200 256 400 512 1024
fct
numpy 0.003739 0.007049 0.013449 0.027419 0.043305 0.054512 0.085771 0.109165 0.168375 0.086897 0.171395
ort 0.001398 0.002418 0.004797 0.009071 0.011402 0.016346 0.026896 0.046692 0.065527 0.076966 0.282709
torch 0.050920 0.060297 0.089363 0.079260 0.090303 0.102800 0.251826 0.277575 0.274097 0.109143 0.136044


Conclusion#

Some of the configurations should be investigated. l-reducesum-problem1. The reduction on tensorflow in one dimension seems to be lazy.

merged = pandas.concat(dfs)
name = "reducesum"
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()
plot op reducesum

Total running time of the script: ( 6 minutes 8.070 seconds)

Gallery generated by Sphinx-Gallery

previous

When to parallelize?

next

Compares implementations of Where