Merge pull request #4 from DeepLabCut/jls_benchmarking

a bit more info in sys_info

Author: AlexEMG

dlclive/bench.py

@@ -12,43 +12,66 @@
 import sys
 import argparse
 import pickle
-
+import subprocess
+import typing
+import warnings
+
+try:
+    from pip._internal.operations import freeze
+except ImportError:
+    from pip.operations import freeze
 import cpuinfo
 from tqdm import tqdm
 import numpy as np
 import tensorflow as tf
 import cv2
 
 from dlclive import DLCLive
+from dlclive import VERSION
+from dlclive import __file__ as dlcfile
 
 
-def get_system_info():
+def get_system_info() -> dict:
     """ Return summary info for system running benchmark
 
     Returns
     -------
-    str
-        name of machine
-    str
-        operating system
-    str
-        path to python (which conda/virtual environment)
-    str
-        device name
+    dict
+        Dictionary containing the following system information:
+
+        * ``host_name`` (str): name of machine
+        * ``op_sys`` (str): operating system
+        * ``python`` (str): path to python (which conda/virtual environment)
+        * ``device`` (tuple): (device type (``'GPU'`` or ``'CPU'```), device information)
+        * ``freeze`` (list): list of installed packages and versions
+        * ``python_version`` (str): python version
+        * ``git_hash`` (str, None): If installed from git repository, hash of HEAD commit
+        * ``dlclive_version`` (str): dlclive version from :data:`dlclive.VERSION`
     """
 
 
     ### get os
 
     op_sys = platform.platform()
-    host_name = platform.node()
+    host_name = platform.node().replace(' ', '')
 
     # A string giving the absolute path of the executable binary for the Python interpreter, on systems where this makes sense.
     if platform.system() == 'Windows':
         host_python = sys.executable.split(os.path.sep)[-2]
     else:
         host_python = sys.executable.split(os.path.sep)[-3]
 
+    # try to get git hash if possible
+    dlc_basedir = os.path.dirname(os.path.dirname(dlcfile))
+    git_hash = None
+    try:
+        git_hash = subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=dlc_basedir)
+        git_hash = git_hash.decode('utf-8').rstrip('\n')
+    except subprocess.CalledProcessError:
+        # not installed from git repo, eg. pypi
+        # fine, pass quietly
+        pass
+
     ### get device info (GPU or CPU)
 
     dev = None
@@ -63,10 +86,21 @@ def get_system_info():
         dev = get_cpu_info() #[get_cpu_info()['brand']]
         dev_type = "CPU"
 
-    return host_name, op_sys, host_python, (dev_type, dev)
-
-
-def run_benchmark(model_path, video_path, resize=None, pixels=None, n_frames=10000, print_rate=False):
+    # return a dictionary rather than a tuple for inspectability's sake
+    return {
+        'host_name': host_name,
+        'op_sys'   : op_sys,
+        'python': host_python,
+        'device_type': dev_type,
+        'device': dev,
+        'freeze': list(freeze.freeze()), # pip freeze to get versions of all packages
+        'python_version': sys.version,
+        'git_hash': git_hash,
+        'dlclive_version': VERSION
+    }
+
+
+def run_benchmark(model_path, video_path, resize=None, pixels=None, n_frames=10000, print_rate=False) -> typing.Tuple[np.ndarray, int, bool]:
     """ Benchmark on inference times for a given DLC model and video
 
     Parameters
@@ -141,18 +175,24 @@ def run_benchmark(model_path, video_path, resize=None, pixels=None, n_frames=100
     return inf_times, resize*im_size[0] * resize*im_size[1], TFGPUinference
 
 
-def save_benchmark(sys_info, inf_times, pixels, iter , TFGPUinference, model=None, out_dir=None):
+def save_benchmark(sys_info: dict,
+                   inf_times: np.ndarray,
+                   pixels: typing.Union[np.ndarray, float],
+                   iter: int,
+                   TFGPUinference: bool,
+                   model: str = None,
+                   out_dir: str = None):
     """ Save benchmarking data with system information to a pickle file
 
     Parameters
     ----------
-    sys_info : tuple
+    sys_info : dict
         system information generated by :func:`get_system_info`
     inf_times : :class:`numpy.ndarray`
         array of inference times generated by :func:`run_benchmark`
     pixels : float or :class:`numpy.ndarray`
         number of pixels for each benchmark run. If an array, each index corresponds to a row in inf_times
-    iter: integer
+    iter: int
         number of the specific instance of experiment (so every part is saved individually)
     TFGPUinference: bool
         flag if using tensorflow inference or numpy inference DLC model
@@ -169,8 +209,8 @@ def save_benchmark(sys_info, inf_times, pixels, iter , TFGPUinference, model=Non
     """
 
     out_dir = out_dir if out_dir is not None else os.getcwd()
-    host_name, op_sys, host_python, dev = sys_info
-    host_name = host_name.replace(" ", "")
+    # host_name, op_sys, host_python, dev = sys_info
+
 
     model_type = None
     if model is not None:
@@ -182,23 +222,19 @@ def save_benchmark(sys_info, inf_times, pixels, iter , TFGPUinference, model=Non
             model_type = None
 
     fn_ind = 0
-    base_name = "benchmark_{}_{}_{}_{}.pickle".format(host_name, dev[0], fn_ind,iter)
+    base_name = "benchmark_{}_{}_{}_{}.pickle".format(sys_info['host_name'], sys_info['device_type'], fn_ind, iter)
     while os.path.isfile(os.path.normpath(out_dir + '/' + base_name)):
         fn_ind += 1
-        base_name = "benchmark_{}_{}_{}_{}.pickle".format(host_name, dev[0], fn_ind,iter)
+        base_name = "benchmark_{}_{}_{}_{}.pickle".format(sys_info['host_name'], sys_info['device_type'], fn_ind,iter)
 
     fn = os.path.normpath(out_dir)
 
-    data = {'host_name' : host_name,
-            'op_sys' : op_sys,
-            'python' : host_python,
-            'device_type' : dev[0],
-            'device' : dev[1],
-            'model' : model,
+    data = {'model' : model,
             'model_type' : model_type,
             'TFGPUinference' : TFGPUinference,
             'pixels' : pixels,
             'inference_times' : inf_times}
+    data.update(sys_info)
 
     pickle.dump(data, open(os.path.normpath(out_dir + '/' + base_name), 'wb'))
 
@@ -250,21 +286,25 @@ def benchmark_model_by_size(model_path, video_path, output=None, n_frames=10000,
     #inf_times = np.zeros((len(resize), n_frames))
     #pixels_out = np.zeros(len(resize))
     print(resize)
+
+    # get system info once, shouldn't change between runs
+    sys_info = get_system_info()
+
     for i in range(len(resize)):
 
         print("\nRun {:d} / {:d}\n".format(i+1, len(resize)))
 
         inf_times, pixels_out, TFGPUinference = run_benchmark(model_path,
-                                                                    video_path,
-                                                                    resize=resize[i],
-                                                                    pixels=pixels[i],
-                                                                    n_frames=n_frames,
-                                                                    print_rate=print_rate)
+                                                              video_path,
+                                                              resize=resize[i],
+                                                              pixels=pixels[i],
+                                                              n_frames=n_frames,
+                                                              print_rate=print_rate)
 
         #TODO: check if a part has already been complted?
 
         ### saving results intermediately
-        sys_info = get_system_info()
+
         #print("Your system info:", sys_info)
         save_benchmark(sys_info, inf_times, pixels_out, i, TFGPUinference, model=os.path.basename(model_path), out_dir=output)