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)
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])