deepspeech2_quanteval.py

#!/usr/bin/env python3
# pylint: disable=E0401,E1101,W0621,R0915,R0914,R0912
# -*- mode: python -*-
# =============================================================================
#  @@-COPYRIGHT-START-@@
#
#  Copyright (c) 2022 of Qualcomm Innovation Center, Inc. All rights reserved.
#
#  @@-COPYRIGHT-END-@@
# =============================================================================

"""
This script will run AIMET QuantSim and evaluate WER using the DeepSpeech2 model
from the SeanNaren repo with manual configuration of quantization ops.
"""

import os
import argparse
import urllib.request

import torch

from aimet_common.defs import QuantScheme
import aimet_torch
from aimet_torch.quantsim import QuantizationSimModel
from aimet_zoo_torch.common.utils.utils import get_device

import deepspeech_pytorch.model
from deepspeech_pytorch.configs.inference_config import LMConfig
from deepspeech_pytorch.decoder import GreedyDecoder
from deepspeech_pytorch.loader.data_loader import SpectrogramDataset, AudioDataLoader
from deepspeech_pytorch.utils import load_model, load_decoder
from deepspeech_pytorch.testing import run_evaluation


def run_quantsim_evaluation(args):
    """function to run quantization evaluation"""
    device = get_device(args)

    def wrapped_forward_function(self, x, lengths=None):
        if lengths is None:
            lengths = torch.IntTensor([_x.shape[0] for _x in x])
        return self.infer(x, lengths)

    deepspeech_pytorch.model.DeepSpeech.infer = (
        deepspeech_pytorch.model.DeepSpeech.forward
    )
    deepspeech_pytorch.model.DeepSpeech.forward = wrapped_forward_function

    model = load_model(device=device, model_path=args.model_path, use_half=False)

    decoder = load_decoder(labels=model.labels, cfg=LMConfig)

    target_decoder = GreedyDecoder(model.labels, blank_index=model.labels.index("_"))

    def eval_func(model, iterations=None, device=device):
        """evaluation function"""
        model = model.to(device)
        test_dataset = SpectrogramDataset(
            audio_conf=model.audio_conf,
            manifest_filepath=args.test_manifest,
            labels=model.labels,
            normalize=True,
        )

        if iterations is not None:
            test_dataset.size = iterations

        test_loader = AudioDataLoader(
            test_dataset, batch_size=args.batch_size, num_workers=args.num_workers
        )

        wer, cer, output_data = run_evaluation(
            test_loader=test_loader,
            device=device,
            model=model,
            decoder=decoder,
            target_decoder=target_decoder,
            save_output=False,
            verbose=False,
            use_half=False,
        )
        return wer, cer, output_data

    quant_scheme = QuantScheme.post_training_tf_enhanced

    # Test original model on GPU
    # pylint: disable=W0612
    wer, cer, output_data = eval_func(model, None)
    print(f"Original Model | 32-bit Environment | Average WER {wer:.4f}")

    # Simulate original model on device
    sim = QuantizationSimModel(
        model,
        dummy_input=torch.randn(tuple([1, 1, 161, 500])).to(device),
        quant_scheme=quant_scheme,
        default_param_bw=args.default_param_bw,
        default_output_bw=args.default_output_bw,
        config_file=args.quantsim_config_file,
    )

    sim.model.to(device)
    sim.compute_encodings(
        eval_func, forward_pass_callback_args=args.encodings_iterations
    )

    wer, cer, output_data = eval_func(sim.model, None)
    print(
        f"Original Model | {args.default_output_bw}-bit Environment | Average WER {wer:.4f}"
    )

    # Optimize model
    sim = QuantizationSimModel(
        model,
        dummy_input=torch.randn(tuple([1, 1, 161, 500])).to(device),
        quant_scheme=quant_scheme,
        default_param_bw=args.default_param_bw,
        default_output_bw=args.default_output_bw,
        config_file=args.quantsim_config_file,
    )

    manually_configure_quant_ops(sim)

    sim.model.to(device)
    sim.compute_encodings(
        eval_func, forward_pass_callback_args=args.encodings_iterations
    )

    # Test optimized model on device
    wer, cer, output_data = eval_func(sim.model, None)
    print(
        f"Optimized Model | {args.default_output_bw}-bit Environment | Average WER {wer:.4f}"
    )


def manually_configure_quant_ops(sim):
    """
    Manually configure Quantization Ops. Please see documentation for further explanation of quant op placement.
    """

    manual_config = {
        "conv.seq_module.0": {  # Conv2d
            "input_quantizer": True,
            "output_quantizer": False,
            "weight_quantizer": True,
            "bias_quantizer": False,
        },
        "conv.seq_module.1": {  # BatchNorm
            "input_quantizer": False,
            "output_quantizer": False,
            "weight_quantizer": False,
            "bias_quantizer": False,
        },
        "conv.seq_module.2": {  # HardTanh
            "input_quantizer": True,
            "output_quantizer": False,
        },
        "conv.seq_module.3": {  # Conv2d
            "input_quantizer": True,
            "output_quantizer": False,
            "weight_quantizer": True,
            "bias_quantizer": False,
        },
        "conv.seq_module.4": {  # BatchNorm
            "input_quantizer": False,
            "output_quantizer": False,
            "weight_quantizer": False,
            "bias_quantizer": False,
        },
        "conv.seq_module.5": {  # HardTanh
            "input_quantizer": True,
            "output_quantizer": False,
        },
        "rnns.0.rnn": {
            "input_l0_quantizer": True,
            "initial_h_l0_quantizer": False,
            "initial_c_l0_quantizer": False,
            "h_l0_quantizer": True,
            "c_l0_quantizer": False,
            "weight_ih_l0_quantizer": True,
            "weight_hh_l0_quantizer": True,
            "bias_ih_l0_quantizer": False,
            "bias_hh_l0_quantizer": False,
            "weight_ih_l0_reverse_quantizer": True,
            "weight_hh_l0_reverse_quantizer": True,
            "bias_ih_l0_reverse_quantizer": False,
            "bias_hh_l0_reverse_quantizer": False,
        },
        "rnns.1.batch_norm.module": {
            "input_quantizer": False,
            "output_quantizer": False,
            "weight_quantizer": False,
            "bias_quantizer": False,
        },
        "rnns.1.rnn": {
            "input_l0_quantizer": True,
            "initial_h_l0_quantizer": False,
            "initial_c_l0_quantizer": False,
            "h_l0_quantizer": True,
            "c_l0_quantizer": False,
            "weight_ih_l0_quantizer": True,
            "weight_hh_l0_quantizer": True,
            "bias_ih_l0_quantizer": False,
            "bias_hh_l0_quantizer": False,
            "weight_ih_l0_reverse_quantizer": True,
            "weight_hh_l0_reverse_quantizer": True,
            "bias_ih_l0_reverse_quantizer": False,
            "bias_hh_l0_reverse_quantizer": False,
        },
        "rnns.2.batch_norm.module": {
            "input_quantizer": False,
            "output_quantizer": False,
            "weight_quantizer": False,
            "bias_quantizer": False,
        },
        "rnns.2.rnn": {
            "input_l0_quantizer": True,
            "initial_h_l0_quantizer": False,
            "initial_c_l0_quantizer": False,
            "h_l0_quantizer": True,
            "c_l0_quantizer": False,
            "weight_ih_l0_quantizer": True,
            "weight_hh_l0_quantizer": True,
            "bias_ih_l0_quantizer": False,
            "bias_hh_l0_quantizer": False,
            "weight_ih_l0_reverse_quantizer": True,
            "weight_hh_l0_reverse_quantizer": True,
            "bias_ih_l0_reverse_quantizer": False,
            "bias_hh_l0_reverse_quantizer": False,
        },
        "rnns.3.batch_norm.module": {
            "input_quantizer": False,
            "output_quantizer": False,
            "weight_quantizer": False,
            "bias_quantizer": False,
        },
        "rnns.3.rnn": {
            "input_l0_quantizer": True,
            "initial_h_l0_quantizer": False,
            "initial_c_l0_quantizer": False,
            "h_l0_quantizer": True,
            "c_l0_quantizer": False,
            "weight_ih_l0_quantizer": True,
            "weight_hh_l0_quantizer": True,
            "bias_ih_l0_quantizer": False,
            "bias_hh_l0_quantizer": False,
            "weight_ih_l0_reverse_quantizer": True,
            "weight_hh_l0_reverse_quantizer": True,
            "bias_ih_l0_reverse_quantizer": False,
            "bias_hh_l0_reverse_quantizer": False,
        },
        "rnns.4.batch_norm.module": {
            "input_quantizer": False,
            "output_quantizer": False,
            "weight_quantizer": False,
            "bias_quantizer": False,
        },
        "rnns.4.rnn": {
            "input_l0_quantizer": True,
            "initial_h_l0_quantizer": False,
            "initial_c_l0_quantizer": False,
            "h_l0_quantizer": True,
            "c_l0_quantizer": False,
            "weight_ih_l0_quantizer": True,
            "weight_hh_l0_quantizer": True,
            "bias_ih_l0_quantizer": False,
            "bias_hh_l0_quantizer": False,
            "weight_ih_l0_reverse_quantizer": True,
            "weight_hh_l0_reverse_quantizer": True,
            "bias_ih_l0_reverse_quantizer": False,
            "bias_hh_l0_reverse_quantizer": False,
        },
        "fc.0.module.0": {
            "input_quantizer": True,
            "output_quantizer": False,
            "weight_quantizer": False,
            "bias_quantizer": False,
        },
        "fc.0.module.1": {
            "input_quantizer": True,
            "output_quantizer": False,
            "weight_quantizer": True,
        },
        "inference_softmax": {
            "input_quantizer": False,
            "output_quantizer": True,
        },
    }
    # pylint: disable=W0212
    quant_ops = QuantizationSimModel._get_qc_quantized_layers(sim.model)
    for name, op in quant_ops:
        mc = manual_config[name]
        if isinstance(op, aimet_torch.qc_quantize_op.QcPostTrainingWrapper):
            op.input_quantizer.enabled = mc["input_quantizer"]
            op.output_quantizer.enabled = mc["output_quantizer"]
            for q_name, param_quantizer in op.param_quantizers.items():
                param_quantizer.enabled = mc[q_name + "_quantizer"]
        elif isinstance(op, aimet_torch.qc_quantize_recurrent.QcQuantizeRecurrent):
            for q_name, input_quantizer in op.input_quantizers.items():
                input_quantizer.enabled = mc[q_name + "_quantizer"]
            for q_name, output_quantizer in op.output_quantizers.items():
                output_quantizer.enabled = mc[q_name + "_quantizer"]
            for q_name, param_quantizer in op.param_quantizers.items():
                param_quantizer.enabled = mc[q_name + "_quantizer"]


def arguments():
    """argument parser"""
    parser = argparse.ArgumentParser(
        description="Evaluation script for an DeepSpeech2 network."
    )
    parser.add_argument(
        "--test-manifest",
        help="Path to csv to do eval on.",
        type=str,
        default="libri_test_clean_manifest.csv",
    )
    parser.add_argument("--batch-size", help="Batch size.", type=int, default=20)
    parser.add_argument("--num-workers", help="Number of workers.", type=int, default=1)
    parser.add_argument(
        "--default-output-bw",
        help="Default output bitwidth for quantization.",
        type=int,
        default=8,
    )
    parser.add_argument(
        "--default-param-bw",
        help="Default parameter bitwidth for quantization.",
        type=int,
        default=8,
    )
    parser.add_argument("--use-cuda", help="Whether to use cuda", default=True)
    args = parser.parse_args()
    return args


def download_weights():
    """Downloading model weights"""
    # Download original model
    if not os.path.exists("./librispeech_pretrained_v2.pth"):
        urllib.request.urlretrieve(
            "https://github.com/SeanNaren/deepspeech.pytorch/releases/download/v2.0/librispeech_pretrained_v2.pth",
            "librispeech_pretrained_v2.pth",
        )

    # Download config file
    if not os.path.exists("./default_config.json"):
        urllib.request.urlretrieve(
            "https://raw.githubusercontent.com/quic/aimet/release-aimet-1.22.1/TrainingExtensions/common/src/python/aimet_common/quantsim_config/default_config_per_channel.json",
            "default_config.json",
        )


class ModelConfig:
    """Hardcoded model configuration"""

    def __init__(self, args):
        self.model_path = "librispeech_pretrained_v2.pth"
        self.quantsim_config_file = "default_config.json"
        self.encodings_iterations = 500
        self.round_mode = "nearest"
        for arg in vars(args):
            setattr(self, arg, getattr(args, arg))


def main():
    """Main function"""
    args = arguments()
    config = ModelConfig(args)
    run_quantsim_evaluation(config)


if __name__ == "__main__":
    download_weights()
    main()