# -*- coding: UTF-8 -*-
import os
import sys
import subprocess
import argparse
import re
import configparser
import logging
import copy
import psutil

if sys.version_info[0] < 3:
    import struct

import tensorflow as tf
from tensorflow.core.framework import graph_pb2

# level=logging.INFO only logging.debug info,level=logging.DEBUG may logging.debug debug and info
logging.basicConfig(level=logging.INFO)


# Global configuration mainly got from test.cfg
TRANSFORMER_PATH = ''
TF_SRC_PATH = ''
TF_SLIM_PATH= '' 

NUM_THREADS = 1
EPSILON = 0.0
LOOPS = 1

testcases = []

class TestCase(object):
    def __init__(self):
        self.model_name = ''
        self.model_type = ''
        self.url = ''
        self.output_node = ''

        # Some checkpoint cannot be frozen without fix
        self.fix_graph = False

        # File path of checkpoint (.ckpt)
        self.ckpt_file = None
        # File path of inference graph, generated by export_inference_graph.py 
        self.graph_file = None
        # File path of frozen pb, generated by freeze_graph
        self.frozen_file = None
        # File path where to save the transformed model
        self.save_model_dir = None


    def __repr__(self):
        return '[%s]\ntype=%s\nurl=%s\noutput_node=%s\nfix_graph=%r' % \
               (self.model_name, self.model_type, self.url, self.output_node, self.fix_graph)


def exec_cmd(cmd, title, check_output=True):
    logging.info(title)
    logging.debug(cmd)
    if check_output:
        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
        return out       
    else:
        text = os.popen(cmd)
        out = text.read()
        return out


def get_extract_command(filename, target_path):
    if filename.endswith('.tar.gz'):
        extract_file = 'tar -xzf %s -C %s' % (filename, target_path)
    elif filename.endswith('.tar.bz2'):
        extract_file = 'tar -xjf %s -C %s' % (filename, target_path)
    else:
        extract_file = 'tar -xf %s -C %s' % (filename, target_path)
    return extract_file
    

def decode_string(str):
    str = str.decode('UTF-8')
    return str


# Find a ckpt file in a directory, and return the path of ckpt file
def find_ckpt(ckpt_dir):
    ckpt_file = None

    files = os.listdir(ckpt_dir)
    for f in files:
        if f.endswith('.ckpt'):
            ckpt_file = f
            break

    return os.path.join(ckpt_dir, ckpt_file) if ckpt_file else None
    

def download_ckpt(tc):
    ckpt_dir = '%s/checkpoints/%s' % (TRANSFORMER_PATH, tc.model_name)
    tc.graph_file = '%s/%s_inf_graph.pb' % (ckpt_dir, tc.model_name)
    tc.frozen_file = '%s/frozen_%s.pb' % (ckpt_dir, tc.model_name)
    tc.save_model_dir = '%s/saved_model/%s' % (TRANSFORMER_PATH, tc.model_name)

    # Look for existing ckpt file
    if os.path.exists(ckpt_dir):
        tc.ckpt_file = find_ckpt(ckpt_dir)
    else:
        os.makedirs(ckpt_dir)

    # Already exist
    if tc.ckpt_file:
        logging.debug("ckpt file (%s) already exist!" % tc.ckpt_file)
        return

    tar_name = tc.url.split('/')[-1]
    download_cmd = 'wget -c %s' % tc.url
    extract_cmd = get_extract_command(tar_name, ckpt_dir)
    rm_cmd = 'rm -f %s' % tar_name
   
    exec_cmd(download_cmd, "download ckpt file ...")
    exec_cmd(extract_cmd, "untar ckpt file ...")
    exec_cmd(rm_cmd, "delete ckpt archieve ...")
    logging.debug("ckpt file has been download!")

    tc.ckpt_file = find_ckpt(ckpt_dir)

 
def export_inference_graph(tc):
    if os.path.isfile(tc.graph_file):
        logging.debug("graph file of %s already exist!" % tc.model_name)
    else:
       export_cmd = 'cd %s && python export_inference_graph.py --alsologtostderr --model_name=%s --output_file=%s' % (TF_SLIM_PATH, tc.model_name, tc.graph_file)
       exec_cmd(export_cmd, "export graph.pb")
       logging.debug("%s has been exported!" % tc.graph_file)


def summarize_graph(tc):
    bazel_build = 'cd %s && bazel build tensorflow/tools/graph_transforms:summarize_graph' % TF_SRC_PATH
    summarize_graph = 'cd %s && bazel-bin/tensorflow/tools/graph_transforms/summarize_graph --in_graph=%s' % (TF_SRC_PATH, tc.graph_file)
    exec_cmd(bazel_build, "bazel build ...")
    exec_cmd(summarize_graph, "summarize graph ...")
    logging.debug("summarize graph has been done!")


def load_graph(filename):
    graph_def = tf.GraphDef()
    with tf.gfile.FastGFile(filename, 'rb') as f:
        graph_def.ParseFromString(f.read())
    return graph_def


def change_tensor_shape(tensor_shape):
    dims = len(tensor_shape.dim)
    if dims == 4 and tensor_shape.dim[3].size == 1001:
        tensor_shape.dim[3].size = 1000
        #print("shape changed, shape=%s" % str(tensor_shape))
    if dims == 1 and tensor_shape.dim[0].size == 1001:
        tensor_shape.dim[0].size = 1000
        #print("shape changed, shape=%s" % str(tensor_shape))
    

def int_to_bytes(val):
    if sys.version_info[0] >= 3:
        return val.to_bytes(4, 'little')
    else:
        return struct.pack("<L", val)


def int_from_bytes(barray):
    if sys.version_info[0] >= 3:
        return int.from_bytes(barray, byteorder='little')
    else:
        return struct.unpack("<L", barray)[0]

# Designed to fix the error like:
# Assign requires shapes of both tensors to match. lhs shape= [1,1,4096,1001] rhs shape= [1,1,4096,1000]
def fix_graph_1001_to_1000(tc):
    pb_file = tc.graph_file

    graph_def = load_graph(pb_file)
    
    new_graph_def = graph_pb2.GraphDef()
    
    for node in graph_def.node:
        # Check the value of const node
        if node.op == 'Const':
            tensor = node.attr.get('value').tensor
            # DataType value got from https://github.com/tensorflow/tensorflow/blob/r1.10/tensorflow/core/framework/types.proto
            if tensor and tensor.dtype >= 3 and tensor.dtype <= 6: # data type is int
                tensor_shape = tensor.tensor_shape
                # Change tensor like: {"tensor":{"dtype":"DT_INT32","tensor_shape":{"dim":[{"size":4}]},"tensor_content":"\\001\\000\\000\\000\\001\\000\\000\\000\\000\\020\\000\\000\\350\\003\\000\\000"}}
                if len(tensor_shape.dim) == 1 and tensor_shape.dim[0].size == 4:
                    element_size = (int)(len(tensor.tensor_content) / 4)
                    shape = [0, 0, 0, 0]
                    for i in range(4):
                        shape[i] = int_from_bytes(tensor.tensor_content[i*element_size: (i+1)*element_size])
                    # 1x1x2048x1001 -> 1x1x2048x1000, etc.
                    if shape[3] == 1001:
                        shape[3] = 1000
                        content = int_to_bytes(shape[0]) + int_to_bytes(shape[1]) + int_to_bytes(shape[2]) + int_to_bytes(shape[3]);
                        tensor.tensor_content = content
                # Change tensor like: {"tensor":{"dtype":"DT_INT32","tensor_shape":{"dim":[{"size":2}]},"tensor_content":"\\377\\377\\377\\377\\351\\003\\000\\000"}}
                if len(tensor_shape.dim) == 1 and tensor_shape.dim[0].size == 2:
                    element_size = (int)(len(tensor.tensor_content) / 2)
                    shape = [0, 0]
                    for i in range(2):
                        shape[i] = int_from_bytes(tensor.tensor_content[i*element_size: (i+1)*element_size])
                    # -1x1001 -> -1x1000, etc.
                    if shape[1] == 1001:
                        shape[1] = 1000
                        content = int_to_bytes(shape[0]) + int_to_bytes(shape[1]);
                        tensor.tensor_content = content
                # Change tensor like: {"tensor":{"dtype":"DT_INT32","tensor_shape":{"dim":[{"size":1}]},"int_val":1000}}
                if len(tensor_shape.dim) == 1 and tensor_shape.dim[0].size == 1 and len(tensor.int_val) == 1:
                    if tensor.int_val[0] == 1001:
                        tensor.int_val[0] = 1000
    
        # Check shape attribute
        shape_value = node.attr.get('shape')
        if shape_value:
            change_tensor_shape(shape_value.shape)
    
        new_graph_def.node.extend([copy.deepcopy(node)])
    
    # save new graph
    with tf.gfile.GFile(pb_file, "wb") as f:
        f.write(new_graph_def.SerializeToString())


def frozen_pb(tc):
    if os.path.exists(tc.frozen_file):
         logging.debug("frozen pb file exist")
    else:
        bazel_build = 'cd %s && bazel build tensorflow/python/tools:freeze_graph' % TF_SRC_PATH
        frozen_cmd = 'cd %s && bazel-bin/tensorflow/python/tools/freeze_graph --input_graph=%s --input_checkpoint=%s --input_binary=true --output_graph=%s --output_node_names=%s' % \
                     (TF_SRC_PATH, tc.graph_file, tc.ckpt_file, tc.frozen_file, tc.output_node)
        exec_cmd(bazel_build, "bazel build ...")
        exec_cmd(frozen_cmd, "frozen pb ...")
        logging.debug("frozen has been done!")



def tf_inference(tc, inference_input):
    if inference_input == 'data':
        do_inference = 'cd %s/tests && OMP_NUM_THREADS=%d python pb_inference.py --pb_file=%s --output_node=%s --batch_size=1 --loop=%d' % \
                       (TRANSFORMER_PATH, NUM_THREADS, tc.frozen_file, tc.output_node, LOOPS)
    else:
        do_inference = 'cd %s/tests && OMP_NUM_THREADS=%d python pb_inference.py --pb_file=%s --output_node=% --batch_size=1 --loop=%d --picture=%s' % \
                       (TRANSFORMER_PATH, NUM_THREADS, tc.frozen_file, tc.output_node, LOOPS, inference_input)
 
    output = exec_cmd(do_inference,"tensorflow do inference!")
    output = decode_string(output)
    
    # Parse output info from cmd out
    tf_info = "tensorflow output:\s\[([\-?\d+\.?\d*e?-?\d*?\s]+)"
    tf_result = re.findall(tf_info, output)
    logging.debug("tensorflow output: %s" % tf_result)
 
    tf_time_info = "TF time used per loop is: (\d+\.?\d*) ms"
    tf_time_used = re.findall(tf_time_info, output)
    
    return tf_result, tf_time_used



def mkldnn_inference(tc, inference_input):
    if not os.path.exists(tc.save_model_dir):   
         os.makedirs(tc.save_model_dir)
         
    tf2topo = 'cd %s/ && python tf2topo.py --input_model_filename=%s --weights_file=%s/weights.bin --pkl_file=%s/weights.pkl --topo_file=%s/topo.txt' % \
              (TRANSFORMER_PATH, tc.frozen_file, tc.save_model_dir, tc.save_model_dir, tc.save_model_dir)
    exec_cmd(tf2topo, "convert tf pb file to topo")

    topo2mkldnn = 'cd %s/ && python topo2code.py --topo=%s/topo.txt' % (TRANSFORMER_PATH, tc.save_model_dir)
    exec_cmd(topo2mkldnn, "convert topo to inference code")

    run_mkldnn = 'cd %s/inference_code/ && sh build.sh && OMP_NUM_THREADS=%d ./test -W %s/weights.bin -b 1 -l %d' % \
                 (TRANSFORMER_PATH, NUM_THREADS, tc.save_model_dir, LOOPS)
    output = exec_cmd(run_mkldnn, "build and run inference code")
    output = decode_string(output)

    #search info from cmd ouput
    out_info = "Last_output >> \[([\-?\d+\.?\d*e?-?\d*?\s]+)"
    inference_result = re.findall(out_info, output)
    logging.debug("mkldnn output:%s" % inference_result)

    mkldnn_time_info = "AVG Time: (\d+\.?\d*) ms"
    mkldnn_time_used = re.findall(mkldnn_time_info, output)
    
    return inference_result,mkldnn_time_used 


def str2list(input_str):
    output_list = input_str.split()
    
    return output_list


def compare(list1,list2):
    length = min(len(list1),len(list2))
  
    num = 0    
    for i in range(length):
       if abs(float(list1[i]) - float(list2[i])) <= float(EPSILON):
           num = num + 1

    if num == length:
        logging.debug("mkldnn inference outputs equal tensorflow outputs, num=%d" % num)
        return True
    else:
        logging.debug("mkldnn inference outputs are different from tensorflow outputs!")
        return False



def init_config():
    global TRANSFORMER_PATH, TF_SRC_PATH, TF_SLIM_PATH, NUM_THREADS, EPSILON, LOOPS

    TRANSFORMER_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
    NUM_THREADS = psutil.cpu_count(logical = False)

    config = configparser.ConfigParser()
    config.read('test.cfg')

    TF_SRC_PATH = config.get('path', 'tensorflow')
    TF_SLIM_PATH = config.get('path', 'tensorflow_slim')
    EPSILON = config.get('control', 'epsilon')
    LOOPS = int(config.get('control', 'loops'))

    str_models = config.get('models', 'names')
    model_list = str_models.split(',')

    for model in model_list:
        tc = TestCase()
        tc.model_name = model
        tc.model_type = config.get(model, 'type')
        tc.url = config.get(model, 'url')
        tc.output_node = config.get(model, 'output_node')
        try:
            tc.fix_graph = config.get(model, 'fix_graph')
        except:
            pass
        testcases.append(tc)



def model_test(tc):        
    print(" Test model: %s start ............" % tc.model_name)
    
    if tc.model_type == "ckpt":
        download_ckpt(tc)
        export_inference_graph(tc)
        if tc.fix_graph:
            fix_graph_1001_to_1000(tc)
        # if need to get the output name, could call this func
        summarize_graph(tc)
        frozen_pb(tc)
    else:
        logger.debug("model type error!")
        exit()

    tf_output, tf_time = tf_inference(tc, args.inference_input)
    mkldnn_output, mkldnn_time = mkldnn_inference(tc, args.inference_input)
    
    tf_output_list = str2list(tf_output[0])
    mkldnn_output_list = str2list(mkldnn_output[0])
   
    logging.debug("tensorflow ouput: %s" % tf_output_list)
    logging.debug("mkldnn output: %s" % mkldnn_output_list)

    result = compare(mkldnn_output_list, tf_output_list) 
    
    if result:
        print("             %s passed! tensorflow used time: %s ms, mkldnn used time: %s ms." % (tc.model_name, tf_time, mkldnn_time)) 
    else:
        print("             %s failed! tensorflow used time: %s ms, mkldnn used time: %s ms." % (tc.model_name, tf_time, mkldnn_time))



if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", "-n", default="all", type=str, help="which model to test, default means to test all configured models in test.cfg.")
    parser.add_argument("--inference_input", "-i", default="data", type=str, help="input: 'data' or an image file path. Currently only support 'data', which means to use emulated data.")
    
    args = parser.parse_args() 

    init_config()

    if args.model_name == "all":  
        for tc in testcases:
            model_test(tc)
    else:
        for tc in testcases:
            if tc.model_name == args.model_name:
                model_test(tc)
    
    print(" All tests done!")

# -*- coding: UTF-8 -*-
import os
import sys
import subprocess
import argparse
import re
import configparser
import logging
import copy
import psutil

if sys.version_info[0] < 3:
    import struct

import tensorflow as tf
from tensorflow.core.framework import graph_pb2

# level=logging.INFO only logging.debug info,level=logging.DEBUG may logging.debug debug and info
logging.basicConfig(level=logging.INFO)


# Global configuration mainly got from test.cfg
TRANSFORMER_PATH = ''
TF_SRC_PATH = ''
TF_SLIM_PATH= '' 

NUM_THREADS = 1
EPSILON = 0.0
LOOPS = 1

testcases = []

class TestCase(object):
    def __init__(self):
        self.model_name = ''
        self.model_type = ''
        self.url = ''
        self.output_node = ''

        # Some checkpoint cannot be frozen without fix
        self.fix_graph = False

        # File path of checkpoint (.ckpt)
        self.ckpt_file = None
        # File path of inference graph, generated by export_inference_graph.py 
        self.graph_file = None
        # File path of frozen pb, generated by freeze_graph
        self.frozen_file = None
        # File path where to save the transformed model
        self.save_model_dir = None


    def __repr__(self):
        return '[%s]\ntype=%s\nurl=%s\noutput_node=%s\nfix_graph=%r' % \
               (self.model_name, self.model_type, self.url, self.output_node, self.fix_graph)


def exec_cmd(cmd, title, check_output=True):
    logging.info(title)
    logging.debug(cmd)
    if check_output:
        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
        return out       
    else:
        text = os.popen(cmd)
        out = text.read()
        return out


def get_extract_command(filename, target_path):
    if filename.endswith('.tar.gz'):
        extract_file = 'tar -xzf %s -C %s' % (filename, target_path)
    elif filename.endswith('.tar.bz2'):
        extract_file = 'tar -xjf %s -C %s' % (filename, target_path)
    else:
        extract_file = 'tar -xf %s -C %s' % (filename, target_path)
    return extract_file
    

def decode_string(str):
    str = str.decode('UTF-8')
    return str


# Find a ckpt file in a directory, and return the path of ckpt file
def find_ckpt(ckpt_dir):
    ckpt_file = None

    files = os.listdir(ckpt_dir)
    for f in files:
        if f.endswith('.ckpt'):
            ckpt_file = f
            break

    return os.path.join(ckpt_dir, ckpt_file) if ckpt_file else None
    

def download_ckpt(tc):
    ckpt_dir = '%s/checkpoints/%s' % (TRANSFORMER_PATH, tc.model_name)
    tc.graph_file = '%s/%s_inf_graph.pb' % (ckpt_dir, tc.model_name)
    tc.frozen_file = '%s/frozen_%s.pb' % (ckpt_dir, tc.model_name)
    tc.save_model_dir = '%s/saved_model/%s' % (TRANSFORMER_PATH, tc.model_name)

    # Look for existing ckpt file
    if os.path.exists(ckpt_dir):
        tc.ckpt_file = find_ckpt(ckpt_dir)
    else:
        os.makedirs(ckpt_dir)

    # Already exist
    if tc.ckpt_file:
        logging.debug("ckpt file (%s) already exist!" % tc.ckpt_file)
        return

    tar_name = tc.url.split('/')[-1]
    download_cmd = 'wget -c %s' % tc.url
    extract_cmd = get_extract_command(tar_name, ckpt_dir)
    rm_cmd = 'rm -f %s' % tar_name
   
    exec_cmd(download_cmd, "download ckpt file ...")
    exec_cmd(extract_cmd, "untar ckpt file ...")
    exec_cmd(rm_cmd, "delete ckpt archieve ...")
    logging.debug("ckpt file has been download!")

    tc.ckpt_file = find_ckpt(ckpt_dir)

 
def export_inference_graph(tc):
    if os.path.isfile(tc.graph_file):
        logging.debug("graph file of %s already exist!" % tc.model_name)
    else:
       export_cmd = 'cd %s && python export_inference_graph.py --alsologtostderr --model_name=%s --output_file=%s' % (TF_SLIM_PATH, tc.model_name, tc.graph_file)
       exec_cmd(export_cmd, "export graph.pb")
       logging.debug("%s has been exported!" % tc.graph_file)


def summarize_graph(tc):
    bazel_build = 'cd %s && bazel build tensorflow/tools/graph_transforms:summarize_graph' % TF_SRC_PATH
    summarize_graph = 'cd %s && bazel-bin/tensorflow/tools/graph_transforms/summarize_graph --in_graph=%s' % (TF_SRC_PATH, tc.graph_file)
    exec_cmd(bazel_build, "bazel build ...")
    exec_cmd(summarize_graph, "summarize graph ...")
    logging.debug("summarize graph has been done!")


def load_graph(filename):
    graph_def = tf.GraphDef()
    with tf.gfile.FastGFile(filename, 'rb') as f:
        graph_def.ParseFromString(f.read())
    return graph_def


def change_tensor_shape(tensor_shape):
    dims = len(tensor_shape.dim)
    if dims == 4 and tensor_shape.dim[3].size == 1001:
        tensor_shape.dim[3].size = 1000
        #print("shape changed, shape=%s" % str(tensor_shape))
    if dims == 1 and tensor_shape.dim[0].size == 1001:
        tensor_shape.dim[0].size = 1000
        #print("shape changed, shape=%s" % str(tensor_shape))
    

def int_to_bytes(val):
    if sys.version_info[0] >= 3:
        return val.to_bytes(4, 'little')
    else:
        return struct.pack("<L", val)


def int_from_bytes(barray):
    if sys.version_info[0] >= 3:
        return int.from_bytes(barray, byteorder='little')
    else:
        return struct.unpack("<L", barray)[0]

# Designed to fix the error like:
# Assign requires shapes of both tensors to match. lhs shape= [1,1,4096,1001] rhs shape= [1,1,4096,1000]
def fix_graph_1001_to_1000(tc):
    pb_file = tc.graph_file

    graph_def = load_graph(pb_file)
    
    new_graph_def = graph_pb2.GraphDef()
    
    for node in graph_def.node:
        # Check the value of const node
        if node.op == 'Const':
            tensor = node.attr.get('value').tensor
            # DataType value got from https://github.com/tensorflow/tensorflow/blob/r1.10/tensorflow/core/framework/types.proto
            if tensor and tensor.dtype >= 3 and tensor.dtype <= 6: # data type is int
                tensor_shape = tensor.tensor_shape
                # Change tensor like: {"tensor":{"dtype":"DT_INT32","tensor_shape":{"dim":[{"size":4}]},"tensor_content":"\\001\\000\\000\\000\\001\\000\\000\\000\\000\\020\\000\\000\\350\\003\\000\\000"}}
                if len(tensor_shape.dim) == 1 and tensor_shape.dim[0].size == 4:
                    element_size = (int)(len(tensor.tensor_content) / 4)
                    shape = [0, 0, 0, 0]
                    for i in range(4):
                        shape[i] = int_from_bytes(tensor.tensor_content[i*element_size: (i+1)*element_size])
                    # 1x1x2048x1001 -> 1x1x2048x1000, etc.
                    if shape[3] == 1001:
                        shape[3] = 1000
                        content = int_to_bytes(shape[0]) + int_to_bytes(shape[1]) + int_to_bytes(shape[2]) + int_to_bytes(shape[3]);
                        tensor.tensor_content = content
                # Change tensor like: {"tensor":{"dtype":"DT_INT32","tensor_shape":{"dim":[{"size":2}]},"tensor_content":"\\377\\377\\377\\377\\351\\003\\000\\000"}}
                if len(tensor_shape.dim) == 1 and tensor_shape.dim[0].size == 2:
                    element_size = (int)(len(tensor.tensor_content) / 2)
                    shape = [0, 0]
                    for i in range(2):
                        shape[i] = int_from_bytes(tensor.tensor_content[i*element_size: (i+1)*element_size])
                    # -1x1001 -> -1x1000, etc.
                    if shape[1] == 1001:
                        shape[1] = 1000
                        content = int_to_bytes(shape[0]) + int_to_bytes(shape[1]);
                        tensor.tensor_content = content
                # Change tensor like: {"tensor":{"dtype":"DT_INT32","tensor_shape":{"dim":[{"size":1}]},"int_val":1000}}
                if len(tensor_shape.dim) == 1 and tensor_shape.dim[0].size == 1 and len(tensor.int_val) == 1:
                    if tensor.int_val[0] == 1001:
                        tensor.int_val[0] = 1000
    
        # Check shape attribute
        shape_value = node.attr.get('shape')
        if shape_value:
            change_tensor_shape(shape_value.shape)
    
        new_graph_def.node.extend([copy.deepcopy(node)])
    
    # save new graph
    with tf.gfile.GFile(pb_file, "wb") as f:
        f.write(new_graph_def.SerializeToString())


def frozen_pb(tc):
    if os.path.exists(tc.frozen_file):
         logging.debug("frozen pb file exist")
    else:
        bazel_build = 'cd %s && bazel build tensorflow/python/tools:freeze_graph' % TF_SRC_PATH
        frozen_cmd = 'cd %s && bazel-bin/tensorflow/python/tools/freeze_graph --input_graph=%s --input_checkpoint=%s --input_binary=true --output_graph=%s --output_node_names=%s' % \
                     (TF_SRC_PATH, tc.graph_file, tc.ckpt_file, tc.frozen_file, tc.output_node)
        exec_cmd(bazel_build, "bazel build ...")
        exec_cmd(frozen_cmd, "frozen pb ...")
        logging.debug("frozen has been done!")



def tf_inference(tc, inference_input):
    if inference_input == 'data':
        do_inference = 'cd %s/tests && OMP_NUM_THREADS=%d python pb_inference.py --pb_file=%s --output_node=%s --batch_size=1 --loop=%d' % \
                       (TRANSFORMER_PATH, NUM_THREADS, tc.frozen_file, tc.output_node, LOOPS)
    else:
        do_inference = 'cd %s/tests && OMP_NUM_THREADS=%d python pb_inference.py --pb_file=%s --output_node=% --batch_size=1 --loop=%d --picture=%s' % \
                       (TRANSFORMER_PATH, NUM_THREADS, tc.frozen_file, tc.output_node, LOOPS, inference_input)
 
    output = exec_cmd(do_inference,"tensorflow do inference!")
    output = decode_string(output)
    
    # Parse output info from cmd out
    tf_info = "tensorflow output:\s\[([\-?\d+\.?\d*e?-?\d*?\s]+)"
    tf_result = re.findall(tf_info, output)
    logging.debug("tensorflow output: %s" % tf_result)
 
    tf_time_info = "TF time used per loop is: (\d+\.?\d*) ms"
    tf_time_used = re.findall(tf_time_info, output)
    
    return tf_result, tf_time_used



def mkldnn_inference(tc, inference_input):
    if not os.path.exists(tc.save_model_dir):   
         os.makedirs(tc.save_model_dir)
         
    tf2topo = 'cd %s/ && python tf2topo.py --input_model_filename=%s --weights_file=%s/weights.bin --pkl_file=%s/weights.pkl --topo_file=%s/topo.txt' % \
              (TRANSFORMER_PATH, tc.frozen_file, tc.save_model_dir, tc.save_model_dir, tc.save_model_dir)
    exec_cmd(tf2topo, "convert tf pb file to topo")

    topo2mkldnn = 'cd %s/ && python topo2code.py --topo=%s/topo.txt' % (TRANSFORMER_PATH, tc.save_model_dir)
    exec_cmd(topo2mkldnn, "convert topo to inference code")

    run_mkldnn = 'cd %s/inference_code/ && sh build.sh && OMP_NUM_THREADS=%d ./test -W %s/weights.bin -b 1 -l %d' % \
                 (TRANSFORMER_PATH, NUM_THREADS, tc.save_model_dir, LOOPS)
    output = exec_cmd(run_mkldnn, "build and run inference code")
    output = decode_string(output)

    #search info from cmd ouput
    out_info = "Last_output >> \[([\-?\d+\.?\d*e?-?\d*?\s]+)"
    inference_result = re.findall(out_info, output)
    logging.debug("mkldnn output:%s" % inference_result)

    mkldnn_time_info = "AVG Time: (\d+\.?\d*) ms"
    mkldnn_time_used = re.findall(mkldnn_time_info, output)
    
    return inference_result,mkldnn_time_used 


def str2list(input_str):
    output_list = input_str.split()
    
    return output_list


def compare(list1,list2):
    length = min(len(list1),len(list2))
  
    num = 0    
    for i in range(length):
       if abs(float(list1[i]) - float(list2[i])) <= float(EPSILON):
           num = num + 1

    if num == length:
        logging.debug("mkldnn inference outputs equal tensorflow outputs, num=%d" % num)
        return True
    else:
        logging.debug("mkldnn inference outputs are different from tensorflow outputs!")
        return False



def init_config():
    global TRANSFORMER_PATH, TF_SRC_PATH, TF_SLIM_PATH, NUM_THREADS, EPSILON, LOOPS

    TRANSFORMER_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
    NUM_THREADS = psutil.cpu_count(logical = False)

    config = configparser.ConfigParser()
    config.read('test.cfg')

    TF_SRC_PATH = config.get('path', 'tensorflow')
    TF_SLIM_PATH = config.get('path', 'tensorflow_slim')
    EPSILON = config.get('control', 'epsilon')
    LOOPS = int(config.get('control', 'loops'))

    str_models = config.get('models', 'names')
    model_list = str_models.split(',')

    for model in model_list:
        tc = TestCase()
        tc.model_name = model
        tc.model_type = config.get(model, 'type')
        tc.url = config.get(model, 'url')
        tc.output_node = config.get(model, 'output_node')
        try:
            tc.fix_graph = config.get(model, 'fix_graph')
        except:
            pass
        testcases.append(tc)



def model_test(tc):        
    print(" Test model: %s start ............" % tc.model_name)
    
    if tc.model_type == "ckpt":
        download_ckpt(tc)
        export_inference_graph(tc)
        if tc.fix_graph:
            fix_graph_1001_to_1000(tc)
        # if need to get the output name, could call this func
        summarize_graph(tc)
        frozen_pb(tc)
    else:
        logger.debug("model type error!")
        exit()

    tf_output, tf_time = tf_inference(tc, args.inference_input)
    mkldnn_output, mkldnn_time = mkldnn_inference(tc, args.inference_input)
    
    tf_output_list = str2list(tf_output[0])
    mkldnn_output_list = str2list(mkldnn_output[0])
   
    logging.debug("tensorflow ouput: %s" % tf_output_list)
    logging.debug("mkldnn output: %s" % mkldnn_output_list)

    result = compare(mkldnn_output_list, tf_output_list) 
    
    if result:
        print("             %s passed! tensorflow used time: %s ms, mkldnn used time: %s ms." % (tc.model_name, tf_time, mkldnn_time)) 
    else:
        print("             %s failed! tensorflow used time: %s ms, mkldnn used time: %s ms." % (tc.model_name, tf_time, mkldnn_time))



if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", "-n", default="all", type=str, help="which model to test, default means to test all configured models in test.cfg.")
    parser.add_argument("--inference_input", "-i", default="data", type=str, help="input: 'data' or an image file path. Currently only support 'data', which means to use emulated data.")
    
    args = parser.parse_args() 

    init_config()

    if args.model_name == "all":  
        for tc in testcases:
            model_test(tc)
    else:
        for tc in testcases:
            if tc.model_name == args.model_name:
                model_test(tc)
    
    print(" All tests done!")

# -*- coding: UTF-8 -*- # level=logging.INFO only logging.debug info,level=logging.DEBUG may logging.debug debug and info # Global configuration mainly got from test.cfg # Some checkpoint cannot be frozen without fix # File path of checkpoint (.ckpt) # File path of inference graph, generated by export_inference_graph.py # File path of frozen pb, generated by freeze_graph # File path where to save the transformed model # Find a ckpt file in a directory, and return the path of ckpt file # Look for existing ckpt file # Already exist #print("shape changed, shape=%s" % str(tensor_shape)) #print("shape changed, shape=%s" % str(tensor_shape)) # Designed to fix the error like: # Assign requires shapes of both tensors to match. lhs shape= [1,1,4096,1001] rhs shape= [1,1,4096,1000] # Check the value of const node # DataType value got from https://github.com/tensorflow/tensorflow/blob/r1.10/tensorflow/core/framework/types.proto # data type is int # Change tensor like: {"tensor":{"dtype":"DT_INT32","tensor_shape":{"dim":[{"size":4}]},"tensor_content":"\\001\\000\\000\\000\\001\\000\\000\\000\\000\\020\\000\\000\\350\\003\\000\\000"}} # 1x1x2048x1001 -> 1x1x2048x1000, etc. # Change tensor like: {"tensor":{"dtype":"DT_INT32","tensor_shape":{"dim":[{"size":2}]},"tensor_content":"\\377\\377\\377\\377\\351\\003\\000\\000"}} # -1x1001 -> -1x1000, etc. # Change tensor like: {"tensor":{"dtype":"DT_INT32","tensor_shape":{"dim":[{"size":1}]},"int_val":1000}} # Check shape attribute # save new graph # Parse output info from cmd out #search info from cmd ouput # if need to get the output name, could call this func