AIEnergyScore/launch-computation-example

1import argparse2 3from optimum_benchmark.experiment import launch, ExperimentConfig4from optimum_benchmark.backends.pytorch.config import PyTorchConfig5from optimum_benchmark.launchers.torchrun.config import TorchrunConfig6from optimum_benchmark.benchmarks.inference.config import InferenceConfig7 8if __name__ == "__main__":9    parser = argparse.ArgumentParser(description='Run optimum-benchmark')10    parser.add_argument('--config-name', dest='experiment_name', type=str,11                        help='experiment name (text classification, etc.)')12    parser.add_argument('--backend-model', dest='backend_model', type=str,13                        help='model name')14    parser.add_argument('--hydra-run-dir', dest='run_dir', type=str)15 16    args = parser.parse_args()17    backend_model = args.backend_model18    run_dir = args.run_dir19    experiment_name = args.experiment_name20 21    launcher_config = TorchrunConfig(nproc_per_node=2)22    benchmark_config = InferenceConfig(latency=True, memory=True)23    backend_config = PyTorchConfig(model=backend_model)24    experiment_config = ExperimentConfig(25        experiment_name=experiment_name,26        benchmark=benchmark_config,27        launcher=launcher_config,28        backend=backend_config,29    )30 31    benchmark_report = launch(experiment_config)32 33    # push artifacts to the hub34    experiment_config.push_to_hub("EnergyStarAI/benchmarksDebug")35    benchmark_report.push_to_hub("EnergyStarAI/benchmarksDebug")36 37