Skip to content

Instantly share code, notes, and snippets.

@pwinston
Last active July 10, 2020 23:09
Show Gist options
  • Save pwinston/0e62b5444dd1aa386eeb22f17d4bc4f8 to your computer and use it in GitHub Desktop.
Save pwinston/0e62b5444dd1aa386eeb22f17d4bc4f8 to your computer and use it in GitHub Desktop.
#!/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)
"""
Sample Output
Accessing stack_10
stack[0] = 9.158942ms
stack[1] = 1.586809ms
stack[2] = 1.527705ms
Accessing stack_100
stack[0] = 2.825355ms
stack[1] = 2.451523ms
stack[2] = 2.498721ms
Accessing stack_1000
stack[0] = 17.305835ms
stack[1] = 14.243399ms
stack[2] = 13.712595ms
Accessing stack_10000
stack[0] = 262.396553ms
stack[1] = 259.603917ms
stack[2] = 174.144542ms
"""
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment