108 lines
4.7 KiB
Python
108 lines
4.7 KiB
Python
# *****************************************************************************
|
|
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are met:
|
|
# * Redistributions of source code must retain the above copyright
|
|
# notice, this list of conditions and the following disclaimer.
|
|
# * Redistributions in binary form must reproduce the above copyright
|
|
# notice, this list of conditions and the following disclaimer in the
|
|
# documentation and/or other materials provided with the distribution.
|
|
# * Neither the name of the NVIDIA CORPORATION nor the
|
|
# names of its contributors may be used to endorse or promote products
|
|
# derived from this software without specific prior written permission.
|
|
#
|
|
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
|
|
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
|
|
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
|
|
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
|
|
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
|
|
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
|
|
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
|
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
|
|
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
#
|
|
# *****************************************************************************
|
|
|
|
# Edited from https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/SpeechSynthesis/Tacotron2/trt/export_onnx2trt.py
|
|
|
|
import pycuda.driver as cuda
|
|
import pycuda.autoinit
|
|
import tensorrt as trt
|
|
import onnx
|
|
import argparse
|
|
|
|
import sys
|
|
sys.path.append('./')
|
|
|
|
def parse_args(parser):
|
|
"""
|
|
Parse commandline arguments.
|
|
"""
|
|
parser.add_argument('-o', '--output', required=True,
|
|
help='output folder to save audio (file per phrase)')
|
|
parser.add_argument('--waveglow', type=str, default="",
|
|
help='full path to the WaveGlow ONNX')
|
|
parser.add_argument('--fp16', action='store_true',
|
|
help='inference with FP16')
|
|
parser.add_argument('-b', '--batch_size', default=1, type=int,
|
|
help='batch size for inference.')
|
|
parser.add_argument('-w', '--max_ws', default=1, type=int,
|
|
help='max workspace size in GB.')
|
|
|
|
return parser
|
|
|
|
|
|
def build_engine(model_file, shapes, max_ws=512*1024*1024, fp16=False):
|
|
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
|
|
builder = trt.Builder(TRT_LOGGER)
|
|
builder.fp16_mode = fp16
|
|
|
|
config = builder.create_builder_config()
|
|
config.max_workspace_size = max_ws
|
|
if fp16:
|
|
config.flags |= 1 << int(trt.BuilderFlag.FP16)
|
|
profile = builder.create_optimization_profile()
|
|
for s in shapes:
|
|
profile.set_shape(s['name'], min=s['min'], opt=s['opt'], max=s['max'])
|
|
config.add_optimization_profile(profile)
|
|
explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
|
|
network = builder.create_network(explicit_batch)
|
|
|
|
with trt.OnnxParser(network, TRT_LOGGER) as parser:
|
|
with open(model_file, 'rb') as model:
|
|
parsed = parser.parse(model.read())
|
|
for i in range(parser.num_errors):
|
|
print("TensorRT ONNX parser error:", parser.get_error(i))
|
|
engine = builder.build_engine(network, config=config)
|
|
|
|
return engine
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description='Export from ONNX to TensorRT for WaveGlow')
|
|
parser = parse_args(parser)
|
|
args = parser.parse_args()
|
|
|
|
engine_prec = ".fp16" if args.fp16 else ".fp32"
|
|
|
|
# WaveGlow
|
|
batch_size = args.batch_size
|
|
shapes=[{"name": "mel", "min": (batch_size,80,32,1), "opt": (batch_size,80,768,1), "max": (batch_size,80,1024,1)},
|
|
{"name": "z", "min": (batch_size,8,1024,1), "opt": (batch_size,8,24576,1), "max": (batch_size,8,32768,1)}]
|
|
if args.waveglow != "":
|
|
print("Building WaveGlow ...")
|
|
waveglow_engine = build_engine(args.waveglow, shapes=shapes, fp16=args.fp16, max_ws=args.max_ws * 1<<30)
|
|
if waveglow_engine is not None:
|
|
with open(args.output+"/"+"waveglow"+engine_prec+".b"+str(batch_size), 'wb') as f:
|
|
f.write(waveglow_engine.serialize())
|
|
else:
|
|
print("Failed to build engine from", args.waveglow)
|
|
sys.exit()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|