add more info back to BenchResult (pytorch#21502)

Summary:
Pull Request resolved: pytorch#21502

In BenchResult, we keep name, avg_fwd, std_fwd, avg_bwd, and std_bwd. There is no information about the number of each iteration. In this diff, I am adding more info to BenchResult to include the number reported from each iteration.

Reviewed By: wanchaol

Differential Revision: D15706306

fbshipit-source-id: 3f14be4ba91f1f6da473995783bd7af1d067938d

Author: hl475

benchmarks/fastrnns/bench.py

@@ -11,7 +11,7 @@
 
 
 BenchResult = namedtuple('BenchResult', [
-    'name', 'avg_fwd', 'std_fwd', 'avg_bwd', 'std_bwd',
+    'name', 'avg_fwd', 'std_fwd', 'info_fwd', 'avg_bwd', 'std_bwd', 'info_bwd',
 ])
 
 
@@ -97,12 +97,13 @@ def train_batch(modeldef):
 
     fwd_times = torch.tensor(fwd_times)
     bwd_times = torch.tensor(bwd_times)
-
     return BenchResult(name=name,
                        avg_fwd=fwd_times.mean().item(),
                        std_fwd=fwd_times.std().item(),
+                       info_fwd=fwd_times,
                        avg_bwd=bwd_times.mean().item(),
-                       std_bwd=bwd_times.std().item())
+                       std_bwd=bwd_times.std().item(),
+                       info_bwd=bwd_times)
 
 
 def print_stderr(*args, **kwargs):
@@ -121,23 +122,22 @@ def print_json_oss_format(results):
     print(json.dumps(oss_results))
 
 
-def print_json_pep_format(num_iters, results):
+def print_json_pep_format(results):
     # print the AI-PEP format json string for each model
     for group_name, group_val in results.items():
         for model_name, run_time in group_val.items():
             # Output for AI-PEP
-            print("Caffe2Observer " + json.dumps(
-                {
-                    "type": "NET",
-                    "metric": group_name + "-" + model_name,
-                    "unit": "ms",
-                    "num_runs": num_iters,
-                    "summary": {
-                        "mean": run_time['avg'],
-                        "stdev": run_time['std']
+            num_iters = len(run_time['info'])
+            info = run_time['info'].tolist()
+            for i in range(num_iters):
+                print("Caffe2Observer " + json.dumps(
+                    {
+                        "type": "NET",
+                        "metric": group_name + "-" + model_name,
+                        "unit": "ms",
+                        "value": str(info[i])
                     }
-                }
-            ))
+                ))
 
 
 def bench(rnn_runners, group_name, print_json=False, sep=' ', **params):
@@ -147,15 +147,18 @@ def bench(rnn_runners, group_name, print_json=False, sep=' ', **params):
         with context():
             try:
                 result = trainbench(name, creator, **params)
-                print_stderr(pretty_print(result, sep=sep))
+                # Replace the value of info_fwd and info_bwd to None
+                result_with_no_info = result._replace(
+                    info_fwd='None', info_bwd='None')
+                print_stderr(pretty_print(result_with_no_info, sep=sep))
                 results[name] = result
             except Exception as e:
                 if not print_json:
                     raise
 
     return {
-        group_name: {k: {"avg": v.avg_fwd, "std": v.std_fwd} for k, v in results.items()},
-        group_name + '-backward': {k: {"avg": v.avg_bwd, "std": v.std_bwd} for k, v in results.items()},
+        group_name: {k: {"avg": v.avg_fwd, "std": v.std_fwd, "info": v.info_fwd} for k, v in results.items()},
+        group_name + '-backward': {k: {"avg": v.avg_bwd, "std": v.std_bwd, "info": v.info_bwd} for k, v in results.items()},
     }
 
 
@@ -230,4 +233,4 @@ def bench_group(model_list, bench_name, bench_group, bench_args):
     if args.print_json == 'oss':
         print_json_oss_format(results)
     elif args.print_json == 'pep':
-        print_json_pep_format(args.nloops, results)
+        print_json_pep_format(results)