Add benchmarks.py to run all benchmarks, add new file with all torchbench model names (#94146)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/94146
Approved by: https://github.com/ezyang

benchmarks/dynamo/all_torchbench_models_list.txt

@@ -0,0 +1,73 @@
+BERT_pytorch
+Background_Matting
+DALLE2_pytorch
+LearningToPaint
+Super_SloMo
+alexnet
+attention_is_all_you_need_pytorch
+dcgan
+demucs
+densenet121
+detectron2_fasterrcnn_r_101_c4
+detectron2_fasterrcnn_r_101_dc5
+detectron2_fasterrcnn_r_101_fpn
+detectron2_fasterrcnn_r_50_c4
+detectron2_fasterrcnn_r_50_dc5
+detectron2_fasterrcnn_r_50_fpn
+detectron2_fcos_r_50_fpn
+detectron2_maskrcnn
+detectron2_maskrcnn_r_101_c4
+detectron2_maskrcnn_r_101_fpn
+detectron2_maskrcnn_r_50_c4
+detectron2_maskrcnn_r_50_fpn
+dlrm
+drq
+fambench_dlrm
+fambench_xlmr
+fastNLP_Bert
+hf_Albert
+hf_Bart
+hf_Bert
+hf_BigBird
+hf_DistilBert
+hf_GPT2
+hf_Longformer
+hf_Reformer
+hf_T5
+maml
+maml_omniglot
+mnasnet1_0
+mobilenet_v2
+mobilenet_v2_quantized_qat
+mobilenet_v3_large
+moco
+nvidia_deeprecommender
+opacus_cifar10
+pplbench_beanmachine
+pyhpc_equation_of_state
+pyhpc_isoneutral_mixing
+pyhpc_turbulent_kinetic_energy
+pytorch_CycleGAN_and_pix2pix
+pytorch_stargan
+pytorch_struct
+pytorch_unet
+resnet18
+resnet50
+resnet50_quantized_qat
+resnext50_32x4d
+shufflenet_v2_x1_0
+soft_actor_critic
+speech_transformer
+squeezenet1_1
+tacotron2
+timm_efficientdet
+timm_efficientnet
+timm_nfnet
+timm_regnet
+timm_resnest
+timm_vision_transformer
+timm_vovnet
+tts_angular
+vgg16
+vision_maskrcnn
+yolov3

benchmarks/dynamo/benchmarks.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+import argparse
+import os
+
+from typing import Set
+
+# Note - hf and timm have their own version of this, torchbench does not
+# TOOD(voz): Someday, consolidate all the files into one runner instead of a shim like this...
+def model_names(filename: str) -> Set[str]:
+    names = set()
+    with open(filename, "r") as fh:
+        lines = fh.readlines()
+        lines = [line.rstrip() for line in lines]
+        for line in lines:
+            line_parts = line.split(" ")
+            if len(line_parts) == 1:
+                line_parts = line.split(",")
+            model_name = line_parts[0]
+            names.add(model_name)
+    return names
+
+
+TIMM_MODEL_NAMES = model_names(
+    os.path.join(os.path.dirname(__file__), "timm_models_list.txt")
+)
+HF_MODELS_FILE_NAME = model_names(
+    os.path.join(os.path.dirname(__file__), "huggingface_models_list.txt")
+)
+TORCHBENCH_MODELS_FILE_NAME = model_names(
+    os.path.join(os.path.dirname(__file__), "all_torchbench_models_list.txt")
+)
+
+# timm <> HF disjoint
+assert TIMM_MODEL_NAMES.isdisjoint(HF_MODELS_FILE_NAME)
+# timm <> torch disjoint
+assert TIMM_MODEL_NAMES.isdisjoint(TORCHBENCH_MODELS_FILE_NAME)
+# torch <> hf disjoint
+assert TORCHBENCH_MODELS_FILE_NAME.isdisjoint(HF_MODELS_FILE_NAME)
+
+
+def parse_args(args=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--only",
+        help="""Run just one model from whichever model suite it belongs to. Or
+        specify the path and class name of the model in format like:
+        --only=path:<MODEL_FILE_PATH>,class:<CLASS_NAME>
+
+        Due to the fact that dynamo changes current working directory,
+        the path should be an absolute path.
+
+        The class should have a method get_example_inputs to return the inputs
+        for the model. An example looks like
+        ```
+        class LinearModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(10, 10)
+
+            def forward(self, x):
+                return self.linear(x)
+
+            def get_example_inputs(self):
+                return (torch.randn(2, 10),)
+        ```
+    """,
+    )
+    return parser.parse_known_args(args)
+
+
+if __name__ == "__main__":
+    args, unknown = parse_args()
+    if args.only:
+        name = args.only
+        if name in TIMM_MODEL_NAMES:
+            import timm_models
+
+            timm_models.timm_main()
+        elif name in HF_MODELS_FILE_NAME:
+            import huggingface
+
+            huggingface.huggingface_main()
+        elif name in TORCHBENCH_MODELS_FILE_NAME:
+            import torchbench
+
+            torchbench.torchbench_main()
+        else:
+            print(f"Illegal model name? {name}")
+            exit(-1)
+    else:
+        import torchbench
+
+        torchbench.torchbench_main()
+
+        import huggingface
+
+        huggingface.huggingface_main()
+
+        import timm_models
+
+        timm_models.timm_main()

benchmarks/dynamo/huggingface.py

@@ -582,10 +582,14 @@ def refresh_model_names_and_batch_sizes():
             log.warning(f"Failed to find suitable batch size for {model_name}")
 
 
-if __name__ == "__main__":
+def huggingface_main():
     # Code to refresh model names and batch sizes
     # if "--find-batch-sizes" not in sys.argv:
     #     refresh_model_names_and_batch_sizes()
     logging.basicConfig(level=logging.WARNING)
     warnings.filterwarnings("ignore")
     main(HuggingfaceRunner())
+
+
+if __name__ == "__main__":
+    huggingface_main()

benchmarks/dynamo/run_all.sh

@@ -32,10 +32,7 @@ WORK="$PWD"
 
 cd "$(dirname "$BASH_SOURCE")"/../..
 
-python benchmarks/dynamo/torchbench.py --output "$WORK"/torchbench.csv "${BASE_FLAGS[@]}" "$@" 2>&1 | tee "$WORK"/torchbench.log
-python benchmarks/dynamo/huggingface.py --output "$WORK"/huggingface.csv "${BASE_FLAGS[@]}" "$@" 2>&1 | tee "$WORK"/huggingface.log
-python benchmarks/dynamo/timm_models.py --output "$WORK"/timm_models.csv "${BASE_FLAGS[@]}" "$@" 2>&1 | tee "$WORK"/timm_models.log
-cat "$WORK"/torchbench.log "$WORK"/huggingface.log "$WORK"/timm_models.log | tee "$WORK"/sweep.log
+python benchmarks/dynamo/benchmarks.py --output "$WORK"/benchmarks.csv "${BASE_FLAGS[@]}" "$@" 2>&1 | tee "$WORK"/sweep.log
 gh gist create -d "Sweep logs for $(git rev-parse --abbrev-ref HEAD) $* - $(git rev-parse HEAD) $DATE" "$WORK"/sweep.log | tee -a "$WORK"/sweep.log
 python benchmarks/dynamo/parse_logs.py "$WORK"/sweep.log > "$WORK"/final.csv
 gh gist create "$WORK"/final.csv

benchmarks/dynamo/timm_models.py

@@ -337,7 +337,11 @@ class TimmRunnner(BenchmarkRunner):
         return None
 
 
-if __name__ == "__main__":
+def timm_main():
     logging.basicConfig(level=logging.WARNING)
     warnings.filterwarnings("ignore")
     main(TimmRunnner())
+
+
+if __name__ == "__main__":
+    timm_main()

benchmarks/dynamo/torchbench.py

@@ -374,9 +374,12 @@ class TorchBenchmarkRunner(BenchmarkRunner):
         return None
 
 
-if __name__ == "__main__":
-
+def torchbench_main():
     original_dir = setup_torchbench_cwd()
     logging.basicConfig(level=logging.WARNING)
     warnings.filterwarnings("ignore")
     main(TorchBenchmarkRunner(), original_dir)
+
+
+if __name__ == "__main__":
+    torchbench_main()