export_torchscript.py

from utils.hparams import HParam
from dataset.texts import valid_symbols
import utils.fastspeech2_script as fs2
import configargparse
import torch
import sys


def get_parser():

    parser = configargparse.ArgumentParser(
        description='Train a new text-to-speech (TTS) model on one CPU, one or multiple GPUs',
        config_file_parser_class=configargparse.YAMLConfigFileParser,
        formatter_class=configargparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('-c', '--config', type=str, required=True,
                        help="yaml file for configuration")
    parser.add_argument('-n', '--name', type=str, required=True,
                        help="name of the model for logging, saving checkpoint")
    parser.add_argument('--outdir', type=str, required=True,
                        help='Output directory')
    parser.add_argument('-t', '--trace', action='store_true', help="For JIT Trace Module")

    return parser


def main(cmd_args):

    parser = get_parser()
    args, _ = parser.parse_known_args(cmd_args)

    args = parser.parse_args(cmd_args)

    hp = HParam(args.config)

    idim = len(valid_symbols)
    odim = hp.audio.num_mels
    model = fs2.FeedForwardTransformer(idim, odim, hp)
    my_script_module = torch.jit.script(model)
    print("Scripting")
    my_script_module.save("{}/{}.pt".format(args.outdir, args.name))
    print("Script done")
    if args.trace:
        print("Tracing")
        model.eval()
        with torch.no_grad():
            my_trace_module = torch.jit.trace(model, torch.ones(50).to(dtype=torch.int64))
        my_trace_module.save("{}/trace_{}.pt".format(args.outdir, args.name))
        print("Trace Done")


if __name__ == "__main__":
    main(sys.argv[1:])