add profile

dorianhenning · Sep 2, 2018 · 429e255 · 429e255
1 parent f857d4d
commit 429e255
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 13 deletions.
diff --git a/demo.py b/demo.py
@@ -97,18 +97,18 @@
             hm = torch.cat(hm)
             ckpt_time, pose_time = getTime(ckpt_time)
             runtime_profile['pt'].append(pose_time)
-
-            hm = hm.cpu().data
+            hm = hm.cpu()
             writer.save(boxes, scores, hm, pt1, pt2, orig_img, im_name.split('/')[-1])
 
             ckpt_time, post_time = getTime(ckpt_time)
             runtime_profile['pn'].append(post_time)
-
-        # TQDM
-        im_names_desc.set_description(
+
+        if args.profile:
+            # TQDM
+            im_names_desc.set_description(
             'det time: {dt:.3f} | pose time: {pt:.2f} | post processing: {pn:.4f}'.format(
                 dt=np.mean(runtime_profile['dt']), pt=np.mean(runtime_profile['pt']), pn=np.mean(runtime_profile['pn']))
-        )
+            )
 
     print('===========================> Finish Model Running.')
     if (args.save_img or args.save_video) and not args.vis_fast:

diff --git a/opt.py b/opt.py
@@ -20,6 +20,8 @@
                     help='Train with data augmentation')
 parser.add_argument('--sp', default=False, action='store_true',
                     help='Use single process for pytorch')
+parser.add_argument('--profile', default=False, action='store_true',
+                    help='add speed profiling at screen output')
 
 "----------------------------- Model options -----------------------------"
 parser.add_argument('--netType', default='hgPRM', type=str,

diff --git a/video_demo.py b/video_demo.py
@@ -99,11 +99,12 @@
             ckpt_time, post_time = getTime(ckpt_time)
             runtime_profile['pn'].append(post_time)
 
-        # TQDM
-        im_names_desc.set_description(
+        if args.profile:
+            # TQDM
+            im_names_desc.set_description(
             'det time: {dt:.3f} | pose time: {pt:.2f} | post processing: {pn:.4f}'.format(
                 dt=np.mean(runtime_profile['dt']), pt=np.mean(runtime_profile['pt']), pn=np.mean(runtime_profile['pn']))
-        )
+            )
 
     print('===========================> Finish Model Running.')
     if (args.save_img or args.save_video) and not args.vis_fast:

diff --git a/webcam_demo.py b/webcam_demo.py
@@ -104,12 +104,12 @@ def loop():
 
                 ckpt_time, post_time = getTime(ckpt_time)
                 runtime_profile['pn'].append(post_time)
-
-            # TQDM
-            im_names_desc.set_description(
+            if args.profile:
+                # TQDM
+                im_names_desc.set_description(
                 'det time: {dt:.3f} | pose time: {pt:.2f} | post processing: {pn:.4f}'.format(
                     dt=np.mean(runtime_profile['dt']), pt=np.mean(runtime_profile['pt']), pn=np.mean(runtime_profile['pn']))
-            )
+                )
         except KeyboardInterrupt:
             break