Merge pull request #5 from DeepLabCut/starts-picking-up-wherestopped

benchmarking with restart

Author: AlexEMG

benchmarking/run.py

@@ -22,12 +22,12 @@
 pixels = [2500, 10000, 40000, 160000, 320000, 640000]
 ind = 1
 
-for m in dog_models:
-    print("\n\nMODEL {:d} / 8".format(ind))
-    benchmark_model_by_size(m, dog_video, output=out_dir, n_frames=n_frames, pixels=pixels)
-    ind += 1
-
-for m in mouse_models:
-    print("\n\nMODEL {:d} / 8".format(ind))
-    benchmark_model_by_size(m, mouse_video, output=out_dir, n_frames=n_frames, pixels=pixels)
-    ind += 1
+for ind_m, m in enumerate(dog_models):
+    print("\n\nMODEL {:d} / 8".format(ind_m))
+    benchmark_model_by_size(m, dog_video, ind_m, out_dir=out_dir, n_frames=n_frames, pixels=pixels)
+
+offset=ind_m+1
+
+for ind_m, m in enumerate(mouse_models):
+    print("\n\nMODEL {:d} / 8".format(ind_m))
+    benchmark_model_by_size(m, mouse_video, ind_m + offset, out_dir=out_dir, n_frames=n_frames, pixels=pixels)

dlclive/bench.py

@@ -174,14 +174,16 @@ def run_benchmark(model_path, video_path, resize=None, pixels=None, n_frames=100
     resize = resize if resize is not None else 1
     return inf_times, resize*im_size[0] * resize*im_size[1], TFGPUinference
 
+def get_savebenchmarkfn(sys_info ,i, fn_ind, out_dir=None):
+    ''' get filename to save data (definitions see save_benchmark)'''
+    out_dir = out_dir if out_dir is not None else os.getcwd()
+    host_name, op_sys, host_python, dev = sys_info
+
+    base_name = "benchmark_{}_{}_{}_{}.pickle".format(host_name, sys_info['device_type'], fn_ind, i)
+    datafilename = out_dir + '/' + base_name
+    return datafilename
 
-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):
+def save_benchmark(sys_info, inf_times, pixels, i, fn_ind, TFGPUinference, model=None, out_dir=None,datafilename=None):
     """ Save benchmarking data with system information to a pickle file
 
     Parameters
@@ -192,7 +194,7 @@ def save_benchmark(sys_info: dict,
         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: int
+    i: integer
         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
@@ -208,9 +210,12 @@ def save_benchmark(sys_info: dict,
         flag indicating successful save
     """
 
-    out_dir = out_dir if out_dir is not None else os.getcwd()
-    # host_name, op_sys, host_python, dev = sys_info
+    if datafilename is None:
+        #out_dir = out_dir if out_dir is not None else os.getcwd()
+        datafilename=get_savebenchmarkfn(sys_info ,iter, fn_ind, out_dir=out_dir)
 
+    host_name, op_sys, host_python, dev = sys_info
+    host_name = host_name.replace(" ", "")
 
     model_type = None
     if model is not None:
@@ -221,22 +226,15 @@ def save_benchmark(sys_info: dict,
         else:
             model_type = None
 
-    fn_ind = 0
-    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(sys_info['host_name'], sys_info['device_type'], fn_ind,iter)
-
-    fn = os.path.normpath(out_dir)
-
     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'))
+    pickle.dump(data, open(os.path.normpath(datafilename), 'wb'))
 
     return True
 
@@ -245,7 +243,7 @@ def read_pickle(filename):
     with open(filename, "rb") as handle:
         return pickle.load(handle)
 
-def benchmark_model_by_size(model_path, video_path, output=None, n_frames=10000, resize=None, pixels=None, print_rate=False):
+def benchmark_model_by_size(model_path, video_path, fn_ind, out_dir=None, n_frames=10000, resize=None, pixels=None, print_rate=False):
     """Benchmark DLC model by image size
 
     Parameters
@@ -254,7 +252,9 @@ def benchmark_model_by_size(model_path, video_path, output=None, n_frames=10000,
         path to exported DLC model
     video_path : str
         path to video file
-    output : str, optional
+    fn_ind : integer
+        auxiliary variable for creating a unique identifier for saving the models
+    out_dir : str, optional
         directory to save data, will not save data if None, by default None
     n_frames : int, optional
         number of frames to run, by default 10000
@@ -264,7 +264,7 @@ def benchmark_model_by_size(model_path, video_path, output=None, n_frames=10000,
         list of pixel image sizes (as ints), can only use one of resize or pixels, if both specified will use pixels, by default None
     print_rate : bool, optional
         flag to print frame by frame inference rate, by default False
-    
+
     Example
     --------
     Linux/MacOs
@@ -292,21 +292,24 @@ def benchmark_model_by_size(model_path, video_path, output=None, n_frames=10000,
 
     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)
-
-        #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)
+        datafilename=get_savebenchmarkfn(sys_info ,i, fn_ind, out_dir=out_dir)
+
+        #Check if a subset was already been completed?
+        if os.path.isfile(os.path.normpath(datafilename)):
+            print("\nAlready ran {:d} / {:d}\n".format(i+1, len(resize)))
+        else:
+            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)
+
+            ### saving results intermediately
+            save_benchmark(sys_info, inf_times, pixels_out, i, fn_ind, TFGPUinference, model=os.path.basename(model_path), datafilename=datafilename)
 
 
 def main():