def create_benchmark(params='--num_steps 8'):
    """Register and create benchmark."""
    # Register the FakeModelBenchmark benchmark.
    BenchmarkRegistry.register_benchmark(
        'pytorch-fake-model',
        FakeModelBenchmark,
        parameters='--hidden_size 2',
        platform=Platform.CUDA,
    )
    context = BenchmarkRegistry.create_benchmark_context(
        'fake-model',
        platform=Platform.CUDA,
        parameters=params,
        framework=Framework.PYTORCH)
    name = BenchmarkRegistry._BenchmarkRegistry__get_benchmark_name(context)
    assert (name)
    (benchmark_class, predefine_params
     ) = BenchmarkRegistry._BenchmarkRegistry__select_benchmark(
         name, context.platform)
    assert (benchmark_class)
    return benchmark_class(name, predefine_params + ' ' + context.parameters)
Example #2
0
def test_get_benchmark_name():
    """Test interface BenchmarkRegistry.get_benchmark_name()."""
    # Register benchmarks for testing.
    benchmark_names = [
        'accumulation', 'pytorch-accumulation', 'tf1-accumulation',
        'onnxruntime-accumulation'
    ]
    for name in benchmark_names:
        BenchmarkRegistry.register_benchmark(name, AccumulationBenchmark)

    # Test benchmark name for different Frameworks.
    benchmark_frameworks = [
        Framework.NONE, Framework.PYTORCH, Framework.TENSORFLOW1,
        Framework.ONNXRUNTIME
    ]
    for i in range(len(benchmark_names)):
        context = BenchmarkRegistry.create_benchmark_context(
            'accumulation',
            platform=Platform.CPU,
            framework=benchmark_frameworks[i])
        name = BenchmarkRegistry._BenchmarkRegistry__get_benchmark_name(
            context)
        assert (name == benchmark_names[i])