add --cuda_only; fix some nvfuser flags (#960)

Summary:
- Added an option `--cuda_only` to only run cuda benchmarks
- For the nvfuser flags, we've already updated those settings in TIMM. See the latest at https://github.com/rwightman/pytorch-image-models/blob/2f2b22d8c7889174dbf11b92c2d72d8587f9164b/timm/utils/jit.py

Pull Request resolved: #960

Reviewed By: george-qi

Differential Revision: D37187084

Pulled By: xuzhao9

fbshipit-source-id: 75cdaaced599b1816f005f3345b803f3d2400e7c

conftest.py

@@ -16,6 +16,8 @@ def pytest_addoption(parser):
  help="The best attempt to check results for inference runs. Not all models support this!")
  parser.addoption("--cpu_only", action='store_true',
  help="Run benchmarks on cpu only and ignore machine configuration checks")
+ parser.addoption("--cuda_only", action='store_true',
+ help="Run benchmarks on cuda only and ignore machine configuration checks")
 
 def set_fuser(fuser):
  if fuser == "te":
@@ -28,9 +30,7 @@ def set_fuser(fuser):
  torch._C._jit_override_can_fuse_on_gpu(True)
  torch._C._jit_set_texpr_fuser_enabled(False)
  elif fuser == "nvfuser":
- os.environ['PYTORCH_CUDA_FUSER_DISABLE_FALLBACK'] = '1'
- os.environ['PYTORCH_CUDA_FUSER_DISABLE_FMA'] = '1'
- os.environ['PYTORCH_CUDA_FUSER_JIT_OPT_LEVEL'] = '0'
+ os.environ['PYTORCH_NVFUSER_DISABLE_FALLBACK'] = '1'
  torch._C._jit_set_texpr_fuser_enabled(False)
  torch._C._jit_set_profiling_executor(True)
  torch._C._jit_set_profiling_mode(True)

test_bench.py

@@ -31,6 +31,9 @@ def pytest_generate_tests(metafunc):
  if metafunc.config.option.cpu_only:
  devices = ['cpu']
 
+ if metafunc.config.option.cuda_only:
+ devices = ['cuda']
+
  if metafunc.cls and metafunc.cls.__name__ == "TestBenchNetwork":
  paths = _list_model_paths()
  metafunc.parametrize(