Note
Go to the end to download the full example code. or to run this example in your browser via Binder
Benchmark#
Numerical backends performace comparison, with GPU acceleration, as run on Kaggle (https://www.kaggle.com/code/benjaminvial/nannos-cpu)
import os
import matplotlib.pyplot as plt
import numpy as np
Load data
absolute_path = os.path.dirname("__file__")
full_path = os.path.join(absolute_path, "kaggle/nannos_cpu/results_kaggle_cpu.npz")
arch = np.load(full_path, allow_pickle=True)
num_harmo = arch["num_harmo_real"]
timings = arch["timedict"].item()
# cases = arch["cases"].tolist()
full_path = os.path.join(absolute_path, "kaggle/nannos_gpu/results_kaggle_gpu.npz")
arch = np.load(full_path, allow_pickle=True)
timings.update(arch["timedict"].item())
# cases += arch["cases"].tolist()
skip_first = True
ifirst = 1 if skip_first else 0
cases = ["numpy", "scipy", "autograd", "torch cpu", "torch gpu", "jax cpu", "jax gpu"]
Time vs. number of harmonics#
tav = {}
tstd = {}
for case in cases:
t = np.array(timings[case])[:, ifirst:]
tav[case] = np.mean(t, axis=1)
tstd[case] = np.std(t, axis=1)
# colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
colors = ["#467ccc", "#8746cc", "#cc4646", "#e4731d", "#46cc96", "#525252"]
colors = colors[:3] + [colors[4], colors[4]] + [colors[5], colors[5]]
plt.figure(figsize=(2, 2))
for i, case in enumerate(cases):
ls = "--" if case.split(" ")[-1] == "gpu" else "-"
plt.errorbar(
num_harmo,
tav[case],
yerr=tstd[case],
label=case,
ls=ls,
color=colors[i],
fmt="s",
capsize=1,
ms=1,
)
plt.xlabel("Number of harmonics")
plt.ylabel("CPU time (s)")
plt.xscale("log")
plt.yscale("log")
plt.legend()
plt.tight_layout()
Speedup vs. number of harmonics#
speedup_av = {}
speedup_std = {}
for case in cases[1:]:
s = np.array(timings["numpy"])[:, ifirst:] / np.array(timings[case])[:, ifirst:]
speedup_av[case] = np.mean(s, axis=1)
speedup_std[case] = np.std(s, axis=1)
plt.figure(figsize=(2, 2))
for i, case in enumerate(cases[1:]):
ls = "--" if case.split(" ")[-1] == "gpu" else "-"
plt.errorbar(
num_harmo,
speedup_av[case],
yerr=speedup_std[case],
label=case,
ls=ls,
color=colors[i + 1],
fmt="s",
capsize=1,
ms=1,
)
plt.xlabel("number of harmonics")
plt.ylabel("Speedup vs. numpy")
plt.legend()
plt.tight_layout()
Total running time of the script: (0 minutes 1.404 seconds)
Estimated memory usage: 519 MB