mirror of
https://github.com/storytold/torchscript-models.git
synced 2026-10-18 07:33:20 +00:00
Brandon Thomas
72 lines
2.5 KiB
Python
Executable File
72 lines
2.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# From: https://github.com/pytorch/pytorch/issues/20356#issuecomment-545572400
|
|
import torch
|
|
|
|
from collections import OrderedDict
|
|
import argparse
|
|
import re
|
|
import os
|
|
import sys
|
|
import yaml
|
|
|
|
from model_melgan.generator import load_melgan
|
|
|
|
#torch.set_default_tensor_type('torch.DoubleTensor')
|
|
|
|
DEFAULT_NAME = 'melgan-vocodes.torchjit'
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser('JITswap the melgan')
|
|
parser.add_argument('--model_file', type=str, required=True) # Input
|
|
parser.add_argument('--output_file', type=str, required=False) # Explicitly set the filename
|
|
parser.add_argument('--copy_name', default=False, action='store_true') # Save as '{model_name}.torchjit'
|
|
parser.add_argument('--output_directory', type=str, required=False, default='./OUTPUT/')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not args.model_file.endswith('.pt'):
|
|
raise Exception('Melgan model filename `{}` does not end with extention `.pt`.'.format(args.model_file))
|
|
|
|
if args.output_file and not args.output_file.endswith('.torchjit'):
|
|
raise Exception('(Optional) Melgan JIT output filename `{}` does not end with extention `.torchjit`.'.format(args.output_file))
|
|
|
|
if args.copy_name:
|
|
if args.output_file:
|
|
raise Exception('Cannot --copy_name and set an explicit --output_file')
|
|
basename = os.path.basename(args.model_file)
|
|
args.output_file = os.path.splitext(basename)[0] + '.torchjit'
|
|
else:
|
|
args.output_file = DEFAULT_NAME
|
|
|
|
return args
|
|
|
|
args = parse_args()
|
|
|
|
|
|
model_checkpoint_path = args.model_file
|
|
output_filename = os.path.join(args.output_directory, args.output_file)
|
|
|
|
print('model file', model_checkpoint_path)
|
|
print('output file', output_filename)
|
|
|
|
melgan_model = load_melgan(model_checkpoint_path)
|
|
|
|
|
|
# NB: Tracing evaluates the model on input and unrolls and hardcodes branching
|
|
# and loops. Scripting allows these to remain by converting the entire program
|
|
# to TorchScript. Scripting seems MUCH faster, is also harder get working.
|
|
trace_model = True
|
|
if trace_model:
|
|
print('Tracing model and saving as: {}'.format(output_filename))
|
|
mel_file = './sample/LJ002-0320.mel'
|
|
mel = torch.load(mel_file, map_location=torch.device('cpu'))
|
|
# NB: Getting `Tracing failed sanity checks! Graphs differed across invocations!`
|
|
traced_script_module = torch.jit.trace(melgan_model, mel, check_trace=True)
|
|
traced_script_module.save(output_filename)
|
|
else:
|
|
print('Saving script model as: {}'.format(output_filename))
|
|
container = torch.jit.script(melgan_model)
|
|
container.save(output_filename)
|
|
|
|
|