Log nvidia hardware, software info

Reviewed By: ir413

Differential Revision: D9360044

fbshipit-source-id: c44c392d065a2cb53db3bbd1d0f83b8ecd80a430

detectron/utils/c2.py

@@ -22,11 +22,13 @@ from __future__ import unicode_literals
 
 from six import string_types
 import contextlib
+import subprocess
 
 from caffe2.proto import caffe2_pb2
 from caffe2.python import core
 from caffe2.python import dyndep
 from caffe2.python import scope
+from caffe2.python import workspace
 
 import detectron.utils.env as envu
 
@@ -145,3 +147,19 @@ def gauss_fill(std):
 def const_fill(value):
     """Constant fill helper to reduce verbosity."""
     return ('ConstantFill', {'value': value})
+
+
+def get_nvidia_info():
+    return (
+        get_nvidia_smi_output(),
+        workspace.GetCUDAVersion(),
+        workspace.GetCuDNNVersion(),
+    )
+
+
+def get_nvidia_smi_output():
+    try:
+        info = subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT)
+    except Exception as e:
+        info = "Executing nvidia-smi failed: " + str(e)
+    return info.strip()

tools/train_net.py

@@ -99,6 +99,10 @@ def main():
     if args.opts is not None:
         merge_cfg_from_list(args.opts)
     assert_and_infer_cfg()
+    smi_output, cuda_ver, cudnn_ver = c2_utils.get_nvidia_info()
+    logger.info("cuda version : {}".format(cuda_ver))
+    logger.info("cudnn version: {}".format(cudnn_ver))
+    logger.info("nvidia-smi output:\n{}".format(smi_output))
     logger.info('Training with config:')
     logger.info(pprint.pformat(cfg))
     # Note that while we set the numpy random seed network training will not be