def main()

in petastorm/benchmark/dummy_reader.py [0:0]

• 29 lines of code
• 7 McCabe index (conditional complexity)

def main(device='cpu', batch=1000, dim=64):
    print("Testing DataLoader on cpu")
    reader = DummyReader(int(batch), int(dim))

    for batch_size in [10, 100, 1000]:
        iterations = 100
        loader = DataLoader(reader, shuffling_queue_capacity=batch_size * 10, batch_size=batch_size)
        it = iter(loader)

        # Warmup
        for _ in range(iterations):
            next(it)
        print("Done warming up")

        tstart = time.time()
        for _ in range(iterations):
            next(it)
        print("Samples per second for batch {}: {:.4g}".format(
            batch_size, (iterations * batch_size) / (time.time() - tstart)))

    print("Testing BatchedDataLoader on", device)
    for batch_size in [10, 100, 1000, 100000]:
        iterations = 100
        loader = BatchedDataLoader(reader, shuffling_queue_capacity=batch_size * 10, batch_size=batch_size,
                                   transform_fn=partial(torch.as_tensor, device=device))
        it = iter(loader)

        # Warmup
        for _ in range(iterations):
            next(it)
        print("Done warming up")

        tstart = time.time()
        for _ in range(iterations):
            next(it)
        print("Samples per second for batch {}: {:.4g}".format(
            batch_size, (iterations * batch_size) / (time.time() - tstart)))