auracaster/piper
3
0
Fork 0
mirror of https://github.com/pstrueb/piper.git synced 2026-04-19 14:54:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Clean up export onnx

Browse Source
This commit is contained in:
Michael Hansen
2022-11-11 16:59:57 -05:00
parent 042718a481
commit 11b294a461
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/python/larynx_train/export_onnx.py
View File

@@ -8,16 +8,16 @@ import torch
from .vits.lightning import VitsModel
_LOGGER = logging.getLogger("mimic3_train.export_onnx")
_LOGGER = logging.getLogger("larynx_train.export_onnx")
OPSET_VERSION = 15
def main():
"""Main entry point"""
torch.manual_seed(12345)
torch.manual_seed(1234)
parser = argparse.ArgumentParser(prog="mimic3_train.export_onnx")
parser = argparse.ArgumentParser()
parser.add_argument("checkpoint", help="Path to model checkpoint (.ckpt)")
parser.add_argument("output", help="Path to output model (.onnx)")
@@ -70,7 +70,10 @@ def main():
model_g.forward = infer_forward
sequences = torch.randint(low=0, high=num_symbols, size=(1, 50), dtype=torch.long)
dummy_input_length = 50
sequences = torch.randint(
low=0, high=num_symbols, size=(1, dummy_input_length), dtype=torch.long
)
sequence_lengths = torch.LongTensor([sequences.size(1)])
sid: Optional[int] = None
@@ -79,7 +82,6 @@ def main():
# noise, noise_w, length
scales = torch.FloatTensor([0.667, 1.0, 0.8])
dummy_input = (sequences, sequence_lengths, scales, sid)
# Export