path: root/compiler/visq/visq

#!/usr/bin/env bash
''''export SCRIPT_PATH="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)" # '''
''''export PY_PATH=${SCRIPT_PATH}/venv/bin/python                                       # '''
''''test -f ${PY_PATH} && exec ${PY_PATH} "$0" "$@"                                     # '''
''''echo "Error: Virtual environment not found. Please run 'one-prepare-venv' command." # '''
''''exit 255                                                                            # '''

# Copyright (c) 2022 Samsung Electronics Co., Ltd. All Rights Reserved
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import subprocess
import tempfile
import json
import os
import math
import sys

import h5py as h5
import numpy as np

from shutil import copyfile
from pathlib import Path

from visqlib.Palette import YLORRD9Palette
from visqlib.QErrorComputer import MPEIRComputer, MSEComputer, TAEComputer
from visqlib.Util import valid_attr, pretty_float
from visqlib.DotBuilder import DotBuilder


def _get_parser():
    parser = argparse.ArgumentParser(
        description='Command line tool to visualize layer-wise quantization errors')
    parser.add_argument(
        "-f",
        "--fp32_circle",
        type=str,
        help="Path to the fp32 circle model.",
        required=True)
    parser.add_argument(
        "-q",
        "--q_circle",
        type=str,
        help="Path to the quantized circle model.",
        required=True)
    parser.add_argument(
        "-d",
        "--data",
        type=str,
        help=
        "Path to the data used for inference. Random data will be used if this option is not given.",
        required=False)
    parser.add_argument(
        "--mpeir_output",
        type=str,
        help="Path to the output json file (qerror metric = MPEIR).",
        required=False)
    parser.add_argument(
        "--mse_output",
        type=str,
        help="Path to the output json file (qerror metric = MSE).",
        required=False)
    parser.add_argument(
        "--tae_output",
        type=str,
        help="Path to the output json file (qerror metric = TAE).",
        required=False)
    parser.add_argument(
        "--dump_dot_graph", action="store_true", help="Dump dot graph.", required=False)
    parser.add_argument(
        "-b",
        "--batch_size",
        type=int,
        help="Batch size to process large datasets.",
        required=False)

    return parser


def _verify_args(args):
    """Verify the given arguments"""

    valid_outputs = ['mpeir_output', 'mse_output', 'tae_output']

    # Check if at least one output option is given
    num_outputs = 0
    for output_name in valid_outputs:
        if valid_attr(args, output_name):
            num_outputs += 1

    if num_outputs == 0:
        raise RuntimeError("At least one output should be given.")


def _run_dalgona(model, data, analysis, save_dir):
    dir_path = Path(__file__).parent.resolve()
    dalgona_path = os.path.join(dir_path, 'dalgona')
    cmd = [dalgona_path]
    cmd += ['--input_model', str(model)]
    cmd += ['--analysis', str(analysis)]
    if data != None:
        cmd += ['--input_data', str(data)]
    cmd += ['--analysis_args', str(save_dir)]

    try:
        subprocess.run(cmd, capture_output=True, check=True)
    except subprocess.CalledProcessError as e:
        print('Error raised while running the below command')
        print(' '.join(cmd))
        print(e.output)
        raise


# Generate h5 file that contains a dataset of a single batch
# This is for batch execution of visq
def gen_batch_h5(inputs_data, inputs_path):
    # Create h5 file
    output_path = inputs_path + "/inputs.h5"
    h5_file = h5.File(output_path, 'w')
    group = h5_file.create_group("value")
    group.attrs['desc'] = "Input data"

    for i in range(len(inputs_data)):
        sample = group.create_group(str(i))
        for j in range(len(inputs_data[i])):
            sample.create_dataset(str(j), data=inputs_data[i][j])

    h5_file.close()
    return output_path


# Aggregate intermediate results for a given data
def advance_on_data(fp32_model, fq_model, data, computers):

    curr_dir = Path(__file__).parent.resolve()
    dump_fp32_py = curr_dir / 'visqlib' / 'DumpFP32FM.py'
    dump_fq_py = curr_dir / 'visqlib' / 'DumpFakeQuantFM.py'

    with tempfile.TemporaryDirectory() as fp32_dir, \
         tempfile.TemporaryDirectory() as fq_dir:

        _run_dalgona(fp32_model, data, dump_fp32_py, fp32_dir)
        copyfile(fp32_dir + '/tensors.txt', fq_dir + '/tensors.txt')
        _run_dalgona(fq_model, data, dump_fq_py, fq_dir)

        for metric_key in computers:
            computers[metric_key][0].advance_on(fp32_dir, fq_dir)


def _run_batch(fp32_model, fq_model, data, computers, batch_size):
    with tempfile.TemporaryDirectory() as inputs_dir:
        with h5.File(data, 'r') as f:
            dataset = f['value']

            inputs = []
            for data_index in dataset:
                cur_inputs = []
                for input_index in dataset[data_index]:
                    d = dataset[data_index][input_index][:]
                    cur_inputs.append(np.array(d, np.float32))

                inputs.append(cur_inputs)
                if len(inputs) >= batch_size:
                    input_path = gen_batch_h5(inputs, inputs_dir)
                    advance_on_data(fp32_model, fq_model, input_path, computers)
                    inputs = []

            if len(inputs) > 0:
                input_path = gen_batch_h5(inputs, inputs_dir)
                advance_on_data(fp32_model, fq_model, input_path, computers)


def _fake_quantize(input_model, output_model):
    dir_path = Path(__file__).parent.resolve()
    circle_quantizer_path = os.path.join(dir_path, 'circle-quantizer')
    cmd = [circle_quantizer_path]
    cmd += ['--fake_quantize']
    cmd += [str(input_model)]
    cmd += [str(output_model)]

    try:
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        print('Error raised while running the below command')
        print(' '.join(cmd))
        print(e.output)
        raise


# Recursively visit items and check if there is Infinity or NaN
def _check_float(item):
    if isinstance(item, dict):
        for v in item.values():
            _check_float(v)
    if isinstance(item, list):
        for v in item:
            _check_float(v)
    if isinstance(item, float):
        if item == -float('inf') or item == float('inf'):
            raise RuntimeError('Infinite value detected. Value must be float')
        if math.isnan(item):
            raise RuntimeError('NaN value detected. Value must be float')


def _build_json(model, metric, colorscheme, error):
    # model: string
    # metric: string
    # colorscheme: list ['b': begin, 'e': end, 'c':color]
    # error: dict {tensor_name:error}

    meta = {}
    meta["model"] = model
    meta["metric"] = metric
    meta["colorscheme"] = pretty_float(colorscheme)
    result = {}
    result["meta"] = meta
    # Why list? To support multiple subgraphs
    result["error"] = [pretty_float(error)]

    # Invariants
    _check_float(meta["colorscheme"])
    _check_float(result["error"])
    return result


def _save_dot(circle_path: str, dot_path: str, metric: str, colors: list, qerror: dict):
    # circle_path: Path to the circle model (required to build graph)
    # dot_path: Path to the output dot file
    # metric: Metric name (ex: MPEIR, MSE)
    # colors: list [{'b': begin, 'e': end, 'c':color}, ..]
    # qerror: dict {tensor_name (str) -> qerror (float)}
    builder = DotBuilder(
        circle_path=circle_path, dot_path=dot_path, metric=metric, colors=colors)

    builder.save(qerror)


def run_on_data_batchwise(fp32_model, q_model, data, dump_dot_graph, computers,
                          batch_size):

    with tempfile.TemporaryDirectory() as model_dir:
        fq_model = model_dir + '/fq_model.circle'

        # Step 1. Fake quantize quantized circle model
        _fake_quantize(q_model, fq_model)

        # process the whole dataset batch by batch
        _run_batch(fp32_model, fq_model, data, computers, batch_size)

        #compute the final results
        for metric_key in computers:
            cur_computer = computers[metric_key][0]
            output = computers[metric_key][1]
            if metric_key == 'MPEIR':
                qerror_map = cur_computer.get_final_result()
                q_min = 0.0
                q_max = 1.0
            elif metric_key == 'MSE' or metric_key == 'TAE':
                qerror_map, q_min, q_max = cur_computer.get_final_result()

            palette = YLORRD9Palette(qerror_min=q_min, qerror_max=q_max)
            result = _build_json(
                metric=metric_key,
                model=Path(fp32_model).name,
                colorscheme=palette.colorscheme(),
                error=qerror_map)
            with open(output, "w") as f:
                json.dump(result, f)

            if dump_dot_graph:
                _save_dot(
                    circle_path=fp32_model,
                    dot_path=output + '.dot',
                    metric=metric_key,
                    colors=palette.colorscheme(),
                    qerror=qerror_map)


def run_on_data(fp32_model, q_model, data, dump_dot_graph, computers):
    curr_dir = Path(__file__).parent.resolve()
    dump_fp32_py = curr_dir / 'visqlib' / 'DumpFP32FM.py'
    dump_fq_py = curr_dir / 'visqlib' / 'DumpFakeQuantFM.py'

    with tempfile.TemporaryDirectory() as model_dir, \
         tempfile.TemporaryDirectory() as fp32_dir, \
         tempfile.TemporaryDirectory() as fq_dir:
        fq_model = model_dir + '/fq_model.circle'

        # Step 1. Fake quantize quantized circle model
        _fake_quantize(q_model, fq_model)

        # Step 2. Run dalgona to dump intermediate FMs in FP32 model
        _run_dalgona(fp32_model, data, dump_fp32_py, fp32_dir)

        # Copy list of dumped tensors
        copyfile(fp32_dir + '/tensors.txt', fq_dir + '/tensors.txt')

        # Step 3. Run dalgona to dump intermediate FMs in fq model
        _run_dalgona(fq_model, data, dump_fq_py, fq_dir)

        # Step 4. Read results and compute qerror
        for metric_key in computers:
            cur_computer = computers[metric_key][0]
            output = computers[metric_key][1]
            cur_computer.advance_on(fp32_dir, fq_dir)
            if metric_key == 'MPEIR':
                qerror_map = cur_computer.get_final_result()
                q_min = 0.0
                q_max = 1.0
            elif metric_key == 'MSE' or metric_key == 'TAE':
                qerror_map, q_min, q_max = cur_computer.get_final_result()

            palette = YLORRD9Palette(qerror_min=q_min, qerror_max=q_max)
            result = _build_json(
                metric=metric_key,
                model=Path(fp32_model).name,
                colorscheme=palette.colorscheme(),
                error=qerror_map)
            with open(output, "w") as f:
                json.dump(result, f)

            if dump_dot_graph:
                _save_dot(
                    circle_path=fp32_model,
                    dot_path=output + '.dot',
                    metric=metric_key,
                    colors=palette.colorscheme(),
                    qerror=qerror_map)


def main():
    # parse arguments
    parser = _get_parser()
    args = parser.parse_args()
    _verify_args(args)

    fp32_model = args.fp32_circle
    q_model = args.q_circle
    data = None
    if valid_attr(args, 'data'):
        data = args.data
    dump_dot_graph = args.dump_dot_graph
    batch_size = None
    if valid_attr(args, 'batch_size'):
        batch_size = args.batch_size

    computers = {}
    if args.mpeir_output:
        computers['MPEIR'] = (MPEIRComputer(None, None), args.mpeir_output)

    if args.mse_output:
        computers['MSE'] = (MSEComputer(None, None), args.mse_output)

    if args.tae_output:
        computers['TAE'] = (TAEComputer(None, None), args.tae_output)

    if batch_size == None:
        run_on_data(fp32_model, q_model, data, dump_dot_graph, computers)
    else:
        run_on_data_batchwise(fp32_model, q_model, data, dump_dot_graph, computers,
                              batch_size)


if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        prog_name = os.path.basename(__file__)
        print(f"{prog_name}: {type(e).__name__}: " + str(e), file=sys.stderr)
        sys.exit(255)