xref: /aosp_15_r20/external/executorch/examples/qualcomm/scripts/wav2letter.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)

# Copyright (c) Qualcomm Innovation Center, Inc.
# All rights reserved
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import json
import os
import sys
from multiprocessing.connection import Client

import numpy as np

import torch
from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype
from executorch.examples.models.wav2letter import Wav2LetterModel
from executorch.examples.qualcomm.utils import (
    build_executorch_binary,
    make_output_dir,
    parse_skip_delegation_node,
    setup_common_args_and_variables,
    SimpleADB,
)


class Conv2D(torch.nn.Module):
    def __init__(self, stride, padding, weight, bias=None):
        super().__init__()
        use_bias = bias is not None
        self.conv = torch.nn.Conv2d(
            in_channels=weight.shape[1],
            out_channels=weight.shape[0],
            kernel_size=[weight.shape[2], 1],
            stride=[*stride, 1],
            padding=[*padding, 0],
            bias=use_bias,
        )
        self.conv.weight = torch.nn.Parameter(weight.unsqueeze(-1))
        if use_bias:
            self.conv.bias = torch.nn.Parameter(bias)

    def forward(self, x):
        return self.conv(x)


def get_dataset(data_size, artifact_dir):
    from torch.utils.data import DataLoader
    from torchaudio.datasets import LIBRISPEECH

    def collate_fun(batch):
        waves, labels = [], []

        for wave, _, text, *_ in batch:
            waves.append(wave.squeeze(0))
            labels.append(text)
        # need padding here for static ouput shape
        waves = torch.nn.utils.rnn.pad_sequence(waves, batch_first=True)
        return waves, labels

    dataset = LIBRISPEECH(artifact_dir, url="test-clean", download=True)
    data_loader = DataLoader(
        dataset=dataset,
        batch_size=data_size,
        shuffle=True,
        collate_fn=lambda x: collate_fun(x),
    )
    # prepare input data
    inputs, targets, input_list = [], [], ""
    for wave, label in data_loader:
        for index in range(data_size):
            # reshape input tensor to NCHW
            inputs.append((wave[index].reshape(1, 1, -1, 1),))
            targets.append(label[index])
            input_list += f"input_{index}_0.raw\n"
        # here we only take first batch, i.e. 'data_size' tensors
        break

    return inputs, targets, input_list


def eval_metric(pred, target_str):
    from torchmetrics.text import CharErrorRate, WordErrorRate

    def parse(ids):
        vocab = " abcdefghijklmnopqrstuvwxyz'*"
        return ["".join([vocab[c] for c in id]).replace("*", "").upper() for id in ids]

    pred_str = parse(
        [
            torch.unique_consecutive(pred[i, :, :].argmax(0))
            for i in range(pred.shape[0])
        ]
    )
    wer, cer = WordErrorRate(), CharErrorRate()
    return wer(pred_str, target_str), cer(pred_str, target_str)


def main(args):
    skip_node_id_set, skip_node_op_set = parse_skip_delegation_node(args)

    # ensure the working directory exist
    os.makedirs(args.artifact, exist_ok=True)

    if not args.compile_only and args.device is None:
        raise RuntimeError(
            "device serial is required if not compile only. "
            "Please specify a device serial by -s/--device argument."
        )

    instance = Wav2LetterModel()
    # target labels " abcdefghijklmnopqrstuvwxyz'*"
    instance.vocab_size = 29
    model = instance.get_eager_model().eval()
    model.load_state_dict(torch.load(args.pretrained_weight, weights_only=True))

    # convert conv1d to conv2d in nn.Module level will only introduce 2 permute
    # nodes around input & output, which is more quantization friendly.
    for i in range(len(model.acoustic_model)):
        for j in range(len(model.acoustic_model[i])):
            module = model.acoustic_model[i][j]
            if isinstance(module, torch.nn.Conv1d):
                model.acoustic_model[i][j] = Conv2D(
                    stride=module.stride,
                    padding=module.padding,
                    weight=module.weight,
                    bias=module.bias,
                )

    # retrieve dataset, will take some time to download
    data_num = 100
    inputs, targets, input_list = get_dataset(
        data_size=data_num, artifact_dir=args.artifact
    )
    pte_filename = "w2l_qnn"
    build_executorch_binary(
        model,
        inputs[0],
        args.model,
        f"{args.artifact}/{pte_filename}",
        inputs,
        skip_node_id_set=skip_node_id_set,
        skip_node_op_set=skip_node_op_set,
        quant_dtype=QuantDtype.use_8a8w,
        shared_buffer=args.shared_buffer,
    )

    if args.compile_only:
        sys.exit(0)

    adb = SimpleADB(
        qnn_sdk=os.getenv("QNN_SDK_ROOT"),
        build_path=f"{args.build_folder}",
        pte_path=f"{args.artifact}/{pte_filename}.pte",
        workspace=f"/data/local/tmp/executorch/{pte_filename}",
        device_id=args.device,
        host_id=args.host,
        soc_model=args.model,
        shared_buffer=args.shared_buffer,
    )
    adb.push(inputs=inputs, input_list=input_list)
    adb.execute()

    # collect output data
    output_data_folder = f"{args.artifact}/outputs"
    make_output_dir(output_data_folder)
    adb.pull(output_path=args.artifact)

    predictions = []
    for i in range(data_num):
        predictions.append(
            np.fromfile(
                os.path.join(output_data_folder, f"output_{i}_0.raw"), dtype=np.float32
            )
        )

    # evaluate metrics
    wer, cer = 0, 0
    for i, pred in enumerate(predictions):
        pred = torch.from_numpy(pred).reshape(1, instance.vocab_size, -1)
        wer_eval, cer_eval = eval_metric(pred, targets[i])
        wer += wer_eval
        cer += cer_eval

    if args.ip and args.port != -1:
        with Client((args.ip, args.port)) as conn:
            conn.send(
                json.dumps({"wer": wer.item() / data_num, "cer": cer.item() / data_num})
            )
    else:
        print(f"wer: {wer / data_num}\ncer: {cer / data_num}")


if __name__ == "__main__":
    parser = setup_common_args_and_variables()

    parser.add_argument(
        "-a",
        "--artifact",
        help="path for storing generated artifacts by this example. "
        "Default ./wav2letter",
        default="./wav2letter",
        type=str,
    )

    parser.add_argument(
        "-p",
        "--pretrained_weight",
        help=(
            "Location of pretrained weight, please download via "
            "https://github.com/nipponjo/wav2letter-ctc-pytorch/tree/main?tab=readme-ov-file#wav2letter-ctc-pytorch"
            " for torchaudio.models.Wav2Letter version"
        ),
        default=None,
        type=str,
        required=True,
    )

    args = parser.parse_args()
    try:
        main(args)
    except Exception as e:
        if args.ip and args.port != -1:
            with Client((args.ip, args.port)) as conn:
                conn.send(json.dumps({"Error": str(e)}))
        else:
            raise Exception(e)