A performance testing for Chainer's `Iterator` on various `ndarray`s

chainer_iter_perf.py

import sys
import timeit

sys.stdout.write("numpy array: ")
print(
    timeit.timeit(
        "[iter.next() for i in range(0, 60000, 100)]",
        setup="""
import numpy as np
import chainer
train, test = chainer.datasets.get_mnist()
images = train._datasets[0]
iter = chainer.iterators.SerialIterator(images, 100)
""",
        number=20))

sys.stdout.write("cupy array: ")
print(
    timeit.timeit(
        "[iter.next() for i in range(0, 60000, 100)]",
        setup="""
import numpy as np
import chainer
from chainer.backends import cuda
train, test = chainer.datasets.get_mnist()
images = cuda.cupy.array(train._datasets[0])
iter = chainer.iterators.SerialIterator(images, 100)
""",
        number=20))

sys.stdout.write("chainerx (native) array: ")
print(
    timeit.timeit(
        "[iter.next() for i in range(0, 60000, 100)]",
        setup="""
import numpy as np
import chainer
import chainerx as chx
train, test = chainer.datasets.get_mnist()
chx_cpu_device = chx.get_device("native:0")
images = chx.array(train._datasets[0], device=chx_cpu_device)
iter = chainer.iterators.SerialIterator(images, 100)
""",
        number=20))

sys.stdout.write("chainerx (cuda) array: ")
print(
     timeit.timeit(
         "[iter.next() for i in range(0, 60000, 100)]",
         setup="""
import numpy as np
import chainer
import chainerx as chx
train, test = chainer.datasets.get_mnist()
chx_cuda_device = chx.get_device("cuda:0")
images = chx.array(train._datasets[0], device=chx_cuda_device)
iter = chainer.iterators.SerialIterator(images, 100)
""",
         number=20))

sys.stdout.write("genuine chainerx impl with shuffle (on cuda): ")
print(
    timeit.timeit(
        """
np.random.shuffle(all_indices_np)
all_indices = chx.array(all_indices_np, device=cuda_device)
[images.take(all_indices[i:i + 100], axis=0) for i in range(0, 60000, 100)]
""",
        setup="""
import numpy as np
import chainer
import chainerx as chx
all_indices_np = np.arange(60000, dtype=np.int64)
train, test = chainer.datasets.get_mnist()
cuda_device = chx.get_device("cuda:0")
chx.set_default_device(cuda_device)
images = chx.array(train._datasets[0])
""",
        number=20))