wehub-resource-sync
97 lines
4.0 KiB
Python
Executable File
97 lines
4.0 KiB
Python
Executable File
import argparse
|
|
parser = argparse.ArgumentParser(description='bench mnn/tensorflow/torch on pc')
|
|
parser.add_argument('-f', '--framework', choices=['mnn', 'tf', 'torch'], help='test framework', required=True)
|
|
parser.add_argument('--modeldir', help='test model directory', required=True)
|
|
parser.add_argument('--thread-num', choices=range(1, 5), default=1, help='model dir')
|
|
parser.add_argument('--loop-num', default=10, help='run loop number')
|
|
parser.add_argument('--backend', choices=['cpu', 'cuda'])
|
|
args = parser.parse_args()
|
|
|
|
import os
|
|
from os.path import join, exists, abspath
|
|
import json
|
|
import time
|
|
import numpy as np
|
|
|
|
def bench_mnn(config):
|
|
def add_suffix(exe):
|
|
import platform
|
|
if platform.system() == 'Linux':
|
|
return exe
|
|
if platform.system() == 'Windows':
|
|
return exe + '.exe'
|
|
else:
|
|
return exe + '.out'
|
|
mnn_json = {
|
|
'outputs': config['output_layers'],
|
|
'inputs': [{'name': name, 'shape': shape, 'value': 0} for name, shape in zip(config['input_layers'], config['input_shapes'])]
|
|
}
|
|
import tempfile
|
|
from subprocess import Popen, PIPE, STDOUT
|
|
with tempfile.TemporaryDirectory() as dirname:
|
|
with open(join(dirname, 'input.json'), 'w') as f:
|
|
json.dump(mnn_json, f, indent=4)
|
|
run_exe = abspath(join('MNN', 'build', add_suffix('ModuleBasic')))
|
|
mnn_path = abspath(join(args.modeldir, 'mnn', f"{config['model']}.mnn"))
|
|
backend = 0 if args.backend == 'cpu' else 2
|
|
with open(join('result', f"mnn_pc_{args.backend}.txt"), 'a+') as logfile:
|
|
p = Popen([run_exe, mnn_path, dirname, '0', str(backend), str(args.loop_num), str(args.thread_num)], stdout=logfile, stderr=STDOUT, text=True)
|
|
p.wait()
|
|
logfile.flush()
|
|
|
|
def bench_tf(config):
|
|
import tensorflow as tf
|
|
tf.config.threading.set_inter_op_parallelism_threads(args.thread_num)
|
|
model_path = join(args.modeldir, 'pb', f"{config['model']}.pb")
|
|
model = tf.saved_model.load(model_path)
|
|
dtype_map = {'float': tf.float32, 'int': tf.int32}
|
|
input_dict = {name: tf.zeros(shape, dtype=dtype_map[dtype]) for name, shape, dtype in zip(config['input_layers'], config['input_shapes'], config['input_dtypes'])}
|
|
infer_func = model.signatures["serving_default"]
|
|
outputs = infer_func(**input_dict)
|
|
times = []
|
|
for i in range(args.loop_num):
|
|
start_t = time.time()
|
|
outputs = infer_func(**input_dict)
|
|
times.append((time.time() - start_t) * 1000)
|
|
with open(join('result', f"tf_pc_{args.backend}.txt"), 'a+') as logfile:
|
|
f.writelines([
|
|
f"model: {model_path}, backend: {args.backend}, loop_num: {args.loop_num}",
|
|
f"max: {max(times)}, min: {min(times)}, avg: {sum(times) / args.loop_num}"
|
|
])
|
|
|
|
def bench_torch(config):
|
|
import torch
|
|
model = torch.jit.load(filename, torch.device(args.backend))
|
|
model.eval()
|
|
torch.set_num_threads(args.thread_num)
|
|
dtype_map = {'float': 'torch.FloatTensor', 'int': 'torch.IntTensor'}
|
|
input_list = [torch.rand(shape).type(dtype_map[dtype]) for shape, dtype in zip(config['input_shapes'], config['input_dtypes'])]
|
|
model = ipex.optimize(model, dtype=torch.float32)
|
|
for i in range(10):
|
|
outputs = model.forward(*input_list)
|
|
times = []
|
|
for i in range(args.loop_num):
|
|
start_t = time.time()
|
|
outputs = model.forward(*input_list)
|
|
times.append((time.time() - start_t) * 1000)
|
|
with open(join('result', f"torch_pc_{args.backend}.txt"), 'a+') as logfile:
|
|
f.writelines([
|
|
f"model: {model_path}, backend: {args.backend}, loop_num: {args.loop_num}",
|
|
f"max: {max(times)}, min: {min(times)}, avg: {sum(times) / args.loop_num}"
|
|
])
|
|
|
|
def main():
|
|
with open(join(args.modeldir, 'config.json')) as f:
|
|
configs = json.load(f)
|
|
for config in configs:
|
|
if args.framework == 'mnn':
|
|
bench_mnn(config)
|
|
elif args.framework == 'tf':
|
|
bench_tf(config)
|
|
else:
|
|
bench_torch(config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|