This page was generated from docs/reference/lif-benchmarks.ipynb. Interactive online version:

Performance benchmarks for LIF layers in Rockpool

This notebook runs and collects results for all benchmarks provided in rockpool.utilities.benchmarks.

# - Display a description for the benchmarking system
import rockpool
import os
import datetime
import sys

print(f'Benchmarks run on rockpool version {rockpool.__version__}, on date {datetime.datetime.today()}.')
print(f'Python version {sys.version}')
u = os.uname()
print(f'System descriptor: {u[0]} {u[2]} {u[4]}')

try:
    import jax
    print(f'Jax devices: {jax.devices()}')
except:
    print('Jax not found')

try:
    import torch
    is_cuda_available = torch.cuda.is_available()
    print(f'Torch CUDA available: {is_cuda_available}')
    if is_cuda_available:
        print(f'    Number of devices: {torch.cuda.device_count()}')
        print(f'    CUDA devices: {[torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]}')
except:
    pass

Benchmarks run on rockpool version 2.6.dev, on date 2023-04-12 11:50:19.475143.
Python version 3.8.13 (default, Oct 21 2022, 23:50:54)
[GCC 11.2.0]
System descriptor: Linux 5.15.0-69-generic x86_64
Jax devices: [StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0)]
Torch CUDA available: True
    Number of devices: 1
    CUDA devices: ['NVIDIA GeForce RTX 3090 Ti']

# - Run all benchmarks and collect results
from rockpool.utilities.benchmarking import benchmark_neurons, plot_benchmark_results, all_lif_benchmarks
import warnings
import matplotlib.pyplot as plt
plt.rcParams['figure.dpi'] = 300
plt.rcParams['figure.figsize'] = (17, 5)

bench_results = []

for benchmark in all_lif_benchmarks:
    try:
        p_fn, c_fn, e_fn, bench_name = benchmark()
        print(f'Running benchmark "{bench_name}"...')

        bench_results.append(benchmark_neurons(p_fn, c_fn, e_fn, bench_name))

    except Exception as e:
        warnings.warn(f'Benchmark {bench_name} failed with error {str(e)}.')

Running benchmark "LIFExodus on a CUDA device"...

Running benchmark "LIF with no acceleration (numpy backend)"...

Running benchmark "LIFTorch using CUDA graph replay acceleration"...

Running benchmark "LIFTorch on a CUDA device"...

Running benchmark "LIFTorch on CPU"...

Running benchmark "LIFJax, no JIT"...

Running benchmark "LIFJax, with CPU JIT compilation"...

Running benchmark "LIFJax, with GPU JIT compilation"...

Running benchmark "LIFJax, with TPU JIT compilation"...

/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 1 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 2 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 5 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 10 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 20 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 50 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 100 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 500 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 1000 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 2000 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 5000 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(
/home/dylan/rockpool/rockpool/utilities/benchmarking/benchmark_utils.py:128: UserWarning: Benchmarking for layer size 10000 failed with error Backend 'tpu' failed to initialize: module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'.
  warnings.warn(

# - Display benchmark results that succeeded
for results in bench_results:
    if not all([len(r) == 0 for r in results[0]]):
        _, axes = plt.subplots(1, 2)
        plot_benchmark_results(*results, axes)