mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,160 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. 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.
|
||||
|
||||
"""
|
||||
Script for calibrating a pretrained ASR model for quantization
|
||||
"""
|
||||
|
||||
from argparse import ArgumentParser
|
||||
|
||||
import torch
|
||||
from omegaconf import open_dict
|
||||
|
||||
from nemo.collections.asr.models import EncDecCTCModel
|
||||
from nemo.utils import logging
|
||||
|
||||
try:
|
||||
from pytorch_quantization import calib
|
||||
from pytorch_quantization import nn as quant_nn
|
||||
from pytorch_quantization import quant_modules
|
||||
from pytorch_quantization.tensor_quant import QuantDescriptor
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"pytorch-quantization is not installed. Install from "
|
||||
"https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
|
||||
)
|
||||
|
||||
try:
|
||||
from torch.cuda.amp import autocast
|
||||
except ImportError:
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def autocast(enabled=None):
|
||||
yield
|
||||
|
||||
|
||||
can_gpu = torch.cuda.is_available()
|
||||
|
||||
|
||||
def main():
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'",
|
||||
)
|
||||
parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data")
|
||||
parser.add_argument("--batch_size", type=int, default=256)
|
||||
parser.add_argument(
|
||||
"--dont_normalize_text",
|
||||
default=False,
|
||||
action='store_false',
|
||||
help="Turn off trasnscript normalization. Recommended for non-English.",
|
||||
)
|
||||
parser.add_argument('--num_calib_batch', default=1, type=int, help="Number of batches for calibration.")
|
||||
parser.add_argument('--calibrator', type=str, choices=["max", "histogram"], default="max")
|
||||
parser.add_argument('--percentile', nargs='+', type=float, default=[99.9, 99.99, 99.999, 99.9999])
|
||||
parser.add_argument("--amp", action="store_true", help="Use AMP in calibration.")
|
||||
parser.set_defaults(amp=False)
|
||||
|
||||
args = parser.parse_args()
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
# Initialize quantization
|
||||
quant_desc_input = QuantDescriptor(calib_method=args.calibrator)
|
||||
quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
|
||||
quant_nn.QuantConvTranspose2d.set_default_quant_desc_input(quant_desc_input)
|
||||
quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
|
||||
|
||||
if args.asr_model.endswith('.nemo'):
|
||||
logging.info(f"Using local ASR model from {args.asr_model}")
|
||||
asr_model_cfg = EncDecCTCModel.restore_from(restore_path=args.asr_model, return_config=True)
|
||||
with open_dict(asr_model_cfg):
|
||||
asr_model_cfg.encoder.quantize = True
|
||||
asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model, override_config_path=asr_model_cfg)
|
||||
|
||||
else:
|
||||
logging.info(f"Using NGC cloud ASR model {args.asr_model}")
|
||||
asr_model_cfg = EncDecCTCModel.from_pretrained(model_name=args.asr_model, return_config=True)
|
||||
with open_dict(asr_model_cfg):
|
||||
asr_model_cfg.encoder.quantize = True
|
||||
asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model, override_config_path=asr_model_cfg)
|
||||
|
||||
asr_model.setup_test_data(
|
||||
test_data_config={
|
||||
'sample_rate': 16000,
|
||||
'manifest_filepath': args.dataset,
|
||||
'labels': asr_model.decoder.vocabulary,
|
||||
'batch_size': args.batch_size,
|
||||
'normalize_transcripts': args.dont_normalize_text,
|
||||
'shuffle': True,
|
||||
}
|
||||
)
|
||||
asr_model.preprocessor.featurizer.dither = 0.0
|
||||
asr_model.preprocessor.featurizer.pad_to = 0
|
||||
if can_gpu:
|
||||
asr_model = asr_model.cuda()
|
||||
asr_model.eval()
|
||||
|
||||
# Enable calibrators
|
||||
for name, module in asr_model.named_modules():
|
||||
if isinstance(module, quant_nn.TensorQuantizer):
|
||||
if module._calibrator is not None:
|
||||
module.disable_quant()
|
||||
module.enable_calib()
|
||||
else:
|
||||
module.disable()
|
||||
|
||||
for i, test_batch in enumerate(asr_model.test_dataloader()):
|
||||
if can_gpu:
|
||||
test_batch = [x.cuda() for x in test_batch]
|
||||
if args.amp:
|
||||
with autocast():
|
||||
_ = asr_model(input_signal=test_batch[0], input_signal_length=test_batch[1])
|
||||
else:
|
||||
_ = asr_model(input_signal=test_batch[0], input_signal_length=test_batch[1])
|
||||
if i >= args.num_calib_batch:
|
||||
break
|
||||
|
||||
# Save calibrated model(s)
|
||||
model_name = args.asr_model.replace(".nemo", "") if args.asr_model.endswith(".nemo") else args.asr_model
|
||||
if not args.calibrator == "histogram":
|
||||
compute_amax(asr_model, method="max")
|
||||
asr_model.save_to(F"{model_name}-max-{args.num_calib_batch*args.batch_size}.nemo")
|
||||
else:
|
||||
for percentile in args.percentile:
|
||||
print(F"{percentile} percentile calibration")
|
||||
compute_amax(asr_model, method="percentile")
|
||||
asr_model.save_to(F"{model_name}-percentile-{percentile}-{args.num_calib_batch*args.batch_size}.nemo")
|
||||
|
||||
for method in ["mse", "entropy"]:
|
||||
print(F"{method} calibration")
|
||||
compute_amax(asr_model, method=method)
|
||||
asr_model.save_to(F"{model_name}-{method}-{args.num_calib_batch*args.batch_size}.nemo")
|
||||
|
||||
|
||||
def compute_amax(model, **kwargs):
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, quant_nn.TensorQuantizer):
|
||||
if module._calibrator is not None:
|
||||
if isinstance(module._calibrator, calib.MaxCalibrator):
|
||||
module.load_calib_amax()
|
||||
else:
|
||||
module.load_calib_amax(**kwargs)
|
||||
print(F"{name:40}: {module}")
|
||||
if can_gpu:
|
||||
model.cuda()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main() # noqa pylint: disable=no-value-for-parameter
|
||||
@@ -0,0 +1,216 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. 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.
|
||||
|
||||
"""
|
||||
Script for post training quantization of ASR models
|
||||
"""
|
||||
|
||||
import collections
|
||||
from argparse import ArgumentParser
|
||||
from pprint import pprint
|
||||
|
||||
import torch
|
||||
from omegaconf import open_dict
|
||||
|
||||
from nemo.collections.asr.metrics.wer import WER, word_error_rate
|
||||
from nemo.collections.asr.models import EncDecCTCModel
|
||||
from nemo.utils import logging
|
||||
|
||||
try:
|
||||
from pytorch_quantization import nn as quant_nn
|
||||
from pytorch_quantization import quant_modules
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"pytorch-quantization is not installed. Install from "
|
||||
"https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
from torch.cuda.amp import autocast
|
||||
except ImportError:
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def autocast(enabled=None):
|
||||
yield
|
||||
|
||||
|
||||
can_gpu = torch.cuda.is_available()
|
||||
|
||||
|
||||
def main():
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'",
|
||||
)
|
||||
parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data")
|
||||
parser.add_argument("--wer_target", type=float, default=None, help="used by test")
|
||||
parser.add_argument("--batch_size", type=int, default=4)
|
||||
parser.add_argument("--wer_tolerance", type=float, default=1.0, help="used by test")
|
||||
parser.add_argument(
|
||||
"--dont_normalize_text",
|
||||
default=False,
|
||||
action='store_false',
|
||||
help="Turn off trasnscript normalization. Recommended for non-English.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_cer", default=False, action='store_true', help="Use Character Error Rate as the evaluation metric"
|
||||
)
|
||||
parser.add_argument('--sensitivity', action="store_true", help="Perform sensitivity analysis")
|
||||
parser.add_argument('--onnx', action="store_true", help="Export to ONNX")
|
||||
parser.add_argument('--quant-disable-keyword', type=str, nargs='+', help='disable quantizers by keyword')
|
||||
args = parser.parse_args()
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
quant_modules.initialize()
|
||||
|
||||
if args.asr_model.endswith('.nemo'):
|
||||
logging.info(f"Using local ASR model from {args.asr_model}")
|
||||
asr_model_cfg = EncDecCTCModel.restore_from(restore_path=args.asr_model, return_config=True)
|
||||
with open_dict(asr_model_cfg):
|
||||
asr_model_cfg.encoder.quantize = True
|
||||
asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model, override_config_path=asr_model_cfg)
|
||||
|
||||
else:
|
||||
logging.info(f"Using NGC cloud ASR model {args.asr_model}")
|
||||
asr_model_cfg = EncDecCTCModel.from_pretrained(model_name=args.asr_model, return_config=True)
|
||||
with open_dict(asr_model_cfg):
|
||||
asr_model_cfg.encoder.quantize = True
|
||||
asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model, override_config_path=asr_model_cfg)
|
||||
asr_model.setup_test_data(
|
||||
test_data_config={
|
||||
'sample_rate': 16000,
|
||||
'manifest_filepath': args.dataset,
|
||||
'labels': asr_model.decoder.vocabulary,
|
||||
'batch_size': args.batch_size,
|
||||
'normalize_transcripts': args.dont_normalize_text,
|
||||
}
|
||||
)
|
||||
asr_model.preprocessor.featurizer.dither = 0.0
|
||||
asr_model.preprocessor.featurizer.pad_to = 0
|
||||
if can_gpu:
|
||||
asr_model = asr_model.cuda()
|
||||
asr_model.eval()
|
||||
|
||||
if args.quant_disable_keyword:
|
||||
for name, module in asr_model.named_modules():
|
||||
if isinstance(module, quant_nn.TensorQuantizer):
|
||||
for keyword in args.quant_disable_keyword:
|
||||
if keyword in name:
|
||||
logging.warning(F"Disable {name}")
|
||||
module.disable()
|
||||
|
||||
labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))])
|
||||
wer = WER(vocabulary=asr_model.decoder.vocabulary, use_cer=args.use_cer)
|
||||
wer_quant = evaluate(asr_model, labels_map, wer)
|
||||
logging.info(f'Got WER of {wer_quant}. Tolerance was {args.wer_tolerance}')
|
||||
|
||||
if args.sensitivity:
|
||||
if wer_quant < args.wer_tolerance:
|
||||
logging.info("Tolerance is already met. Skip sensitivity analyasis.")
|
||||
return
|
||||
quant_layer_names = []
|
||||
for name, module in asr_model.named_modules():
|
||||
if isinstance(module, quant_nn.TensorQuantizer):
|
||||
module.disable()
|
||||
layer_name = name.replace("._input_quantizer", "").replace("._weight_quantizer", "")
|
||||
if layer_name not in quant_layer_names:
|
||||
quant_layer_names.append(layer_name)
|
||||
logging.info(F"{len(quant_layer_names)} quantized layers found.")
|
||||
|
||||
# Build sensitivity profile
|
||||
quant_layer_sensitivity = {}
|
||||
for i, quant_layer in enumerate(quant_layer_names):
|
||||
logging.info(F"Enable {quant_layer}")
|
||||
for name, module in asr_model.named_modules():
|
||||
if isinstance(module, quant_nn.TensorQuantizer) and quant_layer in name:
|
||||
module.enable()
|
||||
logging.info(F"{name:40}: {module}")
|
||||
|
||||
# Eval the model
|
||||
wer_value = evaluate(asr_model, labels_map, wer)
|
||||
logging.info(F"WER: {wer_value}")
|
||||
quant_layer_sensitivity[quant_layer] = args.wer_tolerance - wer_value
|
||||
|
||||
for name, module in asr_model.named_modules():
|
||||
if isinstance(module, quant_nn.TensorQuantizer) and quant_layer in name:
|
||||
module.disable()
|
||||
logging.info(F"{name:40}: {module}")
|
||||
|
||||
# Skip most sensitive layers until WER target is met
|
||||
for name, module in asr_model.named_modules():
|
||||
if isinstance(module, quant_nn.TensorQuantizer):
|
||||
module.enable()
|
||||
quant_layer_sensitivity = collections.OrderedDict(sorted(quant_layer_sensitivity.items(), key=lambda x: x[1]))
|
||||
pprint(quant_layer_sensitivity)
|
||||
skipped_layers = []
|
||||
for quant_layer, _ in quant_layer_sensitivity.items():
|
||||
for name, module in asr_model.named_modules():
|
||||
if isinstance(module, quant_nn.TensorQuantizer):
|
||||
if quant_layer in name:
|
||||
logging.info(F"Disable {name}")
|
||||
if not quant_layer in skipped_layers:
|
||||
skipped_layers.append(quant_layer)
|
||||
module.disable()
|
||||
wer_value = evaluate(asr_model, labels_map, wer)
|
||||
if wer_value <= args.wer_tolerance:
|
||||
logging.info(
|
||||
F"WER tolerance {args.wer_tolerance} is met by skipping {len(skipped_layers)} sensitive layers."
|
||||
)
|
||||
print(skipped_layers)
|
||||
export_onnx(args, asr_model)
|
||||
return
|
||||
raise ValueError(f"WER tolerance {args.wer_tolerance} can not be met with any layer quantized!")
|
||||
|
||||
export_onnx(args, asr_model)
|
||||
|
||||
|
||||
def export_onnx(args, asr_model):
|
||||
if args.onnx:
|
||||
if args.asr_model.endswith("nemo"):
|
||||
onnx_name = args.asr_model.replace(".nemo", ".onnx")
|
||||
else:
|
||||
onnx_name = args.asr_model
|
||||
logging.info(F"Export to {onnx_name}")
|
||||
quant_nn.TensorQuantizer.use_fb_fake_quant = True
|
||||
asr_model.export(onnx_name, onnx_opset_version=13)
|
||||
quant_nn.TensorQuantizer.use_fb_fake_quant = False
|
||||
|
||||
|
||||
def evaluate(asr_model, labels_map, wer):
|
||||
# Eval the model
|
||||
hypotheses = []
|
||||
references = []
|
||||
for test_batch in asr_model.test_dataloader():
|
||||
if can_gpu:
|
||||
test_batch = [x.cuda() for x in test_batch]
|
||||
with autocast():
|
||||
log_probs, encoded_len, greedy_predictions = asr_model(
|
||||
input_signal=test_batch[0], input_signal_length=test_batch[1]
|
||||
)
|
||||
hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions)
|
||||
for batch_ind in range(greedy_predictions.shape[0]):
|
||||
seq_len = test_batch[3][batch_ind].cpu().detach().numpy()
|
||||
seq_ids = test_batch[2][batch_ind].cpu().detach().numpy()
|
||||
reference = ''.join([labels_map[c] for c in seq_ids[0:seq_len]])
|
||||
references.append(reference)
|
||||
del test_batch
|
||||
wer_value = word_error_rate(hypotheses=hypotheses, references=references, use_cer=wer.use_cer)
|
||||
|
||||
return wer_value
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main() # noqa pylint: disable=no-value-for-parameter
|
||||
@@ -0,0 +1,226 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. 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.
|
||||
|
||||
"""
|
||||
Script for inference ASR models using TensorRT
|
||||
"""
|
||||
|
||||
import collections
|
||||
import os
|
||||
import time
|
||||
from argparse import ArgumentParser
|
||||
from pprint import pprint
|
||||
|
||||
import numpy as np
|
||||
import pycuda.autoinit
|
||||
import pycuda.driver as cuda
|
||||
import tensorrt as trt
|
||||
import torch
|
||||
from omegaconf import open_dict
|
||||
|
||||
from nemo.collections.asr.metrics.wer import WER, word_error_rate
|
||||
from nemo.collections.asr.models import EncDecCTCModel
|
||||
from nemo.utils import logging
|
||||
|
||||
TRT_LOGGER = trt.Logger()
|
||||
|
||||
|
||||
can_gpu = torch.cuda.is_available()
|
||||
|
||||
try:
|
||||
from torch.cuda.amp import autocast
|
||||
except ImportError:
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def autocast(enabled=None):
|
||||
yield
|
||||
|
||||
|
||||
def main():
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--asr_onnx",
|
||||
type=str,
|
||||
default="./QuartzNet15x5Base-En-max-32.onnx",
|
||||
help="Pass: 'QuartzNet15x5Base-En-max-32.onnx'",
|
||||
)
|
||||
parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data")
|
||||
parser.add_argument("--batch_size", type=int, default=4)
|
||||
parser.add_argument(
|
||||
"--dont_normalize_text",
|
||||
default=False,
|
||||
action='store_false',
|
||||
help="Turn off trasnscript normalization. Recommended for non-English.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_cer", default=False, action='store_true', help="Use Character Error Rate as the evaluation metric"
|
||||
)
|
||||
parser.add_argument('--qat', action="store_true", help="Use onnx file exported from QAT tools")
|
||||
args = parser.parse_args()
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
if args.asr_model.endswith('.nemo'):
|
||||
logging.info(f"Using local ASR model from {args.asr_model}")
|
||||
asr_model_cfg = EncDecCTCModel.restore_from(restore_path=args.asr_model, return_config=True)
|
||||
with open_dict(asr_model_cfg):
|
||||
asr_model_cfg.encoder.quantize = True
|
||||
asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model, override_config_path=asr_model_cfg)
|
||||
|
||||
else:
|
||||
logging.info(f"Using NGC cloud ASR model {args.asr_model}")
|
||||
asr_model_cfg = EncDecCTCModel.from_pretrained(model_name=args.asr_model, return_config=True)
|
||||
with open_dict(asr_model_cfg):
|
||||
asr_model_cfg.encoder.quantize = True
|
||||
asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model, override_config_path=asr_model_cfg)
|
||||
asr_model.setup_test_data(
|
||||
test_data_config={
|
||||
'sample_rate': 16000,
|
||||
'manifest_filepath': args.dataset,
|
||||
'labels': asr_model.decoder.vocabulary,
|
||||
'batch_size': args.batch_size,
|
||||
'normalize_transcripts': args.dont_normalize_text,
|
||||
}
|
||||
)
|
||||
asr_model.preprocessor.featurizer.dither = 0.0
|
||||
asr_model.preprocessor.featurizer.pad_to = 0
|
||||
if can_gpu:
|
||||
asr_model = asr_model.cuda()
|
||||
asr_model.eval()
|
||||
labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))])
|
||||
wer = WER(vocabulary=asr_model.decoder.vocabulary, use_cer=args.use_cer)
|
||||
wer_result = evaluate(asr_model, args.asr_onnx, labels_map, wer, args.qat)
|
||||
logging.info(f'Got WER of {wer_result}.')
|
||||
|
||||
|
||||
def get_min_max_input_shape(asr_model):
|
||||
max_shape = (1, 64, 1)
|
||||
min_shape = (64, 64, 99999)
|
||||
for test_batch in asr_model.test_dataloader():
|
||||
test_batch = [x.cuda() for x in test_batch]
|
||||
processed_signal, processed_signal_length = asr_model.preprocessor(
|
||||
input_signal=test_batch[0], length=test_batch[1]
|
||||
)
|
||||
shape = processed_signal.cpu().numpy().shape
|
||||
if shape[0] > max_shape[0]:
|
||||
max_shape = (shape[0], *max_shape[1:])
|
||||
if shape[0] < min_shape[0]:
|
||||
min_shape = (shape[0], *min_shape[1:])
|
||||
if shape[2] > max_shape[2]:
|
||||
max_shape = (*max_shape[0:2], shape[2])
|
||||
if shape[2] < min_shape[2]:
|
||||
min_shape = (*min_shape[0:2], shape[2])
|
||||
return min_shape, max_shape
|
||||
|
||||
|
||||
def build_trt_engine(asr_model, onnx_path, qat):
|
||||
trt_engine_path = "{}.trt".format(onnx_path)
|
||||
if os.path.exists(trt_engine_path):
|
||||
return trt_engine_path
|
||||
|
||||
min_input_shape, max_input_shape = get_min_max_input_shape(asr_model)
|
||||
workspace_size = 512
|
||||
with trt.Builder(TRT_LOGGER) as builder:
|
||||
network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
|
||||
if qat:
|
||||
network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_PRECISION)
|
||||
with builder.create_network(flags=network_flags) as network, trt.OnnxParser(
|
||||
network, TRT_LOGGER
|
||||
) as parser, builder.create_builder_config() as builder_config:
|
||||
parser.parse_from_file(onnx_path)
|
||||
builder_config.max_workspace_size = workspace_size * (1024 * 1024)
|
||||
if qat:
|
||||
builder_config.set_flag(trt.BuilderFlag.INT8)
|
||||
|
||||
profile = builder.create_optimization_profile()
|
||||
profile.set_shape("audio_signal", min=min_input_shape, opt=max_input_shape, max=max_input_shape)
|
||||
builder_config.add_optimization_profile(profile)
|
||||
|
||||
engine = builder.build_engine(network, builder_config)
|
||||
serialized_engine = engine.serialize()
|
||||
with open(trt_engine_path, "wb") as fout:
|
||||
fout.write(serialized_engine)
|
||||
return trt_engine_path
|
||||
|
||||
|
||||
def trt_inference(stream, trt_ctx, d_input, d_output, input_signal, input_signal_length):
|
||||
print("infer with shape: {}".format(input_signal.shape))
|
||||
|
||||
trt_ctx.set_binding_shape(0, input_signal.shape)
|
||||
assert trt_ctx.all_binding_shapes_specified
|
||||
|
||||
h_output = cuda.pagelocked_empty(tuple(trt_ctx.get_binding_shape(1)), dtype=np.float32)
|
||||
|
||||
h_input_signal = cuda.register_host_memory(np.ascontiguousarray(input_signal.cpu().numpy().ravel()))
|
||||
cuda.memcpy_htod_async(d_input, h_input_signal, stream)
|
||||
trt_ctx.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
|
||||
cuda.memcpy_dtoh_async(h_output, d_output, stream)
|
||||
stream.synchronize()
|
||||
|
||||
greedy_predictions = torch.tensor(h_output).argmax(dim=-1, keepdim=False)
|
||||
return greedy_predictions
|
||||
|
||||
|
||||
def evaluate(asr_model, asr_onnx, labels_map, wer, qat):
|
||||
# Eval the model
|
||||
hypotheses = []
|
||||
references = []
|
||||
stream = cuda.Stream()
|
||||
vocabulary_size = len(labels_map) + 1
|
||||
engine_file_path = build_trt_engine(asr_model, asr_onnx, qat)
|
||||
with open(engine_file_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
|
||||
trt_engine = runtime.deserialize_cuda_engine(f.read())
|
||||
trt_ctx = trt_engine.create_execution_context()
|
||||
|
||||
profile_shape = trt_engine.get_profile_shape(profile_index=0, binding=0)
|
||||
print("profile shape min:{}, opt:{}, max:{}".format(profile_shape[0], profile_shape[1], profile_shape[2]))
|
||||
max_input_shape = profile_shape[2]
|
||||
input_nbytes = trt.volume(max_input_shape) * trt.float32.itemsize
|
||||
d_input = cuda.mem_alloc(input_nbytes)
|
||||
max_output_shape = [max_input_shape[0], vocabulary_size, (max_input_shape[-1] + 1) // 2]
|
||||
output_nbytes = trt.volume(max_output_shape) * trt.float32.itemsize
|
||||
d_output = cuda.mem_alloc(output_nbytes)
|
||||
|
||||
for test_batch in asr_model.test_dataloader():
|
||||
if can_gpu:
|
||||
test_batch = [x.cuda() for x in test_batch]
|
||||
processed_signal, processed_signal_length = asr_model.preprocessor(
|
||||
input_signal=test_batch[0], length=test_batch[1]
|
||||
)
|
||||
|
||||
greedy_predictions = trt_inference(
|
||||
stream,
|
||||
trt_ctx,
|
||||
d_input,
|
||||
d_output,
|
||||
input_signal=processed_signal,
|
||||
input_signal_length=processed_signal_length,
|
||||
)
|
||||
hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions)
|
||||
for batch_ind in range(greedy_predictions.shape[0]):
|
||||
seq_len = test_batch[3][batch_ind].cpu().detach().numpy()
|
||||
seq_ids = test_batch[2][batch_ind].cpu().detach().numpy()
|
||||
reference = ''.join([labels_map[c] for c in seq_ids[0:seq_len]])
|
||||
references.append(reference)
|
||||
del test_batch
|
||||
wer_value = word_error_rate(hypotheses=hypotheses, references=references, use_cer=wer.use_cer)
|
||||
|
||||
return wer_value
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main() # noqa pylint: disable=no-value-for-parameter
|
||||
Reference in New Issue
Block a user