mace_tools.py

#!/usr/bin/env python

# Must run at root dir of libmace project.
# python tools/mace_tools.py \
#     --config=tools/example.yaml \
#     --round=100 \
#     --mode=all

import argparse
import hashlib
import os
import sh
import shutil
import subprocess
import sys
import urllib
import yaml
import re

import sh_commands

from ConfigParser import ConfigParser


def run_command(command):
  print("Run command: {}".format(command))
  result = subprocess.Popen(
      command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  out, err = result.communicate()

  if out:
    print("Stdout msg:\n{}".format(out))
  if err:
    print("Stderr msg:\n{}".format(err))

  if result.returncode != 0:
    raise Exception("Exit not 0 from bash with code: {}, command: {}".format(
        result.returncode, command))


def get_global_runtime(configs):
  runtime_list = []
  for model_name in configs["models"]:
    model_runtime = configs["models"][model_name]["runtime"]
    runtime_list.append(model_runtime.lower())

  global_runtime = ""
  if "dsp" in runtime_list:
    global_runtime = "dsp"
  elif "gpu" in runtime_list:
    global_runtime = "gpu"
  elif "cpu" in runtime_list:
    global_runtime = "cpu"
  elif "neon" in runtime_list:
    global_runtime = "neon"
  else:
    raise Exception("Not found available RUNTIME in config files!")

  return global_runtime


def generate_version_code():
  command = "bash tools/generate_version_code.sh"
  run_command(command)

def generate_opencl_source_code():
  command = "bash tools/generate_opencl_code.sh source"
  run_command(command)

def generate_opencl_binay_code(target_soc, model_output_dirs, pull_or_not):
  cl_bin_dirs = []
  for d in model_output_dirs:
    cl_bin_dirs.append(os.path.join(d, "opencl_bin"))
  cl_bin_dirs_str = ",".join(cl_bin_dirs)
  if not cl_bin_dirs:
    command = "bash tools/generate_opencl_code.sh binary"
  else:
    command = "bash tools/generate_opencl_code.sh {} {} {} {}".format(
      'binary', target_soc, cl_bin_dirs_str, int(pull_or_not))
  run_command(command)

def generate_tuning_param_code(target_soc, model_output_dirs, pull_or_not):
  cl_bin_dirs = []
  for d in model_output_dirs:
    cl_bin_dirs.append(os.path.join(d, "opencl_bin"))
  cl_bin_dirs_str = ",".join(cl_bin_dirs)
  if not cl_bin_dirs:
    command = "bash tools/generate_tuning_param_code.sh"
  else:
    command = "bash tools/generate_tuning_param_code.sh {} {} {}".format(
      target_soc, cl_bin_dirs_str, int(pull_or_not))
  run_command(command)

def generate_code(target_soc, model_output_dirs, pull_or_not):
  generate_opencl_binay_code(target_soc, model_output_dirs, pull_or_not)
  generate_tuning_param_code(target_soc, model_output_dirs, pull_or_not)

def clear_env(target_soc):
  command = "bash tools/clear_env.sh {}".format(target_soc)
  run_command(command)

def input_file_name(input_name):
  return os.environ['INPUT_FILE_NAME'] + '_' + \
         re.sub('[^0-9a-zA-Z]+', '_', input_name)

def generate_random_input(target_soc, model_output_dir,
                          input_names, input_files):
  generate_data_or_not = True
  command = "bash tools/validate_tools.sh {} {} {}".format(
      target_soc, model_output_dir, int(generate_data_or_not))
  run_command(command)

  input_file_list = []
  if isinstance(input_files, list):
    input_file_list.extend(input_files)
  else:
    input_file_list.append(input_files)
  if len(input_file_list) != 0:
    input_name_list = []
    if isinstance(input_names, list):
      input_name_list.extend(input_names)
    else:
      input_name_list.append(input_names)
    if len(input_file_list) != len(input_name_list):
      raise Exception('If input_files set, the input files should match the input names.')
    for i in range(len(input_file_list)):
      if input_file_list[i] is not None:
        dst_input_file = model_output_dir + '/' + input_file_name(input_name_list[i])
        if input_file_list[i].startswith("http://") or \
            input_file_list[i].startswith("https://"):
          urllib.urlretrieve(input_file_list[i], dst_input_file)
        else:
          shutil.copy(input_file_list[i], dst_input_file)

def generate_model_code():
  command = "bash tools/generate_model_code.sh"
  run_command(command)


def build_mace_run(production_mode, model_output_dir, hexagon_mode):
  command = "bash tools/build_mace_run.sh {} {} {}".format(
      int(production_mode), model_output_dir, int(hexagon_mode))
  run_command(command)


def tuning_run(model_name,
               target_runtime,
               target_abi,
               target_soc,
               model_output_dir,
               running_round,
               tuning,
               restart_round,
               option_args=''):
  # TODO(yejianwu) refactoring the hackish code
  stdout_buff = []
  process_output = sh_commands.make_output_processor(stdout_buff)
  p = sh.bash("tools/tuning_run.sh", target_soc, model_output_dir,
              running_round, int(tuning), int(production_mode),
              restart_round, option_args, _out=process_output,
              _bg=True, _err_to_out=True)
  p.wait()
  metrics = {}
  for line in stdout_buff:
    line = line.strip()
    parts = line.split()
    if len(parts) == 6 and parts[0].startswith("time"):
      metrics["%s.create_net_ms" % model_name] = str(float(parts[1]))
      metrics["%s.mace_engine_ctor_ms" % model_name] = str(float(parts[2]))
      metrics["%s.init_ms" % model_name] = str(float(parts[3]))
      metrics["%s.warmup_ms" % model_name] = str(float(parts[4]))
      if float(parts[5]) > 0:
        metrics["%s.avg_latency_ms" % model_name] = str(float(parts[5]))
  tags = {"ro.board.platform": target_soc,
          "abi": target_abi,
          # "runtime": target_runtime, # TODO(yejianwu) Add the actual runtime
          "round": running_round, # TODO(yejianwu) change this to source/binary
          "tuning": tuning}
  sh_commands.falcon_push_metrics(metrics, endpoint="mace_model_benchmark",
                                  tags=tags)

def benchmark_model(target_soc, model_output_dir, option_args=''):
  command = "bash tools/benchmark.sh {} {} \"{}\"".format(
      target_soc, model_output_dir, option_args)
  run_command(command)


def run_model(model_name, target_runtime, target_abi, target_soc,
              model_output_dir, running_round, restart_round, option_args):
  tuning_run(model_name, target_runtime, target_abi, target_soc,
             model_output_dir, running_round, False, False,
             restart_round, option_args)


def generate_production_code(target_soc, model_output_dirs, pull_or_not):
  cl_bin_dirs = []
  for d in model_output_dirs:
    cl_bin_dirs.append(os.path.join(d, "opencl_bin"))
  cl_bin_dirs_str = ",".join(cl_bin_dirs)
  command = "bash tools/generate_production_code.sh {} {} {}".format(
      target_soc, cl_bin_dirs_str, int(pull_or_not))
  run_command(command)


def build_mace_run_prod(model_name, target_runtime, target_abi, target_soc,
                        model_output_dir, tuning):
  if "dsp" == target_runtime:
    hexagon_mode = True
  else:
    hexagon_mode = False

  generate_code(target_soc, [], False)
  production_or_not = False
  build_mace_run(production_or_not, model_output_dir, hexagon_mode)
  tuning_run(
      model_name,
      target_runtime,
      target_abi, 
      target_soc,
      model_output_dir,
      running_round=0,
      tuning=tuning,
      restart_round=1)

  generate_code(target_soc, [model_output_dir], True)
  production_or_not = True
  build_mace_run(production_or_not, model_output_dir, hexagon_mode)


def build_run_throughput_test(target_soc, run_seconds, merged_lib_file,
                              model_input_dir):
  command = "bash tools/build_run_throughput_test.sh {} {} {} {}".format(
      target_soc, run_seconds, merged_lib_file, model_input_dir)
  run_command(command)


def validate_model(target_soc, model_output_dir):
  generate_data_or_not = False
  command = "bash tools/validate_tools.sh {} {} {}".format(
      target_soc, model_output_dir, int(generate_data_or_not))
  run_command(command)


def build_production_code():
  command = "bash tools/build_production_code.sh"
  run_command(command)


def merge_libs_and_tuning_results(target_soc, output_dir, model_output_dirs):
  generate_code(target_soc, model_output_dirs, False)
  build_production_code()

  model_output_dirs_str = ",".join(model_output_dirs)
  print output_dir, model_output_dirs
  command = "bash tools/merge_libs.sh {} {} {}".format(target_soc, output_dir,
                                                       model_output_dirs_str)
  run_command(command)


def packaging_lib_file(output_dir):
  command = "bash tools/packaging_lib.sh {}".format(output_dir)
  run_command(command)

def download_model_files(model_file_path,
                         model_output_dir,
                         weight_file_path=""):
  if model_file_path.startswith("http://") or \
      model_file_path.startswith("https://"):
    os.environ["MODEL_FILE_PATH"] = model_output_dir + "/model.pb"
    urllib.urlretrieve(model_file_path, os.environ["MODEL_FILE_PATH"])

  if weight_file_path.startswith("http://") or \
      weight_file_path.startswith("https://"):
    os.environ[
      "WEIGHT_FILE_PATH"] = model_output_dir + "/model.caffemodel"
    urllib.urlretrieve(weight_file_path,
      os.environ["WEIGHT_FILE_PATH"])

def md5sum(str):
  md5 = hashlib.md5()
  md5.update(str)
  return md5.hexdigest()


def parse_model_configs():
  with open(FLAGS.config) as f:
    configs = yaml.load(f)
    return configs


def parse_args():
  """Parses command line arguments."""
  parser = argparse.ArgumentParser()
  parser.register("type", "bool", lambda v: v.lower() == "true")
  parser.add_argument(
      "--config",
      type=str,
      default="./tool/config",
      help="The global config file of models.")
  parser.add_argument(
      "--output_dir", type=str, default="build", help="The output dir.")
  parser.add_argument(
      "--round", type=int, default=1, help="The model running round.")
  parser.add_argument(
      "--run_seconds",
      type=int,
      default=10,
      help="The model throughput test running seconds.")
  parser.add_argument(
      "--restart_round", type=int, default=1, help="The model restart round.")
  parser.add_argument(
      "--tuning", type="bool", default="true", help="Tune opencl params.")
  parser.add_argument(
      "--mode",
      type=str,
      default="all",
      help="[build|run|validate|merge|all|throughput_test].")
  parser.add_argument(
      "--target_socs",
      type=str,
      default="all",
      help="SoCs to build, comma seperated list (getprop ro.board.platform)")
  return parser.parse_known_args()

def set_environment(configs):
  os.environ["EMBED_MODEL_DATA"] = str(configs["embed_model_data"])
  os.environ["VLOG_LEVEL"] = str(configs["vlog_level"])
  os.environ["PROJECT_NAME"] = os.path.splitext(os.path.basename(
    FLAGS.config))[0]
  os.environ['INPUT_FILE_NAME'] = "model_input"
  os.environ['OUTPUT_FILE_NAME'] = "model_out"

def main(unused_args):
  configs = parse_model_configs()

  if FLAGS.mode == "validate":
    FLAGS.round = 1
    FLAGS.restart_round = 1

  set_environment(configs)

  if FLAGS.mode == "build" or FLAGS.mode == "all":
    # Remove previous output dirs
    if not os.path.exists(FLAGS.output_dir):
      os.makedirs(FLAGS.output_dir)
    elif os.path.exists(os.path.join(FLAGS.output_dir, "libmace")):
      shutil.rmtree(os.path.join(FLAGS.output_dir, os.environ["PROJECT_NAME"]))
      os.makedirs(os.path.join(FLAGS.output_dir, os.environ["PROJECT_NAME"]))

    generate_version_code()
    generate_opencl_source_code()

  option_args = ' '.join([arg for arg in unused_args if arg.startswith('--')])

  available_socs = sh_commands.adb_get_all_socs()
  target_socs = available_socs
  if hasattr(configs, "target_socs"):
    target_socs = set(configs["target_socs"])
    target_socs = target_socs & available_socs

  if FLAGS.target_socs != "all":
    socs = set(FLAGS.target_socs.split(','))
    target_socs = target_socs & socs
    missing_socs = socs.difference(target_socs)
    if len(missing_socs) > 0:
      print("Error: devices with SoCs are not connected %s" % missing_socs)
      exit(1)


  for target_soc in target_socs:
    for target_abi in configs["target_abis"]:
      global_runtime = get_global_runtime(configs)
      # Transfer params by environment
      os.environ["TARGET_ABI"] = target_abi
      model_output_dirs = []
      for model_name in configs["models"]:
        print '=======================', model_name, '======================='
        # Transfer params by environment
        os.environ["MODEL_TAG"] = model_name
        model_config = configs["models"][model_name]
        input_file_list = model_config.get("validation_inputs_data", [])
        for key in model_config:
          if key in ['input_nodes', 'output_nodes'] and isinstance(
              model_config[key], list):
            os.environ[key.upper()] = ",".join(model_config[key])
          elif key in ['input_shapes', 'output_shapes'] and isinstance(
              model_config[key], list):
            os.environ[key.upper()] = ":".join(model_config[key])
          else:
            os.environ[key.upper()] = str(model_config[key])

        # Create model build directory
        model_path_digest = md5sum(model_config["model_file_path"])
        model_output_dir = "%s/%s/%s/%s/%s/%s/%s" % (FLAGS.output_dir,
                                                     os.environ["PROJECT_NAME"],
                                                     "build", model_name,
                                                     model_path_digest,
                                                     target_soc, target_abi)
        model_output_dirs.append(model_output_dir)

        if FLAGS.mode == "build" or FLAGS.mode == "all":
          if os.path.exists(model_output_dir):
            shutil.rmtree(model_output_dir)
          os.makedirs(model_output_dir)
          clear_env(target_soc)

        download_model_files(model_config["model_file_path"],
          model_output_dir, model_config.get("weight_file_path", ""))

        if FLAGS.mode == "build" or FLAGS.mode == "run" or FLAGS.mode == "validate"\
            or FLAGS.mode == "benchmark" or FLAGS.mode == "all":
          generate_random_input(target_soc, model_output_dir,
            model_config['input_nodes'], input_file_list)

        if FLAGS.mode == "build" or FLAGS.mode == "all":
          generate_model_code()
          build_mace_run_prod(model_name, global_runtime, target_abi,
                              target_soc, model_output_dir, FLAGS.tuning)

        if FLAGS.mode == "run" or FLAGS.mode == "validate" or FLAGS.mode == "all":
          run_model(model_name, global_runtime, target_abi, target_soc,
                    model_output_dir, FLAGS.round, FLAGS.restart_round,
                    option_args)

        if FLAGS.mode == "benchmark":
          benchmark_model(target_soc, model_output_dir, option_args)

        if FLAGS.mode == "validate" or FLAGS.mode == "all":
          validate_model(target_soc, model_output_dir)

      if FLAGS.mode == "build" or FLAGS.mode == "merge" or FLAGS.mode == "all":
        merge_libs_and_tuning_results(
            target_soc, FLAGS.output_dir + "/" + os.environ["PROJECT_NAME"],
            model_output_dirs)

      if FLAGS.mode == "throughput_test":
        merged_lib_file = FLAGS.output_dir + "/%s/%s/libmace_%s.%s.a" % \
            (os.environ["PROJECT_NAME"], target_abi, os.environ["PROJECT_NAME"], target_soc)
        generate_random_input(target_soc, FLAGS.output_dir, [], [])
        for model_name in configs["models"]:
          runtime = configs["models"][model_name]["runtime"]
          os.environ["%s_MODEL_TAG" % runtime.upper()] = model_name
        build_run_throughput_test(target_soc, FLAGS.run_seconds,
                                  merged_lib_file, FLAGS.output_dir)

  if FLAGS.mode == "build" or FLAGS.mode == "all":
    packaging_lib_file(FLAGS.output_dir)


if __name__ == "__main__":
  FLAGS, unparsed = parse_args()
  main(unused_args=[sys.argv[0]] + unparsed)