Upload collect_env.py

utils/collect_env.py

@@ -13,7 +13,7 @@ import torch
 def is_rocm_pytorch() -> bool:
     """Check whether the PyTorch is compiled on ROCm."""
     is_rocm = False
-    if TORCH_VERSION != "parrots":
+    if TORCH_VERSION != 'parrots':
         try:
             from torch.utils.cpp_extension import ROCM_HOME
             is_rocm = True if ((torch.version.hip is not None) and
@@ -26,7 +26,7 @@ TORCH_VERSION = torch.__version__
 
 def get_build_config():
     """Obtain the build information of PyTorch or Parrots."""
-    if TORCH_VERSION == "parrots":
+    if TORCH_VERSION == 'parrots':
         from parrots.config import get_build_info
         return get_build_info()
     else:
@@ -46,7 +46,7 @@ def is_cuda_available() -> bool:
     return torch.cuda.is_available()
 
 def _get_cuda_home():
-    if TORCH_VERSION == "parrots":
+    if TORCH_VERSION == 'parrots':
         from parrots.utils.build_extension import CUDA_HOME
     else:
         if is_rocm_pytorch():
@@ -58,7 +58,7 @@ def _get_cuda_home():
 
 
 def _get_musa_home():
-    return os.environ.get("MUSA_HOME")
+    return os.environ.get('MUSA_HOME')
 
 
 def collect_env():
@@ -84,77 +84,77 @@ def collect_env():
     from distutils import errors
 
     env_info = OrderedDict()
-    env_info["sys.platform"] = sys.platform
-    env_info["Python"] = sys.version.replace("\n", "")
+    env_info['sys.platform'] = sys.platform
+    env_info['Python'] = sys.version.replace('\n', '')
 
     cuda_available = is_cuda_available()
     musa_available = is_musa_available()
-    env_info["CUDA available"] = cuda_available
-    env_info["MUSA available"] = musa_available
-    env_info["numpy_random_seed"] = np.random.get_state()[1][0]
+    env_info['CUDA available'] = cuda_available
+    env_info['MUSA available'] = musa_available
+    env_info['numpy_random_seed'] = np.random.get_state()[1][0]
 
     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
+            env_info['GPU ' + ','.join(device_ids)] = name
 
         CUDA_HOME = _get_cuda_home()
-        env_info["CUDA_HOME"] = CUDA_HOME
+        env_info['CUDA_HOME'] = CUDA_HOME
 
         if CUDA_HOME is not None and osp.isdir(CUDA_HOME):
-            if CUDA_HOME == "/opt/rocm":
+            if CUDA_HOME == '/opt/rocm':
                 try:
-                    nvcc = osp.join(CUDA_HOME, "hip/bin/hipcc")
+                    nvcc = osp.join(CUDA_HOME, 'hip/bin/hipcc')
                     nvcc = subprocess.check_output(
-                        f"\"{nvcc}\" --version", shell=True)
-                    nvcc = nvcc.decode("utf-8").strip()
-                    release = nvcc.rfind("HIP version:")
-                    build = nvcc.rfind("")
+                        f'"{nvcc}" --version', shell=True)
+                    nvcc = nvcc.decode('utf-8').strip()
+                    release = nvcc.rfind('HIP version:')
+                    build = nvcc.rfind('')
                     nvcc = nvcc[release:build].strip()
                 except subprocess.SubprocessError:
-                    nvcc = "Not Available"
+                    nvcc = 'Not Available'
             else:
                 try:
-                    nvcc = osp.join(CUDA_HOME, "bin/nvcc")
-                    nvcc = subprocess.check_output(f"\"{nvcc}\" -V", shell=True)
-                    nvcc = nvcc.decode("utf-8").strip()
-                    release = nvcc.rfind("Cuda compilation tools")
-                    build = nvcc.rfind("Build ")
+                    nvcc = osp.join(CUDA_HOME, 'bin/nvcc')
+                    nvcc = subprocess.check_output(f'"{nvcc}" -V', shell=True)
+                    nvcc = nvcc.decode('utf-8').strip()
+                    release = nvcc.rfind('Cuda compilation tools')
+                    build = nvcc.rfind('Build ')
                     nvcc = nvcc[release:build].strip()
                 except subprocess.SubprocessError:
-                    nvcc = "Not Available"
-            env_info["NVCC"] = nvcc
+                    nvcc = 'Not Available'
+            env_info['NVCC'] = nvcc
     elif musa_available:
         devices = defaultdict(list)
         for k in range(torch.musa.device_count()):
             devices[torch.musa.get_device_name(k)].append(str(k))
         for name, device_ids in devices.items():
-            env_info["GPU " + ",".join(device_ids)] = name
+            env_info['GPU ' + ','.join(device_ids)] = name
 
         MUSA_HOME = _get_musa_home()
-        env_info["MUSA_HOME"] = MUSA_HOME
+        env_info['MUSA_HOME'] = MUSA_HOME
 
         if MUSA_HOME is not None and osp.isdir(MUSA_HOME):
             try:
-                mcc = osp.join(MUSA_HOME, "bin/mcc")
-                subprocess.check_output(f"\"{mcc}\" -v", shell=True)
+                mcc = osp.join(MUSA_HOME, 'bin/mcc')
+                subprocess.check_output(f'"{mcc}" -v', shell=True)
             except subprocess.SubprocessError:
-                mcc = "Not Available"
-            env_info["mcc"] = mcc
+                mcc = 'Not Available'
+            env_info['mcc'] = mcc
     try:
         # Check C++ Compiler.
         # For Unix-like, sysconfig has 'CC' variable like 'gcc -pthread ...',
         # indicating the compiler used, we use this to get the compiler name
         import io
         import sysconfig
-        cc = sysconfig.get_config_var("CC")
+        cc = sysconfig.get_config_var('CC')
         if cc:
             cc = osp.basename(cc.split()[0])
-            cc_info = subprocess.check_output(f"{cc} --version", shell=True)
-            env_info["GCC"] = cc_info.decode("utf-8").partition(
-                "\n")[0].strip()
+            cc_info = subprocess.check_output(f'{cc} --version', shell=True)
+            env_info['GCC'] = cc_info.decode('utf-8').partition(
+                '\n')[0].strip()
         else:
             # on Windows, cl.exe is not in PATH. We need to find the path.
             # distutils.ccompiler.new_compiler() returns a msvccompiler
@@ -165,38 +165,38 @@ def collect_env():
             ccompiler = new_compiler()
             ccompiler.initialize()
             cc = subprocess.check_output(
-                f"{ccompiler.cc}", stderr=subprocess.STDOUT, shell=True)
+                f'{ccompiler.cc}', stderr=subprocess.STDOUT, shell=True)
             encoding = os.device_encoding(
                 sys.stdout.fileno()) or locale.getpreferredencoding()
-            env_info["MSVC"] = cc.decode(encoding).partition("\n")[0].strip()
-            env_info["GCC"] = "n/a"
+            env_info['MSVC'] = cc.decode(encoding).partition('\n')[0].strip()
+            env_info['GCC'] = 'n/a'
     except (subprocess.CalledProcessError, errors.DistutilsPlatformError):
-        env_info["GCC"] = "n/a"
+        env_info['GCC'] = 'n/a'
     except io.UnsupportedOperation as e:
         # JupyterLab on Windows changes sys.stdout, which has no `fileno` attr
         # Refer to: https://github.com/open-mmlab/mmengine/issues/931
         # TODO: find a solution to get compiler info in Windows JupyterLab,
         # while preserving backward-compatibility in other systems.
-        env_info["MSVC"] = f"n/a, reason: {str(e)}"
+        env_info['MSVC'] = f'n/a, reason: {str(e)}'
 
-    env_info["PyTorch"] = torch.__version__
-    env_info["PyTorch compiling details"] = get_build_config()
+    env_info['PyTorch'] = torch.__version__
+    env_info['PyTorch compiling details'] = get_build_config()
 
     try:
         import torchvision
-        env_info["TorchVision"] = torchvision.__version__
+        env_info['TorchVision'] = torchvision.__version__
     except ModuleNotFoundError:
         pass
 
     try:
         import cv2
-        env_info["OpenCV"] = cv2.__version__
+        env_info['OpenCV'] = cv2.__version__
     except ImportError:
         pass
 
 
     return env_info
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     for name, val in collect_env().items():
-        print(f"{name}: {val}")
+        print(f'{name}: {val}')