test_dask_simpler.py

#!/usr/bin/env python
import contextlib
import time

import numpy as np

from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import Dataset
from torch import randn

from napari.utils import resize_dask_cache

import dask
import dask.array as da

resize_dask_cache(0)


@contextlib.contextmanager
def perf_timer(name: str):
    start_ns = time.perf_counter_ns()
    yield
    end_ns = time.perf_counter_ns()
    ms = (end_ns - start_ns) / 1e6
    print(f"{name} {ms}ms")


mnist = MNIST('../data/MNIST', download=True, train=False,)


def make_stack(count):
    return da.stack(
        [
            da.from_delayed(
                dask.delayed(lambda i: np.array(mnist[i][0]))(i),
                shape=(1, 28, 28),
                dtype=np.float32,
            ).reshape((28, 28))
            for i in range(count)
        ]
    )


def test_access(label, stack):
    print(f"Accessing {label}")
    for i in range(3):
        with perf_timer(f"stack[{i}] = "):
            np.asarray(stack[0])


stack_10 = make_stack(10)
stack_100 = make_stack(100)
stack_1000 = make_stack(1000)
stack_10000 = make_stack(10000)

test_access("stack_10", stack_10)
test_access("stack_100", stack_100)
test_access("stack_1000", stack_1000)
test_access("stack_10000", stack_10000)