mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-09 10:02:22 +00:00
372 lines
11 KiB
Python
Executable File
372 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
#
|
|
# Copyright 2023 Xiaomi Corporation (Author: Fangjun Kuang, Wei Kang)
|
|
# Copyright 2023 Danqing Fu (danqing.fu@gmail.com)
|
|
|
|
"""
|
|
This script exports a transducer model from PyTorch to ONNX.
|
|
|
|
We use the pre-trained model from
|
|
https://huggingface.co/Zengwei/icefall-asr-librispeech-zipformer-2023-05-15
|
|
as an example to show how to use this file.
|
|
|
|
1. Download the pre-trained model
|
|
|
|
cd egs/librispeech/ASR
|
|
|
|
repo_url=https://huggingface.co/Zengwei/icefall-asr-librispeech-zipformer-2023-05-15
|
|
GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
|
|
repo=$(basename $repo_url)
|
|
|
|
pushd $repo
|
|
git lfs pull --include "exp/pretrained.pt"
|
|
|
|
cd exp
|
|
ln -s pretrained.pt epoch-99.pt
|
|
popd
|
|
|
|
2. Export the model to ONNX
|
|
|
|
./zipformer/export-onnx.py \
|
|
--tokens $repo/data/lang_bpe_500/tokens.txt \
|
|
--use-averaged-model 0 \
|
|
--epoch 99 \
|
|
--avg 1 \
|
|
--exp-dir $repo/exp \
|
|
--num-encoder-layers "2,2,3,4,3,2" \
|
|
--downsampling-factor "1,2,4,8,4,2" \
|
|
--feedforward-dim "512,768,1024,1536,1024,768" \
|
|
--num-heads "4,4,4,8,4,4" \
|
|
--encoder-dim "192,256,384,512,384,256" \
|
|
--query-head-dim 32 \
|
|
--value-head-dim 12 \
|
|
--pos-head-dim 4 \
|
|
--pos-dim 48 \
|
|
--encoder-unmasked-dim "192,192,256,256,256,192" \
|
|
--cnn-module-kernel "31,31,15,15,15,31" \
|
|
--decoder-dim 512 \
|
|
--joiner-dim 512 \
|
|
--causal False \
|
|
--chunk-size "16,32,64,-1" \
|
|
--left-context-frames "64,128,256,-1" \
|
|
--fp16 True
|
|
It will generate the following 3 files inside $repo/exp:
|
|
|
|
- encoder-epoch-99-avg-1.onnx
|
|
- decoder-epoch-99-avg-1.onnx
|
|
- joiner-epoch-99-avg-1.onnx
|
|
|
|
See ./onnx_pretrained.py and ./onnx_check.py for how to
|
|
use the exported ONNX models.
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Dict, Tuple
|
|
|
|
import onnx
|
|
import torch
|
|
import torch.nn as nn
|
|
from onnxconverter_common import float16
|
|
from onnxruntime.quantization import QuantType, quantize_dynamic
|
|
from train import add_model_arguments, get_model, get_params
|
|
|
|
from icefall.checkpoint import (
|
|
average_checkpoints,
|
|
average_checkpoints_with_averaged_model,
|
|
find_checkpoints,
|
|
load_checkpoint,
|
|
)
|
|
from icefall.utils import make_pad_mask, num_tokens, str2bool
|
|
|
|
|
|
def get_parser():
|
|
parser = argparse.ArgumentParser(
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--sampling-rate",
|
|
type=int,
|
|
default=24000,
|
|
help="The sampleing rate of libritts dataset",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--frame-shift",
|
|
type=int,
|
|
default=256,
|
|
help="Frame shift.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--frame-length",
|
|
type=int,
|
|
default=1024,
|
|
help="Frame shift.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--epoch",
|
|
type=int,
|
|
default=28,
|
|
help="""It specifies the checkpoint to use for averaging.
|
|
Note: Epoch counts from 0.
|
|
You can specify --avg to use more checkpoints for model averaging.""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--iter",
|
|
type=int,
|
|
default=0,
|
|
help="""If positive, --epoch is ignored and it
|
|
will use the checkpoint exp_dir/checkpoint-iter.pt.
|
|
You can specify --avg to use more checkpoints for model averaging.
|
|
""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--avg",
|
|
type=int,
|
|
default=15,
|
|
help="Number of checkpoints to average. Automatically select "
|
|
"consecutive checkpoints before the checkpoint specified by "
|
|
"'--epoch' and '--iter'",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--use-averaged-model",
|
|
type=str2bool,
|
|
default=True,
|
|
help="Whether to load averaged model. Currently it only supports "
|
|
"using --epoch. If True, it would decode with the averaged model "
|
|
"over the epoch range from `epoch-avg` (excluded) to `epoch`."
|
|
"Actually only the models with epoch number of `epoch-avg` and "
|
|
"`epoch` are loaded for averaging. ",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--exp-dir",
|
|
type=str,
|
|
default="zipformer/exp",
|
|
help="""It specifies the directory where all training related
|
|
files, e.g., checkpoints, log, etc, are saved
|
|
""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--fp16",
|
|
type=str2bool,
|
|
default=False,
|
|
help="Whether to export models in fp16",
|
|
)
|
|
|
|
add_model_arguments(parser)
|
|
|
|
return parser
|
|
|
|
|
|
def add_meta_data(filename: str, meta_data: Dict[str, str]):
|
|
"""Add meta data to an ONNX model. It is changed in-place.
|
|
|
|
Args:
|
|
filename:
|
|
Filename of the ONNX model to be changed.
|
|
meta_data:
|
|
Key-value pairs.
|
|
"""
|
|
model = onnx.load(filename)
|
|
for key, value in meta_data.items():
|
|
meta = model.metadata_props.add()
|
|
meta.key = key
|
|
meta.value = value
|
|
|
|
onnx.save(model, filename)
|
|
|
|
|
|
def export_model_onnx(
|
|
model: nn.Module,
|
|
model_filename: str,
|
|
opset_version: int = 13,
|
|
) -> None:
|
|
"""Export the joiner model to ONNX format.
|
|
The exported joiner model has two inputs:
|
|
|
|
- encoder_out: a tensor of shape (N, joiner_dim)
|
|
- decoder_out: a tensor of shape (N, joiner_dim)
|
|
|
|
and produces one output:
|
|
|
|
- logit: a tensor of shape (N, vocab_size)
|
|
"""
|
|
input_tensor = torch.rand((2, 80, 100), dtype=torch.float32)
|
|
|
|
torch.onnx.export(
|
|
model,
|
|
(input_tensor,),
|
|
model_filename,
|
|
verbose=False,
|
|
opset_version=opset_version,
|
|
input_names=[
|
|
"features",
|
|
],
|
|
output_names=["audio"],
|
|
dynamic_axes={
|
|
"features": {0: "N", 2: "F"},
|
|
"audio": {0: "N", 1: "T"},
|
|
},
|
|
)
|
|
|
|
meta_data = {
|
|
"model_type": "Vocos",
|
|
"version": "1",
|
|
"model_author": "k2-fsa",
|
|
"comment": "ConvNext Vocos",
|
|
}
|
|
logging.info(f"meta_data: {meta_data}")
|
|
|
|
add_meta_data(filename=model_filename, meta_data=meta_data)
|
|
|
|
|
|
@torch.no_grad()
|
|
def main():
|
|
args = get_parser().parse_args()
|
|
args.exp_dir = Path(args.exp_dir)
|
|
|
|
params = get_params()
|
|
params.update(vars(args))
|
|
|
|
device = torch.device("cpu")
|
|
params.device = device
|
|
logging.info(params)
|
|
|
|
logging.info("About to create model")
|
|
model = get_model(params)
|
|
model.to(device)
|
|
|
|
if not params.use_averaged_model:
|
|
if params.iter > 0:
|
|
filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
|
|
: params.avg
|
|
]
|
|
if len(filenames) == 0:
|
|
raise ValueError(
|
|
f"No checkpoints found for"
|
|
f" --iter {params.iter}, --avg {params.avg}"
|
|
)
|
|
elif len(filenames) < params.avg:
|
|
raise ValueError(
|
|
f"Not enough checkpoints ({len(filenames)}) found for"
|
|
f" --iter {params.iter}, --avg {params.avg}"
|
|
)
|
|
logging.info(f"averaging {filenames}")
|
|
model.to(device)
|
|
model.load_state_dict(average_checkpoints(filenames, device=device))
|
|
elif params.avg == 1:
|
|
load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model)
|
|
else:
|
|
start = params.epoch - params.avg + 1
|
|
filenames = []
|
|
for i in range(start, params.epoch + 1):
|
|
if i >= 1:
|
|
filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
|
|
logging.info(f"averaging {filenames}")
|
|
model.to(device)
|
|
model.load_state_dict(average_checkpoints(filenames, device=device))
|
|
else:
|
|
if params.iter > 0:
|
|
filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
|
|
: params.avg + 1
|
|
]
|
|
if len(filenames) == 0:
|
|
raise ValueError(
|
|
f"No checkpoints found for"
|
|
f" --iter {params.iter}, --avg {params.avg}"
|
|
)
|
|
elif len(filenames) < params.avg + 1:
|
|
raise ValueError(
|
|
f"Not enough checkpoints ({len(filenames)}) found for"
|
|
f" --iter {params.iter}, --avg {params.avg}"
|
|
)
|
|
filename_start = filenames[-1]
|
|
filename_end = filenames[0]
|
|
logging.info(
|
|
"Calculating the averaged model over iteration checkpoints"
|
|
f" from {filename_start} (excluded) to {filename_end}"
|
|
)
|
|
model.to(device)
|
|
model.load_state_dict(
|
|
average_checkpoints_with_averaged_model(
|
|
filename_start=filename_start,
|
|
filename_end=filename_end,
|
|
device=device,
|
|
)
|
|
)
|
|
else:
|
|
assert params.avg > 0, params.avg
|
|
start = params.epoch - params.avg
|
|
assert start >= 1, start
|
|
filename_start = f"{params.exp_dir}/epoch-{start}.pt"
|
|
filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
|
|
logging.info(
|
|
f"Calculating the averaged model over epoch range from "
|
|
f"{start} (excluded) to {params.epoch}"
|
|
)
|
|
model.to(device)
|
|
model.load_state_dict(
|
|
average_checkpoints_with_averaged_model(
|
|
filename_start=filename_start,
|
|
filename_end=filename_end,
|
|
device=device,
|
|
)
|
|
)
|
|
|
|
model.eval()
|
|
vocos = model.generator
|
|
|
|
if params.iter > 0:
|
|
suffix = f"iter-{params.iter}"
|
|
else:
|
|
suffix = f"epoch-{params.epoch}"
|
|
|
|
suffix += f"-avg-{params.avg}"
|
|
|
|
opset_version = 13
|
|
|
|
logging.info("Exporting model")
|
|
model_filename = params.exp_dir / f"vocos-{suffix}.onnx"
|
|
export_model_onnx(
|
|
vocos,
|
|
model_filename,
|
|
opset_version=opset_version,
|
|
)
|
|
logging.info(f"Exported vocos generator to {model_filename}")
|
|
|
|
if params.fp16:
|
|
logging.info("Generate fp16 models")
|
|
|
|
model = onnx.load(model_filename)
|
|
model_fp16 = float16.convert_float_to_float16(model, keep_io_types=True)
|
|
model_filename_fp16 = params.exp_dir / f"vocos-{suffix}.fp16.onnx"
|
|
onnx.save(model_fp16, model_filename_fp16)
|
|
|
|
# Generate int8 quantization models
|
|
# See https://onnxruntime.ai/docs/performance/model-optimizations/quantization.html#data-type-selection
|
|
|
|
logging.info("Generate int8 quantization models")
|
|
|
|
model_filename_int8 = params.exp_dir / f"vocos-{suffix}.int8.onnx"
|
|
quantize_dynamic(
|
|
model_input=model_filename,
|
|
model_output=model_filename_int8,
|
|
op_types_to_quantize=["MatMul"],
|
|
weight_type=QuantType.QInt8,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
|
|
logging.basicConfig(format=formatter, level=logging.INFO)
|
|
main()
|