[feature] Add collect_env() to collect environment info (#517)

* [feature] Add collect_env() to collect environment info

* fix unit tests

* fix the case when ops are not compiled

* fix docstring

mmcv/utils/__init__.py

@@ -26,19 +26,18 @@ except ImportError:
         'get_git_hash'
     ]
 else:
-    from .env import TORCH_VERSION
+    from .env import collect_env
     from .logging import get_logger, print_log
-    from .parrots_wrapper import (CUDA_HOME, BuildExtension, CppExtension,
-                                  CUDAExtension, DataLoader, PoolDataLoader,
-                                  SyncBatchNorm, _AdaptiveAvgPoolNd,
-                                  _AdaptiveMaxPoolNd, _AvgPoolNd, _BatchNorm,
-                                  _ConvNd, _ConvTransposeMixin, _InstanceNorm,
-                                  _MaxPoolNd, get_build_config)
+    from .parrots_wrapper import (
+        CUDA_HOME, TORCH_VERSION, BuildExtension, CppExtension, CUDAExtension,
+        DataLoader, PoolDataLoader, SyncBatchNorm, _AdaptiveAvgPoolNd,
+        _AdaptiveMaxPoolNd, _AvgPoolNd, _BatchNorm, _ConvNd,
+        _ConvTransposeMixin, _InstanceNorm, _MaxPoolNd, get_build_config)
     from .registry import Registry, build_from_cfg
     __all__ = [
-        'Config', 'ConfigDict', 'DictAction', 'get_logger', 'print_log',
-        'is_str', 'iter_cast', 'list_cast', 'tuple_cast', 'is_seq_of',
-        'is_list_of', 'is_tuple_of', 'slice_list', 'concat_list',
+        'Config', 'ConfigDict', 'DictAction', 'collect_env', 'get_logger',
+        'print_log', 'is_str', 'iter_cast', 'list_cast', 'tuple_cast',
+        'is_seq_of', 'is_list_of', 'is_tuple_of', 'slice_list', 'concat_list',
         'check_prerequisites', 'requires_package', 'requires_executable',
         'is_filepath', 'fopen', 'check_file_exist', 'mkdir_or_exist',
         'symlink', 'scandir', 'ProgressBar', 'track_progress',

mmcv/utils/env.py

-# This file holding some environment constant for sharing by other files
+"""This file holding some environment constant for sharing by other files."""
+
+import os.path as osp
+import subprocess
+import sys
+from collections import defaultdict
+
+import cv2
 import torch
 
-TORCH_VERSION = torch.__version__
+import mmcv
+from .parrots_wrapper import get_build_config
+
+
+def collect_env():
+    """Collect the information of the running environments.
+
+    Returns:
+        dict: The environment information. The following fields are contained.
+
+            - sys.platform: The variable of ``sys.platform``.
+            - Python: Python version.
+            - CUDA available: Bool, indicating if CUDA is available.
+            - GPU devices: Device type of each GPU.
+            - CUDA_HOME (optional): The env var ``CUDA_HOME``.
+            - NVCC (optional): NVCC version.
+            - GCC: GCC version, "n/a" if GCC is not installed.
+            - PyTorch: PyTorch version.
+            - PyTorch compiling details: The output of \
+                ``torch.__config__.show()``.
+            - TorchVision (optional): TorchVision version.
+            - OpenCV: OpenCV version.
+            - MMCV: MMCV version.
+            - MMCV Compiler: The GCC version for compiling MMCV ops.
+            - MMCV CUDA Compiler: The CUDA version for compiling MMCV ops.
+    """
+    env_info = {}
+    env_info['sys.platform'] = sys.platform
+    env_info['Python'] = sys.version.replace('\n', '')
+
+    cuda_available = torch.cuda.is_available()
+    env_info['CUDA available'] = cuda_available
+
+    if cuda_available:
+        devices = defaultdict(list)
+        for k in range(torch.cuda.device_count()):
+            devices[torch.cuda.get_device_name(k)].append(str(k))
+        for name, device_ids in devices.items():
+            env_info['GPU ' + ','.join(device_ids)] = name
+
+        from torch.utils.cpp_extension import CUDA_HOME
+        env_info['CUDA_HOME'] = CUDA_HOME
+
+        if CUDA_HOME is not None and osp.isdir(CUDA_HOME):
+            try:
+                nvcc = osp.join(CUDA_HOME, 'bin/nvcc')
+                nvcc = subprocess.check_output(
+                    f'"{nvcc}" -V | tail -n1', shell=True)
+                nvcc = nvcc.decode('utf-8').strip()
+            except subprocess.SubprocessError:
+                nvcc = 'Not Available'
+            env_info['NVCC'] = nvcc
+
+    try:
+        gcc = subprocess.check_output('gcc --version | head -n1', shell=True)
+        gcc = gcc.decode('utf-8').strip()
+        env_info['GCC'] = gcc
+    except subprocess.CalledProcessError:  # gcc is unavailable
+        env_info['GCC'] = 'n/a'
+
+    env_info['PyTorch'] = torch.__version__
+    env_info['PyTorch compiling details'] = get_build_config()
+
+    try:
+        import torchvision
+        env_info['TorchVision'] = torchvision.__version__
+    except ModuleNotFoundError:
+        pass
+
+    env_info['OpenCV'] = cv2.__version__
+
+    env_info['MMCV'] = mmcv.__version__
+
+    try:
+        from mmcv.ops import get_compiler_version, get_compiling_cuda_version
+    except ModuleNotFoundError:
+        env_info['MMCV Compiler'] = 'n/a'
+        env_info['MMCV CUDA Compiler'] = 'n/a'
+    else:
+        env_info['MMCV Compiler'] = get_compiler_version()
+        env_info['MMCV CUDA Compiler'] = get_compiling_cuda_version()
+
+    return env_info

mmcv/utils/parrots_wrapper.py

@@ -2,7 +2,7 @@ from functools import partial
 
 import torch
 
-from .env import TORCH_VERSION
+TORCH_VERSION = torch.__version__
 
 
 def _get_cuda_home():

tests/test_env.py

+import sys
+
+import pytest
+
+import mmcv
+
+
+def test_collect_env():
+    try:
+        import torch  # noqa: F401
+    except ModuleNotFoundError:
+        pytest.skip('skipping tests that require PyTorch')
+
+    from mmcv.utils import collect_env
+    env_info = collect_env()
+    expected_keys = [
+        'sys.platform', 'Python', 'CUDA available', 'PyTorch',
+        'PyTorch compiling details', 'OpenCV', 'MMCV', 'MMCV Compiler',
+        'MMCV CUDA Compiler'
+    ]
+    for key in expected_keys:
+        assert key in env_info
+
+    if env_info['CUDA available']:
+        for key in ['CUDA_HOME', 'NVCC']:
+            assert key in env_info
+
+    if sys.platform != 'win32':
+        assert 'GCC' in env_info
+
+    assert env_info['sys.platform'] == sys.platform
+    assert env_info['Python'] == sys.version.replace('\n', '')
+    assert env_info['MMCV'] == mmcv.__version__