Function iteration_test

dali/test/python/test_RN50_external_source_parallel_data.py:32–99

(args)

def iteration_test(args):
    test_pipe_factories = get_pipe_factories(
        args.test_pipes,
        external_source_parallel_pipeline,
        file_reader_pipeline,
        external_source_pipeline,
    )
    for pipe_factory in test_pipe_factories:
        # TODO(klecki): We don't handle sharding in this test yet, would need to do it manually
        # for External Source pipelines
        pipes = [
            pipe_factory(
                batch_size=args.batch_size,
                num_threads=args.workers,
                device_id=gpu,
                data_path=args.data_path,
                prefetch_queue_depth=args.prefetch,
                reader_queue_depth=args.reader_queue_depth,
                py_start_method=args.worker_init,
                py_num_workers=args.py_workers,
                source_mode=args.source_mode,
                read_encoded=args.dali_decode,
            )
            for gpu in range(args.gpus)
        ]
        # First start the Python workers, so we fork without CUDA context.
        for pipe in pipes:
            pipe.start_py_workers()
        for pipe in pipes:
            pipe.build()

        samples_no = pipes[0].epoch_size("Reader")
        if args.benchmark_iters is None:
            expected_iters = samples_no // args.batch_size + (samples_no % args.batch_size != 0)
        else:
            expected_iters = args.benchmark_iters

        print("RUN {}".format(pipe_factory.__name__))
        for i in range(args.epochs):
            if i == 0:
                print("Warm up")
            else:
                print("Test run " + str(i))
            data_time = AverageMeter()
            end = time.time()
            frequency = 50
            for j in range(expected_iters):
                stop_iter = False
                for pipe in pipes:
                    try:
                        pipe.run()
                    except StopIteration:
                        assert j == expected_iters - 1
                        stop_iter = True
                if stop_iter:
                    break
                if j % frequency == 0 and j != 0:
                    data_time.update((time.time() - end) / frequency)