polish memory opt benchmark (#198)

Tencent · Nov 19, 2020 · ecaf698 · ecaf698
1 parent 674f08d
commit ecaf698
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/benchmark/benchmark.py b/benchmark/benchmark.py
@@ -14,7 +14,7 @@
 turbo-transformers Benchmark Utils
 
 Usage:
-    benchmark <model_name> [--seq_len=<int>] [--framework=<f>] [--batch_size=<int>] [-n <int>] [--enable-random] [--min_seq_len=<int>] [--max_seq_len=<int>] [--use_gpu] [--num_threads=<int>] [--enable_mem_opt=<bool>]
+    benchmark <model_name> [--seq_len=<int>] [--framework=<f>] [--batch_size=<int>] [-n <int>] [--enable-random] [--min_seq_len=<int>] [--max_seq_len=<int>] [--use_gpu] [--num_threads=<int>] [--enable_mem_opt]
 
 Options:
     --framework=<f>      The framework to test in (torch, torch_jit, turbo-transformers,
@@ -27,7 +27,7 @@
     --max_seq_len=<int>  Maximal sequence length generated when enable random [default: 50]
     --use_gpu            Enable GPU.
     --num_threads=<int>  The number of CPU threads. [default: 4]
-    --enable_mem_opt=<bool>  Use memory optimization for BERT. [default: False]
+    --enable_mem_opt     Use model aware memory optimization for BERT.
 """
 
 import json
@@ -54,7 +54,8 @@ def main():
         'use_gpu': True if args['--use_gpu'] else False,
         'enable_mem_opt': True if args['--enable_mem_opt'] else False,
     }
-    if (kwargs['model_name'] != 'bert'):
+    if (kwargs['model_name'] != 'bert'
+            or args['--framework'] != 'turbo-transformers'):
         kwargs['enable_mem_opt'] = False
     if args['--framework'] == 'turbo-transformers':
         benchmark_turbo_transformers(**kwargs)

diff --git a/benchmark/onnx_benchmark_helper.py b/benchmark/onnx_benchmark_helper.py
@@ -153,6 +153,7 @@ def _impl_(model_name: str,
 
             if enable_latency_plot:
                 import time
+                import torch
                 print(f"dump results to onnxrt_latency_{num_threads}.txt")
                 result_list = []
                 with open(f"onnxrt_latency_{num_threads}.txt", "w") as of:

diff --git a/benchmark/run_gpu_variable_benchmark.sh b/benchmark/run_gpu_variable_benchmark.sh
@@ -17,7 +17,7 @@ set -e
 # FRAMEWORKS=("turbo-transformers" "torch" "onnxruntime")
 FRAMEWORKS=("turbo-transformers" "torch")
 # Note Onnx doese not supports Albert
-# FRAMEWORKS=("onnxruntime")
+# FRAMEWORKS=("onnxruntime-gpu")
 
 MAX_SEQ_LEN=(500)