Disable C++ models from being compiled without explicit request (#1535)

* Disable C++ models from being compiled without explicitly being asked for

* Fix import in tests, which are already disabled

setup.py

@@ -88,6 +88,8 @@ def get_extensions():
     sources = main_file + source_cpu
     extension = CppExtension
 
+    compile_cpp_tests = os.getenv('WITH_CPP_MODELS_TEST', '0') == '1'
+    if compile_cpp_tests:
         test_dir = os.path.join(this_dir, 'test')
         models_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'models')
         test_file = glob.glob(os.path.join(test_dir, '*.cpp'))
@@ -96,6 +98,7 @@ def get_extensions():
         test_file = [os.path.join(test_dir, s) for s in test_file]
         source_models = [os.path.join(models_dir, s) for s in source_models]
         tests = test_file + source_models
+        tests_include_dirs = [test_dir, models_dir]
 
     define_macros = []
 
@@ -123,7 +126,6 @@ def get_extensions():
     sources = [os.path.join(extensions_dir, s) for s in sources]
 
     include_dirs = [extensions_dir]
-    tests_include_dirs = [test_dir, models_dir]
 
     ffmpeg_exe = distutils.spawn.find_executable('ffmpeg')
     has_ffmpeg = ffmpeg_exe is not None
@@ -143,15 +145,18 @@ def get_extensions():
             include_dirs=include_dirs,
             define_macros=define_macros,
             extra_compile_args=extra_compile_args,
-        ),
+        )
+    ]
+    if compile_cpp_tests:
+        ext_modules.append(
             extension(
                 'torchvision._C_tests',
                 tests,
                 include_dirs=tests_include_dirs,
                 define_macros=define_macros,
                 extra_compile_args=extra_compile_args,
-        ),
-    ]
+            )
+        )
     if has_ffmpeg:
         ext_modules.append(
             CppExtension(

test/test_cpp_models.py

 import torch
 import os
 import unittest
-from torchvision import models, transforms, _C_tests
+from torchvision import models, transforms
 import sys
 
 from PIL import Image
 import torchvision.transforms.functional as F
 
+try:
+    from torchvision import _C_tests
+except ImportError:
+    _C_tests = None
+
 
 def process_model(model, tensor, func, name):
     model.eval()