mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-10 10:32:17 +00:00
remove baker-zh
This commit is contained in:
Fangjun Kuang
parent
f9bd5ced9d
commit
35578f0593
@ -1,7 +0,0 @@
|
|||||||
# Introduction
|
|
||||||
|
|
||||||
[./symbols.py](./symbols.py) is copied from
|
|
||||||
https://github.com/UEhQZXI/vits_chinese/blob/master/text/symbols.py
|
|
||||||
|
|
||||||
[./pypinyin-local.dict](./pypinyin-local.dict) is copied from
|
|
||||||
https://github.com/UEhQZXI/vits_chinese/blob/master/misc/pypinyin-local.dict
|
|
||||||
@ -1,106 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
# Copyright 2021-2023 Xiaomi Corp. (authors: Fangjun Kuang,
|
|
||||||
# Zengwei Yao)
|
|
||||||
#
|
|
||||||
# See ../../../../LICENSE for clarification regarding multiple authors
|
|
||||||
#
|
|
||||||
# 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.
|
|
||||||
|
|
||||||
|
|
||||||
"""
|
|
||||||
This file computes fbank features of the baker_zh dataset.
|
|
||||||
It looks for manifests in the directory data/manifests.
|
|
||||||
|
|
||||||
The generated spectrogram features are saved in data/spectrogram.
|
|
||||||
"""
|
|
||||||
|
|
||||||
import logging
|
|
||||||
import os
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from lhotse import (
|
|
||||||
CutSet,
|
|
||||||
LilcomChunkyWriter,
|
|
||||||
Spectrogram,
|
|
||||||
SpectrogramConfig,
|
|
||||||
load_manifest,
|
|
||||||
)
|
|
||||||
from lhotse.audio import RecordingSet
|
|
||||||
from lhotse.supervision import SupervisionSet
|
|
||||||
|
|
||||||
from icefall.utils import get_executor
|
|
||||||
|
|
||||||
# Torch's multithreaded behavior needs to be disabled or
|
|
||||||
# it wastes a lot of CPU and slow things down.
|
|
||||||
# Do this outside of main() in case it needs to take effect
|
|
||||||
# even when we are not invoking the main (e.g. when spawning subprocesses).
|
|
||||||
torch.set_num_threads(1)
|
|
||||||
torch.set_num_interop_threads(1)
|
|
||||||
|
|
||||||
|
|
||||||
def compute_spectrogram_baker_zh():
|
|
||||||
src_dir = Path("data/manifests")
|
|
||||||
output_dir = Path("data/spectrogram")
|
|
||||||
num_jobs = min(4, os.cpu_count())
|
|
||||||
|
|
||||||
sampling_rate = 48000
|
|
||||||
frame_length = 1024 / sampling_rate # (in second)
|
|
||||||
frame_shift = 256 / sampling_rate # (in second)
|
|
||||||
use_fft_mag = True
|
|
||||||
|
|
||||||
prefix = "baker_zh"
|
|
||||||
suffix = "jsonl.gz"
|
|
||||||
partition = "all"
|
|
||||||
|
|
||||||
recordings = load_manifest(
|
|
||||||
src_dir / f"{prefix}_recordings_{partition}.{suffix}", RecordingSet
|
|
||||||
)
|
|
||||||
supervisions = load_manifest(
|
|
||||||
src_dir / f"{prefix}_supervisions_{partition}.{suffix}", SupervisionSet
|
|
||||||
)
|
|
||||||
|
|
||||||
config = SpectrogramConfig(
|
|
||||||
sampling_rate=sampling_rate,
|
|
||||||
frame_length=frame_length,
|
|
||||||
frame_shift=frame_shift,
|
|
||||||
use_fft_mag=use_fft_mag,
|
|
||||||
)
|
|
||||||
extractor = Spectrogram(config)
|
|
||||||
|
|
||||||
with get_executor() as ex: # Initialize the executor only once.
|
|
||||||
cuts_filename = f"{prefix}_cuts_{partition}.{suffix}"
|
|
||||||
if (output_dir / cuts_filename).is_file():
|
|
||||||
logging.info(f"{cuts_filename} already exists - skipping.")
|
|
||||||
return
|
|
||||||
logging.info(f"Processing {partition}")
|
|
||||||
cut_set = CutSet.from_manifests(
|
|
||||||
recordings=recordings, supervisions=supervisions
|
|
||||||
)
|
|
||||||
|
|
||||||
cut_set = cut_set.compute_and_store_features(
|
|
||||||
extractor=extractor,
|
|
||||||
storage_path=f"{output_dir}/{prefix}_feats_{partition}",
|
|
||||||
# when an executor is specified, make more partitions
|
|
||||||
num_jobs=num_jobs if ex is None else 80,
|
|
||||||
executor=ex,
|
|
||||||
storage_type=LilcomChunkyWriter,
|
|
||||||
)
|
|
||||||
cut_set.to_file(output_dir / cuts_filename)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
|
|
||||||
|
|
||||||
logging.basicConfig(format=formatter, level=logging.INFO)
|
|
||||||
compute_spectrogram_baker_zh()
|
|
||||||
@ -1,421 +0,0 @@
|
|||||||
# This dict is copied from
|
|
||||||
# https://github.com/UEhQZXI/vits_chinese/blob/master/vits_strings.py
|
|
||||||
pinyin_dict = {
|
|
||||||
"a": ("^", "a"),
|
|
||||||
"ai": ("^", "ai"),
|
|
||||||
"an": ("^", "an"),
|
|
||||||
"ang": ("^", "ang"),
|
|
||||||
"ao": ("^", "ao"),
|
|
||||||
"ba": ("b", "a"),
|
|
||||||
"bai": ("b", "ai"),
|
|
||||||
"ban": ("b", "an"),
|
|
||||||
"bang": ("b", "ang"),
|
|
||||||
"bao": ("b", "ao"),
|
|
||||||
"be": ("b", "e"),
|
|
||||||
"bei": ("b", "ei"),
|
|
||||||
"ben": ("b", "en"),
|
|
||||||
"beng": ("b", "eng"),
|
|
||||||
"bi": ("b", "i"),
|
|
||||||
"bian": ("b", "ian"),
|
|
||||||
"biao": ("b", "iao"),
|
|
||||||
"bie": ("b", "ie"),
|
|
||||||
"bin": ("b", "in"),
|
|
||||||
"bing": ("b", "ing"),
|
|
||||||
"bo": ("b", "o"),
|
|
||||||
"bu": ("b", "u"),
|
|
||||||
"ca": ("c", "a"),
|
|
||||||
"cai": ("c", "ai"),
|
|
||||||
"can": ("c", "an"),
|
|
||||||
"cang": ("c", "ang"),
|
|
||||||
"cao": ("c", "ao"),
|
|
||||||
"ce": ("c", "e"),
|
|
||||||
"cen": ("c", "en"),
|
|
||||||
"ceng": ("c", "eng"),
|
|
||||||
"cha": ("ch", "a"),
|
|
||||||
"chai": ("ch", "ai"),
|
|
||||||
"chan": ("ch", "an"),
|
|
||||||
"chang": ("ch", "ang"),
|
|
||||||
"chao": ("ch", "ao"),
|
|
||||||
"che": ("ch", "e"),
|
|
||||||
"chen": ("ch", "en"),
|
|
||||||
"cheng": ("ch", "eng"),
|
|
||||||
"chi": ("ch", "iii"),
|
|
||||||
"chong": ("ch", "ong"),
|
|
||||||
"chou": ("ch", "ou"),
|
|
||||||
"chu": ("ch", "u"),
|
|
||||||
"chua": ("ch", "ua"),
|
|
||||||
"chuai": ("ch", "uai"),
|
|
||||||
"chuan": ("ch", "uan"),
|
|
||||||
"chuang": ("ch", "uang"),
|
|
||||||
"chui": ("ch", "uei"),
|
|
||||||
"chun": ("ch", "uen"),
|
|
||||||
"chuo": ("ch", "uo"),
|
|
||||||
"ci": ("c", "ii"),
|
|
||||||
"cong": ("c", "ong"),
|
|
||||||
"cou": ("c", "ou"),
|
|
||||||
"cu": ("c", "u"),
|
|
||||||
"cuan": ("c", "uan"),
|
|
||||||
"cui": ("c", "uei"),
|
|
||||||
"cun": ("c", "uen"),
|
|
||||||
"cuo": ("c", "uo"),
|
|
||||||
"da": ("d", "a"),
|
|
||||||
"dai": ("d", "ai"),
|
|
||||||
"dan": ("d", "an"),
|
|
||||||
"dang": ("d", "ang"),
|
|
||||||
"dao": ("d", "ao"),
|
|
||||||
"de": ("d", "e"),
|
|
||||||
"dei": ("d", "ei"),
|
|
||||||
"den": ("d", "en"),
|
|
||||||
"deng": ("d", "eng"),
|
|
||||||
"di": ("d", "i"),
|
|
||||||
"dia": ("d", "ia"),
|
|
||||||
"dian": ("d", "ian"),
|
|
||||||
"diao": ("d", "iao"),
|
|
||||||
"die": ("d", "ie"),
|
|
||||||
"ding": ("d", "ing"),
|
|
||||||
"diu": ("d", "iou"),
|
|
||||||
"dong": ("d", "ong"),
|
|
||||||
"dou": ("d", "ou"),
|
|
||||||
"du": ("d", "u"),
|
|
||||||
"duan": ("d", "uan"),
|
|
||||||
"dui": ("d", "uei"),
|
|
||||||
"dun": ("d", "uen"),
|
|
||||||
"duo": ("d", "uo"),
|
|
||||||
"e": ("^", "e"),
|
|
||||||
"ei": ("^", "ei"),
|
|
||||||
"en": ("^", "en"),
|
|
||||||
"ng": ("^", "en"),
|
|
||||||
"eng": ("^", "eng"),
|
|
||||||
"er": ("^", "er"),
|
|
||||||
"fa": ("f", "a"),
|
|
||||||
"fan": ("f", "an"),
|
|
||||||
"fang": ("f", "ang"),
|
|
||||||
"fei": ("f", "ei"),
|
|
||||||
"fen": ("f", "en"),
|
|
||||||
"feng": ("f", "eng"),
|
|
||||||
"fo": ("f", "o"),
|
|
||||||
"fou": ("f", "ou"),
|
|
||||||
"fu": ("f", "u"),
|
|
||||||
"ga": ("g", "a"),
|
|
||||||
"gai": ("g", "ai"),
|
|
||||||
"gan": ("g", "an"),
|
|
||||||
"gang": ("g", "ang"),
|
|
||||||
"gao": ("g", "ao"),
|
|
||||||
"ge": ("g", "e"),
|
|
||||||
"gei": ("g", "ei"),
|
|
||||||
"gen": ("g", "en"),
|
|
||||||
"geng": ("g", "eng"),
|
|
||||||
"gong": ("g", "ong"),
|
|
||||||
"gou": ("g", "ou"),
|
|
||||||
"gu": ("g", "u"),
|
|
||||||
"gua": ("g", "ua"),
|
|
||||||
"guai": ("g", "uai"),
|
|
||||||
"guan": ("g", "uan"),
|
|
||||||
"guang": ("g", "uang"),
|
|
||||||
"gui": ("g", "uei"),
|
|
||||||
"gun": ("g", "uen"),
|
|
||||||
"guo": ("g", "uo"),
|
|
||||||
"ha": ("h", "a"),
|
|
||||||
"hai": ("h", "ai"),
|
|
||||||
"han": ("h", "an"),
|
|
||||||
"hang": ("h", "ang"),
|
|
||||||
"hao": ("h", "ao"),
|
|
||||||
"he": ("h", "e"),
|
|
||||||
"hei": ("h", "ei"),
|
|
||||||
"hen": ("h", "en"),
|
|
||||||
"heng": ("h", "eng"),
|
|
||||||
"hong": ("h", "ong"),
|
|
||||||
"hou": ("h", "ou"),
|
|
||||||
"hu": ("h", "u"),
|
|
||||||
"hua": ("h", "ua"),
|
|
||||||
"huai": ("h", "uai"),
|
|
||||||
"huan": ("h", "uan"),
|
|
||||||
"huang": ("h", "uang"),
|
|
||||||
"hui": ("h", "uei"),
|
|
||||||
"hun": ("h", "uen"),
|
|
||||||
"huo": ("h", "uo"),
|
|
||||||
"ji": ("j", "i"),
|
|
||||||
"jia": ("j", "ia"),
|
|
||||||
"jian": ("j", "ian"),
|
|
||||||
"jiang": ("j", "iang"),
|
|
||||||
"jiao": ("j", "iao"),
|
|
||||||
"jie": ("j", "ie"),
|
|
||||||
"jin": ("j", "in"),
|
|
||||||
"jing": ("j", "ing"),
|
|
||||||
"jiong": ("j", "iong"),
|
|
||||||
"jiu": ("j", "iou"),
|
|
||||||
"ju": ("j", "v"),
|
|
||||||
"juan": ("j", "van"),
|
|
||||||
"jue": ("j", "ve"),
|
|
||||||
"jun": ("j", "vn"),
|
|
||||||
"ka": ("k", "a"),
|
|
||||||
"kai": ("k", "ai"),
|
|
||||||
"kan": ("k", "an"),
|
|
||||||
"kang": ("k", "ang"),
|
|
||||||
"kao": ("k", "ao"),
|
|
||||||
"ke": ("k", "e"),
|
|
||||||
"kei": ("k", "ei"),
|
|
||||||
"ken": ("k", "en"),
|
|
||||||
"keng": ("k", "eng"),
|
|
||||||
"kong": ("k", "ong"),
|
|
||||||
"kou": ("k", "ou"),
|
|
||||||
"ku": ("k", "u"),
|
|
||||||
"kua": ("k", "ua"),
|
|
||||||
"kuai": ("k", "uai"),
|
|
||||||
"kuan": ("k", "uan"),
|
|
||||||
"kuang": ("k", "uang"),
|
|
||||||
"kui": ("k", "uei"),
|
|
||||||
"kun": ("k", "uen"),
|
|
||||||
"kuo": ("k", "uo"),
|
|
||||||
"la": ("l", "a"),
|
|
||||||
"lai": ("l", "ai"),
|
|
||||||
"lan": ("l", "an"),
|
|
||||||
"lang": ("l", "ang"),
|
|
||||||
"lao": ("l", "ao"),
|
|
||||||
"le": ("l", "e"),
|
|
||||||
"lei": ("l", "ei"),
|
|
||||||
"leng": ("l", "eng"),
|
|
||||||
"li": ("l", "i"),
|
|
||||||
"lia": ("l", "ia"),
|
|
||||||
"lian": ("l", "ian"),
|
|
||||||
"liang": ("l", "iang"),
|
|
||||||
"liao": ("l", "iao"),
|
|
||||||
"lie": ("l", "ie"),
|
|
||||||
"lin": ("l", "in"),
|
|
||||||
"ling": ("l", "ing"),
|
|
||||||
"liu": ("l", "iou"),
|
|
||||||
"lo": ("l", "o"),
|
|
||||||
"long": ("l", "ong"),
|
|
||||||
"lou": ("l", "ou"),
|
|
||||||
"lu": ("l", "u"),
|
|
||||||
"lv": ("l", "v"),
|
|
||||||
"luan": ("l", "uan"),
|
|
||||||
"lve": ("l", "ve"),
|
|
||||||
"lue": ("l", "ve"),
|
|
||||||
"lun": ("l", "uen"),
|
|
||||||
"luo": ("l", "uo"),
|
|
||||||
"ma": ("m", "a"),
|
|
||||||
"mai": ("m", "ai"),
|
|
||||||
"man": ("m", "an"),
|
|
||||||
"mang": ("m", "ang"),
|
|
||||||
"mao": ("m", "ao"),
|
|
||||||
"me": ("m", "e"),
|
|
||||||
"mei": ("m", "ei"),
|
|
||||||
"men": ("m", "en"),
|
|
||||||
"meng": ("m", "eng"),
|
|
||||||
"mi": ("m", "i"),
|
|
||||||
"mian": ("m", "ian"),
|
|
||||||
"miao": ("m", "iao"),
|
|
||||||
"mie": ("m", "ie"),
|
|
||||||
"min": ("m", "in"),
|
|
||||||
"ming": ("m", "ing"),
|
|
||||||
"miu": ("m", "iou"),
|
|
||||||
"mo": ("m", "o"),
|
|
||||||
"mou": ("m", "ou"),
|
|
||||||
"mu": ("m", "u"),
|
|
||||||
"na": ("n", "a"),
|
|
||||||
"nai": ("n", "ai"),
|
|
||||||
"nan": ("n", "an"),
|
|
||||||
"nang": ("n", "ang"),
|
|
||||||
"nao": ("n", "ao"),
|
|
||||||
"ne": ("n", "e"),
|
|
||||||
"nei": ("n", "ei"),
|
|
||||||
"nen": ("n", "en"),
|
|
||||||
"neng": ("n", "eng"),
|
|
||||||
"ni": ("n", "i"),
|
|
||||||
"nia": ("n", "ia"),
|
|
||||||
"nian": ("n", "ian"),
|
|
||||||
"niang": ("n", "iang"),
|
|
||||||
"niao": ("n", "iao"),
|
|
||||||
"nie": ("n", "ie"),
|
|
||||||
"nin": ("n", "in"),
|
|
||||||
"ning": ("n", "ing"),
|
|
||||||
"niu": ("n", "iou"),
|
|
||||||
"nong": ("n", "ong"),
|
|
||||||
"nou": ("n", "ou"),
|
|
||||||
"nu": ("n", "u"),
|
|
||||||
"nv": ("n", "v"),
|
|
||||||
"nuan": ("n", "uan"),
|
|
||||||
"nve": ("n", "ve"),
|
|
||||||
"nue": ("n", "ve"),
|
|
||||||
"nuo": ("n", "uo"),
|
|
||||||
"o": ("^", "o"),
|
|
||||||
"ou": ("^", "ou"),
|
|
||||||
"pa": ("p", "a"),
|
|
||||||
"pai": ("p", "ai"),
|
|
||||||
"pan": ("p", "an"),
|
|
||||||
"pang": ("p", "ang"),
|
|
||||||
"pao": ("p", "ao"),
|
|
||||||
"pe": ("p", "e"),
|
|
||||||
"pei": ("p", "ei"),
|
|
||||||
"pen": ("p", "en"),
|
|
||||||
"peng": ("p", "eng"),
|
|
||||||
"pi": ("p", "i"),
|
|
||||||
"pian": ("p", "ian"),
|
|
||||||
"piao": ("p", "iao"),
|
|
||||||
"pie": ("p", "ie"),
|
|
||||||
"pin": ("p", "in"),
|
|
||||||
"ping": ("p", "ing"),
|
|
||||||
"po": ("p", "o"),
|
|
||||||
"pou": ("p", "ou"),
|
|
||||||
"pu": ("p", "u"),
|
|
||||||
"qi": ("q", "i"),
|
|
||||||
"qia": ("q", "ia"),
|
|
||||||
"qian": ("q", "ian"),
|
|
||||||
"qiang": ("q", "iang"),
|
|
||||||
"qiao": ("q", "iao"),
|
|
||||||
"qie": ("q", "ie"),
|
|
||||||
"qin": ("q", "in"),
|
|
||||||
"qing": ("q", "ing"),
|
|
||||||
"qiong": ("q", "iong"),
|
|
||||||
"qiu": ("q", "iou"),
|
|
||||||
"qu": ("q", "v"),
|
|
||||||
"quan": ("q", "van"),
|
|
||||||
"que": ("q", "ve"),
|
|
||||||
"qun": ("q", "vn"),
|
|
||||||
"ran": ("r", "an"),
|
|
||||||
"rang": ("r", "ang"),
|
|
||||||
"rao": ("r", "ao"),
|
|
||||||
"re": ("r", "e"),
|
|
||||||
"ren": ("r", "en"),
|
|
||||||
"reng": ("r", "eng"),
|
|
||||||
"ri": ("r", "iii"),
|
|
||||||
"rong": ("r", "ong"),
|
|
||||||
"rou": ("r", "ou"),
|
|
||||||
"ru": ("r", "u"),
|
|
||||||
"rua": ("r", "ua"),
|
|
||||||
"ruan": ("r", "uan"),
|
|
||||||
"rui": ("r", "uei"),
|
|
||||||
"run": ("r", "uen"),
|
|
||||||
"ruo": ("r", "uo"),
|
|
||||||
"sa": ("s", "a"),
|
|
||||||
"sai": ("s", "ai"),
|
|
||||||
"san": ("s", "an"),
|
|
||||||
"sang": ("s", "ang"),
|
|
||||||
"sao": ("s", "ao"),
|
|
||||||
"se": ("s", "e"),
|
|
||||||
"sen": ("s", "en"),
|
|
||||||
"seng": ("s", "eng"),
|
|
||||||
"sha": ("sh", "a"),
|
|
||||||
"shai": ("sh", "ai"),
|
|
||||||
"shan": ("sh", "an"),
|
|
||||||
"shang": ("sh", "ang"),
|
|
||||||
"shao": ("sh", "ao"),
|
|
||||||
"she": ("sh", "e"),
|
|
||||||
"shei": ("sh", "ei"),
|
|
||||||
"shen": ("sh", "en"),
|
|
||||||
"sheng": ("sh", "eng"),
|
|
||||||
"shi": ("sh", "iii"),
|
|
||||||
"shou": ("sh", "ou"),
|
|
||||||
"shu": ("sh", "u"),
|
|
||||||
"shua": ("sh", "ua"),
|
|
||||||
"shuai": ("sh", "uai"),
|
|
||||||
"shuan": ("sh", "uan"),
|
|
||||||
"shuang": ("sh", "uang"),
|
|
||||||
"shui": ("sh", "uei"),
|
|
||||||
"shun": ("sh", "uen"),
|
|
||||||
"shuo": ("sh", "uo"),
|
|
||||||
"si": ("s", "ii"),
|
|
||||||
"song": ("s", "ong"),
|
|
||||||
"sou": ("s", "ou"),
|
|
||||||
"su": ("s", "u"),
|
|
||||||
"suan": ("s", "uan"),
|
|
||||||
"sui": ("s", "uei"),
|
|
||||||
"sun": ("s", "uen"),
|
|
||||||
"suo": ("s", "uo"),
|
|
||||||
"ta": ("t", "a"),
|
|
||||||
"tai": ("t", "ai"),
|
|
||||||
"tan": ("t", "an"),
|
|
||||||
"tang": ("t", "ang"),
|
|
||||||
"tao": ("t", "ao"),
|
|
||||||
"te": ("t", "e"),
|
|
||||||
"tei": ("t", "ei"),
|
|
||||||
"teng": ("t", "eng"),
|
|
||||||
"ti": ("t", "i"),
|
|
||||||
"tian": ("t", "ian"),
|
|
||||||
"tiao": ("t", "iao"),
|
|
||||||
"tie": ("t", "ie"),
|
|
||||||
"ting": ("t", "ing"),
|
|
||||||
"tong": ("t", "ong"),
|
|
||||||
"tou": ("t", "ou"),
|
|
||||||
"tu": ("t", "u"),
|
|
||||||
"tuan": ("t", "uan"),
|
|
||||||
"tui": ("t", "uei"),
|
|
||||||
"tun": ("t", "uen"),
|
|
||||||
"tuo": ("t", "uo"),
|
|
||||||
"wa": ("^", "ua"),
|
|
||||||
"wai": ("^", "uai"),
|
|
||||||
"wan": ("^", "uan"),
|
|
||||||
"wang": ("^", "uang"),
|
|
||||||
"wei": ("^", "uei"),
|
|
||||||
"wen": ("^", "uen"),
|
|
||||||
"weng": ("^", "ueng"),
|
|
||||||
"wo": ("^", "uo"),
|
|
||||||
"wu": ("^", "u"),
|
|
||||||
"xi": ("x", "i"),
|
|
||||||
"xia": ("x", "ia"),
|
|
||||||
"xian": ("x", "ian"),
|
|
||||||
"xiang": ("x", "iang"),
|
|
||||||
"xiao": ("x", "iao"),
|
|
||||||
"xie": ("x", "ie"),
|
|
||||||
"xin": ("x", "in"),
|
|
||||||
"xing": ("x", "ing"),
|
|
||||||
"xiong": ("x", "iong"),
|
|
||||||
"xiu": ("x", "iou"),
|
|
||||||
"xu": ("x", "v"),
|
|
||||||
"xuan": ("x", "van"),
|
|
||||||
"xue": ("x", "ve"),
|
|
||||||
"xun": ("x", "vn"),
|
|
||||||
"ya": ("^", "ia"),
|
|
||||||
"yan": ("^", "ian"),
|
|
||||||
"yang": ("^", "iang"),
|
|
||||||
"yao": ("^", "iao"),
|
|
||||||
"ye": ("^", "ie"),
|
|
||||||
"yi": ("^", "i"),
|
|
||||||
"yin": ("^", "in"),
|
|
||||||
"ying": ("^", "ing"),
|
|
||||||
"yo": ("^", "iou"),
|
|
||||||
"yong": ("^", "iong"),
|
|
||||||
"you": ("^", "iou"),
|
|
||||||
"yu": ("^", "v"),
|
|
||||||
"yuan": ("^", "van"),
|
|
||||||
"yue": ("^", "ve"),
|
|
||||||
"yun": ("^", "vn"),
|
|
||||||
"za": ("z", "a"),
|
|
||||||
"zai": ("z", "ai"),
|
|
||||||
"zan": ("z", "an"),
|
|
||||||
"zang": ("z", "ang"),
|
|
||||||
"zao": ("z", "ao"),
|
|
||||||
"ze": ("z", "e"),
|
|
||||||
"zei": ("z", "ei"),
|
|
||||||
"zen": ("z", "en"),
|
|
||||||
"zeng": ("z", "eng"),
|
|
||||||
"zha": ("zh", "a"),
|
|
||||||
"zhai": ("zh", "ai"),
|
|
||||||
"zhan": ("zh", "an"),
|
|
||||||
"zhang": ("zh", "ang"),
|
|
||||||
"zhao": ("zh", "ao"),
|
|
||||||
"zhe": ("zh", "e"),
|
|
||||||
"zhei": ("zh", "ei"),
|
|
||||||
"zhen": ("zh", "en"),
|
|
||||||
"zheng": ("zh", "eng"),
|
|
||||||
"zhi": ("zh", "iii"),
|
|
||||||
"zhong": ("zh", "ong"),
|
|
||||||
"zhou": ("zh", "ou"),
|
|
||||||
"zhu": ("zh", "u"),
|
|
||||||
"zhua": ("zh", "ua"),
|
|
||||||
"zhuai": ("zh", "uai"),
|
|
||||||
"zhuan": ("zh", "uan"),
|
|
||||||
"zhuang": ("zh", "uang"),
|
|
||||||
"zhui": ("zh", "uei"),
|
|
||||||
"zhun": ("zh", "uen"),
|
|
||||||
"zhuo": ("zh", "uo"),
|
|
||||||
"zi": ("z", "ii"),
|
|
||||||
"zong": ("z", "ong"),
|
|
||||||
"zou": ("z", "ou"),
|
|
||||||
"zu": ("z", "u"),
|
|
||||||
"zuan": ("z", "uan"),
|
|
||||||
"zui": ("z", "uei"),
|
|
||||||
"zun": ("z", "uen"),
|
|
||||||
"zuo": ("z", "uo"),
|
|
||||||
}
|
|
||||||
@ -1,53 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
# Copyright 2023 Xiaomi Corp. (authors: Zengwei Yao)
|
|
||||||
#
|
|
||||||
# See ../../../../LICENSE for clarification regarding multiple authors
|
|
||||||
#
|
|
||||||
# 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.
|
|
||||||
|
|
||||||
|
|
||||||
"""
|
|
||||||
This file generates the file that maps tokens to IDs.
|
|
||||||
"""
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import logging
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Dict
|
|
||||||
from symbols import symbols
|
|
||||||
|
|
||||||
|
|
||||||
def get_args():
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--tokens",
|
|
||||||
type=Path,
|
|
||||||
default=Path("data/tokens.txt"),
|
|
||||||
help="Path to the dict that maps the text tokens to IDs",
|
|
||||||
)
|
|
||||||
|
|
||||||
return parser.parse_args()
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
args = get_args()
|
|
||||||
tokens = Path(args.tokens)
|
|
||||||
|
|
||||||
with open(tokens, "w", encoding="utf-8") as f:
|
|
||||||
for token_id, token in enumerate(symbols):
|
|
||||||
f.write(f"{token} {token_id}\n")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@ -1,59 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
# Copyright 2023 Xiaomi Corp. (authors: Zengwei Yao)
|
|
||||||
#
|
|
||||||
# See ../../../../LICENSE for clarification regarding multiple authors
|
|
||||||
#
|
|
||||||
# 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.
|
|
||||||
|
|
||||||
|
|
||||||
"""
|
|
||||||
This file reads the texts in given manifest and save the new cuts with tokens.
|
|
||||||
"""
|
|
||||||
|
|
||||||
import logging
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
from lhotse import CutSet, load_manifest
|
|
||||||
|
|
||||||
from tokenizer import Tokenizer
|
|
||||||
|
|
||||||
|
|
||||||
def prepare_tokens_baker_zh():
|
|
||||||
output_dir = Path("data/spectrogram")
|
|
||||||
prefix = "baker_zh"
|
|
||||||
suffix = "jsonl.gz"
|
|
||||||
partition = "all"
|
|
||||||
|
|
||||||
cut_set = load_manifest(output_dir / f"{prefix}_cuts_{partition}.{suffix}")
|
|
||||||
|
|
||||||
tokenizer = Tokenizer()
|
|
||||||
|
|
||||||
new_cuts = []
|
|
||||||
i = 0
|
|
||||||
for cut in cut_set:
|
|
||||||
# Each cut only contains one supervision
|
|
||||||
assert len(cut.supervisions) == 1, (len(cut.supervisions), cut)
|
|
||||||
text = cut.supervisions[0].normalized_text
|
|
||||||
cut.tokens = tokenizer.text_to_tokens(text)
|
|
||||||
|
|
||||||
new_cuts.append(cut)
|
|
||||||
|
|
||||||
new_cut_set = CutSet.from_cuts(new_cuts)
|
|
||||||
new_cut_set.to_file(output_dir / f"{prefix}_cuts_with_tokens_{partition}.{suffix}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
|
|
||||||
logging.basicConfig(format=formatter, level=logging.INFO)
|
|
||||||
|
|
||||||
prepare_tokens_baker_zh()
|
|
||||||
@ -1,328 +0,0 @@
|
|||||||
姐姐 jie3 jie
|
|
||||||
宝宝 bao3 bao
|
|
||||||
哥哥 ge1 ge
|
|
||||||
妹妹 mei4 mei
|
|
||||||
弟弟 di4 di
|
|
||||||
妈妈 ma1 ma
|
|
||||||
开心哦 kai1 xin1 o
|
|
||||||
爸爸 ba4 ba
|
|
||||||
秘密哟 mi4 mi4 yo
|
|
||||||
哦 o
|
|
||||||
一年 yi4 nian2
|
|
||||||
一夜 yi2 ye4
|
|
||||||
一切 yi2 qie4
|
|
||||||
一座 yi2 zuo4
|
|
||||||
一下 yi2 xia4
|
|
||||||
上一山 shang4 yi2 shan1
|
|
||||||
下一山 xia4 yi2 shan1
|
|
||||||
休息 xiu1 xi2
|
|
||||||
东西 dong1 xi
|
|
||||||
上一届 shang4 yi2 jie4
|
|
||||||
便宜 pian2 yi4
|
|
||||||
加长 jia1 chang2
|
|
||||||
单田芳 shan4 tian2 fang1
|
|
||||||
帧 zhen1
|
|
||||||
长时间 chang2 shi2 jian1
|
|
||||||
长时 chang2 shi2
|
|
||||||
识别 shi2 bie2
|
|
||||||
生命中 sheng1 ming4 zhong1
|
|
||||||
踏实 ta1 shi
|
|
||||||
嗯 en4
|
|
||||||
溜达 liu1 da
|
|
||||||
少儿 shao4 er2
|
|
||||||
爷爷 ye2 ye
|
|
||||||
不是 bu2 shi4
|
|
||||||
一圈 yi1 quan1
|
|
||||||
厜读一声 zui1 du2 yi4 sheng1
|
|
||||||
一种 yi4 zhong3
|
|
||||||
一簇簇 yi2 cu4 cu4
|
|
||||||
一个 yi2 ge4
|
|
||||||
一样 yi2 yang4
|
|
||||||
一跩一跩 yi4 zhuai3 yi4 zhuai3
|
|
||||||
一会儿 yi2 hui4 er
|
|
||||||
一幢 yi2 zhuang4
|
|
||||||
挨了 ai2 le
|
|
||||||
熬菜 ao1 cai4
|
|
||||||
扒鸡 pa2 ji1
|
|
||||||
背枪 bei1 qiang1
|
|
||||||
绷瓷儿 beng4 ci2 er2
|
|
||||||
绷劲儿 beng3 jin4 er
|
|
||||||
绷着脸 beng3 zhe lian3
|
|
||||||
藏医 zang4 yi1
|
|
||||||
噌吰 cheng1 hong2
|
|
||||||
差点儿 cha4 dian3 er
|
|
||||||
差失 cha1 shi1
|
|
||||||
差误 cha1 wu4
|
|
||||||
孱头 can4 tou
|
|
||||||
乘间 cheng2 jian4
|
|
||||||
锄镰棘矜 chu2 lian2 ji2 qin2
|
|
||||||
川藏 chuan1 zang4
|
|
||||||
穿著 chuan1 zhuo2
|
|
||||||
答讪 da1 shan4
|
|
||||||
答言 da1 yan2
|
|
||||||
大伯子 da4 bai3 zi
|
|
||||||
大夫 dai4 fu
|
|
||||||
弹冠 tan2 guan1
|
|
||||||
当间 dang1 jian4
|
|
||||||
当然咯 dang1 ran2 lo
|
|
||||||
点种 dian3 zhong3
|
|
||||||
垛好 duo4 hao3
|
|
||||||
发疟子 fa1 yao4 zi
|
|
||||||
饭熟了 fan4 shou2 le
|
|
||||||
附著 fu4 zhuo2
|
|
||||||
复沓 fu4 ta4
|
|
||||||
供稿 gong1 gao3
|
|
||||||
供养 gong1 yang3
|
|
||||||
骨朵 gu1 duo
|
|
||||||
骨碌 gu1 lu
|
|
||||||
果脯 guo3 fu3
|
|
||||||
哈什玛 ha4 shi2 ma3
|
|
||||||
海蜇 hai3 zhe2
|
|
||||||
呵欠 he1 qian
|
|
||||||
河水汤汤 he2 shui3 shang1 shang1
|
|
||||||
鹄立 hu2 li4
|
|
||||||
鹄望 hu2 wang4
|
|
||||||
混人 hun2 ren2
|
|
||||||
混水 hun2 shui3
|
|
||||||
鸡血 ji1 xie3
|
|
||||||
缉鞋口 qi1 xie2 kou3
|
|
||||||
亟来闻讯 qi4 lai2 wen2 xun4
|
|
||||||
计量 ji4 liang2
|
|
||||||
济水 ji3 shui3
|
|
||||||
间杂 jian4 za2
|
|
||||||
脚跐两只船 jiao3 ci3 liang3 zhi1 chuan2
|
|
||||||
脚儿 jue2 er2
|
|
||||||
口角 kou3 jiao3
|
|
||||||
勒石 le4 shi2
|
|
||||||
累进 lei3 jin4
|
|
||||||
累累如丧家之犬 lei2 lei2 ru2 sang4 jia1 zhi1 quan3
|
|
||||||
累年 lei3 nian2
|
|
||||||
脸涨通红 lian3 zhang4 tong1 hong2
|
|
||||||
踉锵 liang4 qiang1
|
|
||||||
燎眉毛 liao3 mei2 mao2
|
|
||||||
燎头发 liao3 tou2 fa4
|
|
||||||
溜达 liu1 da
|
|
||||||
溜缝儿 liu4 feng4 er
|
|
||||||
馏口饭 liu4 kou3 fan4
|
|
||||||
遛马 liu4 ma3
|
|
||||||
遛鸟 liu4 niao3
|
|
||||||
遛弯儿 liu4 wan1 er
|
|
||||||
楼枪机 lou1 qiang1 ji1
|
|
||||||
搂钱 lou1 qian2
|
|
||||||
鹿脯 lu4 fu3
|
|
||||||
露头 lou4 tou2
|
|
||||||
落魄 luo4 po4
|
|
||||||
捋胡子 lv3 hu2 zi
|
|
||||||
绿地 lv4 di4
|
|
||||||
麦垛 mai4 duo4
|
|
||||||
没劲儿 mei2 jin4 er
|
|
||||||
闷棍 men4 gun4
|
|
||||||
闷葫芦 men4 hu2 lu
|
|
||||||
闷头干 men1 tou2 gan4
|
|
||||||
蒙古 meng3 gu3
|
|
||||||
靡日不思 mi3 ri4 bu4 si1
|
|
||||||
缪姓 miao4 xing4
|
|
||||||
抹墙 mo4 qiang2
|
|
||||||
抹下脸 ma1 xia4 lian3
|
|
||||||
泥子 ni4 zi
|
|
||||||
拗不过 niu4 bu guo4
|
|
||||||
排车 pai3 che1
|
|
||||||
盘诘 pan2 jie2
|
|
||||||
膀肿 pang1 zhong3
|
|
||||||
炮干 bao1 gan1
|
|
||||||
炮格 pao2 ge2
|
|
||||||
碰钉子 peng4 ding1 zi
|
|
||||||
缥色 piao3 se4
|
|
||||||
瀑河 bao4 he2
|
|
||||||
蹊径 xi1 jing4
|
|
||||||
前后相属 qian2 hou4 xiang1 zhu3
|
|
||||||
翘尾巴 qiao4 wei3 ba
|
|
||||||
趄坡儿 qie4 po1 er
|
|
||||||
秦桧 qin2 hui4
|
|
||||||
圈马 juan1 ma3
|
|
||||||
雀盲眼 qiao3 mang2 yan3
|
|
||||||
雀子 qiao1 zi
|
|
||||||
三年五载 san1 nian2 wu3 zai3
|
|
||||||
加载 jia1 zai3
|
|
||||||
山大王 shan1 dai4 wang
|
|
||||||
苫屋草 shan4 wu1 cao3
|
|
||||||
数数 shu3 shu4
|
|
||||||
说客 shui4 ke4
|
|
||||||
思量 si1 liang2
|
|
||||||
伺侯 ci4 hou
|
|
||||||
踏实 ta1 shi
|
|
||||||
提溜 di1 liu
|
|
||||||
调拨 diao4 bo1
|
|
||||||
帖子 tie3 zi
|
|
||||||
铜钿 tong2 tian2
|
|
||||||
头昏脑涨 tou2 hun1 nao3 zhang4
|
|
||||||
褪色 tui4 se4
|
|
||||||
褪着手 tun4 zhe shou3
|
|
||||||
圩子 wei2 zi
|
|
||||||
尾巴 wei3 ba
|
|
||||||
系好船只 xi4 hao3 chuan2 zhi1
|
|
||||||
系好马匹 xi4 hao3 ma3 pi3
|
|
||||||
杏脯 xing4 fu3
|
|
||||||
姓单 xing4 shan4
|
|
||||||
姓葛 xing4 ge3
|
|
||||||
姓哈 xing4 ha3
|
|
||||||
姓解 xing4 xie4
|
|
||||||
姓秘 xing4 bi4
|
|
||||||
姓宁 xing4 ning4
|
|
||||||
旋风 xuan4 feng1
|
|
||||||
旋根车轴 xuan4 gen1 che1 zhou2
|
|
||||||
荨麻 qian2 ma2
|
|
||||||
一幢楼房 yi1 zhuang4 lou2 fang2
|
|
||||||
遗之千金 wei4 zhi1 qian1 jin1
|
|
||||||
殷殷 yin3 yin3
|
|
||||||
应招 ying4 zhao1
|
|
||||||
用称约 yong4 cheng4 yao1
|
|
||||||
约斤肉 yao1 jin1 rou4
|
|
||||||
晕机 yun4 ji1
|
|
||||||
熨贴 yu4 tie1
|
|
||||||
咋办 za3 ban4
|
|
||||||
咋呼 zha1 hu
|
|
||||||
仔兽 zi3 shou4
|
|
||||||
扎彩 za1 cai3
|
|
||||||
扎实 zha1 shi
|
|
||||||
扎腰带 za1 yao1 dai4
|
|
||||||
轧朋友 ga2 peng2 you3
|
|
||||||
爪子 zhua3 zi
|
|
||||||
折腾 zhe1 teng
|
|
||||||
着实 zhuo2 shi2
|
|
||||||
着我旧时裳 zhuo2 wo3 jiu4 shi2 chang2
|
|
||||||
枝蔓 zhi1 man4
|
|
||||||
中鹄 zhong1 hu2
|
|
||||||
中选 zhong4 xuan3
|
|
||||||
猪圈 zhu1 juan4
|
|
||||||
拽住不放 zhuai4 zhu4 bu4 fang4
|
|
||||||
转悠 zhuan4 you
|
|
||||||
庄稼熟了 zhuang1 jia shou2 le
|
|
||||||
酌量 zhuo2 liang2
|
|
||||||
罪行累累 zui4 xing2 lei3 lei3
|
|
||||||
一手 yi4 shou3
|
|
||||||
一去不复返 yi2 qu4 bu2 fu4 fan3
|
|
||||||
一颗 yi4 ke1
|
|
||||||
一件 yi2 jian4
|
|
||||||
一斤 yi4 jin1
|
|
||||||
一点 yi4 dian3
|
|
||||||
一朵 yi4 duo3
|
|
||||||
一声 yi4 sheng1
|
|
||||||
一身 yi4 shen1
|
|
||||||
不要 bu2 yao4
|
|
||||||
一人 yi4 ren2
|
|
||||||
一个 yi2 ge4
|
|
||||||
一把 yi4 ba3
|
|
||||||
一门 yi4 men2
|
|
||||||
一門 yi4 men2
|
|
||||||
一艘 yi4 sou1
|
|
||||||
一片 yi2 pian4
|
|
||||||
一篇 yi2 pian1
|
|
||||||
一份 yi2 fen4
|
|
||||||
好嗲 hao3 dia3
|
|
||||||
随地 sui2 di4
|
|
||||||
扁担长 bian3 dan4 chang3
|
|
||||||
一堆 yi4 dui1
|
|
||||||
不义 bu2 yi4
|
|
||||||
放一放 fang4 yi2 fang4
|
|
||||||
一米 yi4 mi3
|
|
||||||
一顿 yi2 dun4
|
|
||||||
一层楼 yi4 ceng2 lou2
|
|
||||||
一条 yi4 tiao2
|
|
||||||
一件 yi2 jian4
|
|
||||||
一棵 yi4 ke1
|
|
||||||
一小股 yi4 xiao3 gu3
|
|
||||||
一拐一拐 yi4 guai3 yi4 guai3
|
|
||||||
一根 yi4 gen1
|
|
||||||
沆瀣一气 hang4 xie4 yi2 qi4
|
|
||||||
一丝 yi4 si1
|
|
||||||
一毫 yi4 hao2
|
|
||||||
一樣 yi2 yang4
|
|
||||||
处处 chu4 chu4
|
|
||||||
一餐 yi4 can
|
|
||||||
永不 yong3 bu2
|
|
||||||
一看 yi2 kan4
|
|
||||||
一架 yi2 jia4
|
|
||||||
送还 song4 huan2
|
|
||||||
一见 yi2 jian4
|
|
||||||
一座 yi2 zuo4
|
|
||||||
一块 yi2 kuai4
|
|
||||||
一天 yi4 tian1
|
|
||||||
一只 yi4 zhi1
|
|
||||||
一支 yi4 zhi1
|
|
||||||
一字 yi2 zi4
|
|
||||||
一句 yi2 ju4
|
|
||||||
一张 yi4 zhang1
|
|
||||||
一條 yi4 tiao2
|
|
||||||
一场 yi4 chang3
|
|
||||||
一粒 yi2 li4
|
|
||||||
小俩口 xiao3 liang3 kou3
|
|
||||||
一首 yi4 shou3
|
|
||||||
一对 yi2 dui4
|
|
||||||
一手 yi4 shou3
|
|
||||||
又一村 you4 yi4 cun1
|
|
||||||
一概而论 yi2 gai4 er2 lun4
|
|
||||||
一峰峰 yi4 feng1 feng1
|
|
||||||
不但 bu2 dan4
|
|
||||||
一笑 yi2 xiao4
|
|
||||||
挠痒痒 nao2 yang3 yang
|
|
||||||
不对 bu2 dui4
|
|
||||||
拧开 ning3 kai1
|
|
||||||
爱不释手 ai4 bu2 shi4 shou3
|
|
||||||
一念 yi2 nian4
|
|
||||||
夺得 duo2 de2
|
|
||||||
一袭 yi4 xi2
|
|
||||||
一定 yi2 ding4
|
|
||||||
不慎 bu2 shen4
|
|
||||||
剽窃 piao2 qie4
|
|
||||||
一时 yi4 shi2
|
|
||||||
撇开 pie3 kai1
|
|
||||||
一祭 yi2 ji4
|
|
||||||
发卡 fa4 qia3
|
|
||||||
少不了 shao3 bu4 liao3
|
|
||||||
千虑一失 qian1 lv4 yi4 shi1
|
|
||||||
呛得 qiang4 de2
|
|
||||||
切菜 qie1 cai4
|
|
||||||
茄盒 qie2 he2
|
|
||||||
不去 bu2 qu4
|
|
||||||
一大圈 yi2 da4 quan1
|
|
||||||
不再 bu2 zai4
|
|
||||||
一群 yi4 qun2
|
|
||||||
不必 bu2 bi4
|
|
||||||
一些 yi4 xie1
|
|
||||||
一路 yi2 lu4
|
|
||||||
一股 yi4 gu3
|
|
||||||
一到 yi2 dao4
|
|
||||||
一拨 yi4 bo1
|
|
||||||
一排 yi4 pai2
|
|
||||||
一空 yi4 kong1
|
|
||||||
吮吸着 shun3 xi1 zhe
|
|
||||||
不适合 bu2 shi4 he2
|
|
||||||
一串串 yi2 chuan4 chuan4
|
|
||||||
一提起 yi4 ti2 qi3
|
|
||||||
一尘不染 yi4 chen2 bu4 ran3
|
|
||||||
一生 yi4 sheng1
|
|
||||||
一派 yi2 pai4
|
|
||||||
不断 bu2 duan4
|
|
||||||
一次 yi2 ci4
|
|
||||||
不进步 bu2 jin4 bu4
|
|
||||||
娃娃 wa2 wa
|
|
||||||
万户侯 wan4 hu4 hou2
|
|
||||||
一方 yi4 fang1
|
|
||||||
一番话 yi4 fan1 hua4
|
|
||||||
一遍 yi2 bian4
|
|
||||||
不计较 bu2 ji4 jiao4
|
|
||||||
诇 xiong4
|
|
||||||
一边 yi4 bian1
|
|
||||||
一束 yi2 shu4
|
|
||||||
一听到 yi4 ting1 dao4
|
|
||||||
炸鸡 zha2 ji1
|
|
||||||
乍暧还寒 zha4 ai4 huan2 han2
|
|
||||||
我说诶 wo3 shuo1 ei1
|
|
||||||
棒诶 bang4 ei1
|
|
||||||
寒碜 han2 chen4
|
|
||||||
应采儿 ying4 cai3 er2
|
|
||||||
晕车 yun1 che1
|
|
||||||
必应 bi4 ying4
|
|
||||||
应援 ying4 yuan2
|
|
||||||
应力 ying4 li4
|
|
||||||
@ -1,73 +0,0 @@
|
|||||||
# This file is copied from
|
|
||||||
# https://github.com/UEhQZXI/vits_chinese/blob/master/text/symbols.py
|
|
||||||
_pause = ["sil", "eos", "sp", "#0", "#1", "#2", "#3"]
|
|
||||||
|
|
||||||
_initials = [
|
|
||||||
"^",
|
|
||||||
"b",
|
|
||||||
"c",
|
|
||||||
"ch",
|
|
||||||
"d",
|
|
||||||
"f",
|
|
||||||
"g",
|
|
||||||
"h",
|
|
||||||
"j",
|
|
||||||
"k",
|
|
||||||
"l",
|
|
||||||
"m",
|
|
||||||
"n",
|
|
||||||
"p",
|
|
||||||
"q",
|
|
||||||
"r",
|
|
||||||
"s",
|
|
||||||
"sh",
|
|
||||||
"t",
|
|
||||||
"x",
|
|
||||||
"z",
|
|
||||||
"zh",
|
|
||||||
]
|
|
||||||
|
|
||||||
_tones = ["1", "2", "3", "4", "5"]
|
|
||||||
|
|
||||||
_finals = [
|
|
||||||
"a",
|
|
||||||
"ai",
|
|
||||||
"an",
|
|
||||||
"ang",
|
|
||||||
"ao",
|
|
||||||
"e",
|
|
||||||
"ei",
|
|
||||||
"en",
|
|
||||||
"eng",
|
|
||||||
"er",
|
|
||||||
"i",
|
|
||||||
"ia",
|
|
||||||
"ian",
|
|
||||||
"iang",
|
|
||||||
"iao",
|
|
||||||
"ie",
|
|
||||||
"ii",
|
|
||||||
"iii",
|
|
||||||
"in",
|
|
||||||
"ing",
|
|
||||||
"iong",
|
|
||||||
"iou",
|
|
||||||
"o",
|
|
||||||
"ong",
|
|
||||||
"ou",
|
|
||||||
"u",
|
|
||||||
"ua",
|
|
||||||
"uai",
|
|
||||||
"uan",
|
|
||||||
"uang",
|
|
||||||
"uei",
|
|
||||||
"uen",
|
|
||||||
"ueng",
|
|
||||||
"uo",
|
|
||||||
"v",
|
|
||||||
"van",
|
|
||||||
"ve",
|
|
||||||
"vn",
|
|
||||||
]
|
|
||||||
|
|
||||||
symbols = _pause + _initials + [i + j for i in _finals for j in _tones]
|
|
||||||
@ -1,137 +0,0 @@
|
|||||||
# This file is modified from
|
|
||||||
# https://github.com/UEhQZXI/vits_chinese/blob/master/vits_strings.py
|
|
||||||
|
|
||||||
import logging
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import List
|
|
||||||
|
|
||||||
# Note pinyin_dict is from ./pinyin_dict.py
|
|
||||||
from pinyin_dict import pinyin_dict
|
|
||||||
from pypinyin import Style
|
|
||||||
from pypinyin.contrib.neutral_tone import NeutralToneWith5Mixin
|
|
||||||
from pypinyin.converter import DefaultConverter
|
|
||||||
from pypinyin.core import Pinyin, load_phrases_dict
|
|
||||||
|
|
||||||
|
|
||||||
class _MyConverter(NeutralToneWith5Mixin, DefaultConverter):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
class Tokenizer:
|
|
||||||
def __init__(self, tokens: str = ""):
|
|
||||||
self._load_pinyin_dict()
|
|
||||||
self._pinyin_parser = Pinyin(_MyConverter())
|
|
||||||
|
|
||||||
if tokens != "":
|
|
||||||
self._load_tokens(tokens)
|
|
||||||
|
|
||||||
def texts_to_token_ids(self, texts: List[str], **kwargs) -> List[List[int]]:
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
texts:
|
|
||||||
A list of sentences.
|
|
||||||
kwargs:
|
|
||||||
Not used. It is for compatibility with other TTS recipes in icefall.
|
|
||||||
"""
|
|
||||||
tokens = []
|
|
||||||
|
|
||||||
for text in texts:
|
|
||||||
tokens.append(self.text_to_tokens(text))
|
|
||||||
|
|
||||||
return self.tokens_to_token_ids(tokens)
|
|
||||||
|
|
||||||
def tokens_to_token_ids(self, tokens: List[List[str]]) -> List[List[int]]:
|
|
||||||
ans = []
|
|
||||||
|
|
||||||
for token_list in tokens:
|
|
||||||
token_ids = []
|
|
||||||
for t in token_list:
|
|
||||||
if t not in self.token2id:
|
|
||||||
logging.warning(f"Skip OOV {t}")
|
|
||||||
continue
|
|
||||||
token_ids.append(self.token2id[t])
|
|
||||||
ans.append(token_ids)
|
|
||||||
|
|
||||||
return ans
|
|
||||||
|
|
||||||
def text_to_tokens(self, text: str) -> List[str]:
|
|
||||||
# Convert "，" to ["sp", "sil"]
|
|
||||||
# Convert "。" to ["sil"]
|
|
||||||
# append ["eos"] at the end of a sentence
|
|
||||||
phonemes = ["sil"]
|
|
||||||
pinyins = self._pinyin_parser.pinyin(
|
|
||||||
text,
|
|
||||||
style=Style.TONE3,
|
|
||||||
errors=lambda x: [[w] for w in x],
|
|
||||||
)
|
|
||||||
|
|
||||||
new_pinyin = []
|
|
||||||
for p in pinyins:
|
|
||||||
p = p[0]
|
|
||||||
if p == "，":
|
|
||||||
new_pinyin.extend(["sp", "sil"])
|
|
||||||
elif p == "。":
|
|
||||||
new_pinyin.append("sil")
|
|
||||||
else:
|
|
||||||
new_pinyin.append(p)
|
|
||||||
sub_phonemes = self._get_phoneme4pinyin(new_pinyin)
|
|
||||||
sub_phonemes.append("eos")
|
|
||||||
phonemes.extend(sub_phonemes)
|
|
||||||
return phonemes
|
|
||||||
|
|
||||||
def _get_phoneme4pinyin(self, pinyins):
|
|
||||||
result = []
|
|
||||||
for pinyin in pinyins:
|
|
||||||
if pinyin in ("sil", "sp"):
|
|
||||||
result.append(pinyin)
|
|
||||||
elif pinyin[:-1] in pinyin_dict:
|
|
||||||
tone = pinyin[-1]
|
|
||||||
a = pinyin[:-1]
|
|
||||||
a1, a2 = pinyin_dict[a]
|
|
||||||
# every word is appended with a #0
|
|
||||||
result += [a1, a2 + tone, "#0"]
|
|
||||||
|
|
||||||
return result
|
|
||||||
|
|
||||||
def _load_pinyin_dict(self):
|
|
||||||
this_dir = Path(__file__).parent.resolve()
|
|
||||||
my_dict = {}
|
|
||||||
with open(f"{this_dir}/pypinyin-local.dict", "r", encoding="utf-8") as f:
|
|
||||||
content = f.readlines()
|
|
||||||
for line in content:
|
|
||||||
cuts = line.strip().split()
|
|
||||||
hanzi = cuts[0]
|
|
||||||
pinyin = cuts[1:]
|
|
||||||
my_dict[hanzi] = [[p] for p in pinyin]
|
|
||||||
|
|
||||||
load_phrases_dict(my_dict)
|
|
||||||
|
|
||||||
def _load_tokens(self, filename):
|
|
||||||
token2id: Dict[str, int] = {}
|
|
||||||
|
|
||||||
with open(filename, "r", encoding="utf-8") as f:
|
|
||||||
for line in f.readlines():
|
|
||||||
info = line.rstrip().split()
|
|
||||||
if len(info) == 1:
|
|
||||||
# case of space
|
|
||||||
token = " "
|
|
||||||
idx = int(info[0])
|
|
||||||
else:
|
|
||||||
token, idx = info[0], int(info[1])
|
|
||||||
|
|
||||||
assert token not in token2id, token
|
|
||||||
|
|
||||||
token2id[token] = idx
|
|
||||||
|
|
||||||
self.token2id = token2id
|
|
||||||
self.vocab_size = len(self.token2id)
|
|
||||||
self.pad_id = self.token2id["#0"]
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
tokenizer = Tokenizer()
|
|
||||||
tokenizer._sentence_to_ids("你好，好的。")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/local/validate_manifest.py
|
|
||||||
@ -1,124 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
|
|
||||||
# fix segmentation fault reported in https://github.com/k2-fsa/icefall/issues/674
|
|
||||||
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
|
|
||||||
|
|
||||||
set -eou pipefail
|
|
||||||
|
|
||||||
stage=-1
|
|
||||||
stop_stage=100
|
|
||||||
|
|
||||||
dl_dir=$PWD/download
|
|
||||||
|
|
||||||
. shared/parse_options.sh || exit 1
|
|
||||||
|
|
||||||
# All files generated by this script are saved in "data".
|
|
||||||
# You can safely remove "data" and rerun this script to regenerate it.
|
|
||||||
mkdir -p data
|
|
||||||
|
|
||||||
log() {
|
|
||||||
# This function is from espnet
|
|
||||||
local fname=${BASH_SOURCE[1]##*/}
|
|
||||||
echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
|
|
||||||
}
|
|
||||||
|
|
||||||
log "dl_dir: $dl_dir"
|
|
||||||
|
|
||||||
if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
|
|
||||||
log "Stage 0: build monotonic_align lib"
|
|
||||||
if [ ! -d vits/monotonic_align/build ]; then
|
|
||||||
cd vits/monotonic_align
|
|
||||||
python3 setup.py build_ext --inplace
|
|
||||||
cd ../../
|
|
||||||
else
|
|
||||||
log "monotonic_align lib already built"
|
|
||||||
fi
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
|
|
||||||
log "Stage 1: Download data"
|
|
||||||
|
|
||||||
# The directory $dl_dir/BZNSYP will contain 3 sub directories:
|
|
||||||
# - PhoneLabeling
|
|
||||||
# - ProsodyLabeling
|
|
||||||
# - Wave
|
|
||||||
|
|
||||||
# If you have pre-downloaded it to /path/to/BZNSYP, you can create a symlink
|
|
||||||
#
|
|
||||||
# ln -sfv /path/to/BZNSYP $dl_dir/
|
|
||||||
# touch $dl_dir/BZNSYP/.completed
|
|
||||||
#
|
|
||||||
if [ ! -d $dl_dir/BZNSYP ]; then
|
|
||||||
lhotse download baker-zh $dl_dir
|
|
||||||
fi
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
|
|
||||||
log "Stage 2: Prepare baker-zh manifest"
|
|
||||||
# We assume that you have downloaded the baker corpus
|
|
||||||
# to $dl_dir/BZNSYP
|
|
||||||
mkdir -p data/manifests
|
|
||||||
if [ ! -e data/manifests/.baker.done ]; then
|
|
||||||
lhotse prepare baker-zh $dl_dir/BZNSYP data/manifests
|
|
||||||
touch data/manifests/.baker.done
|
|
||||||
fi
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
|
|
||||||
log "Stage 3: Compute spectrogram for baker (may take 3 minutes)"
|
|
||||||
mkdir -p data/spectrogram
|
|
||||||
if [ ! -e data/spectrogram/.baker.done ]; then
|
|
||||||
./local/compute_spectrogram_baker.py
|
|
||||||
touch data/spectrogram/.baker.done
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [ ! -e data/spectrogram/.baker-validated.done ]; then
|
|
||||||
log "Validating data/spectrogram for baker"
|
|
||||||
python3 ./local/validate_manifest.py \
|
|
||||||
data/spectrogram/baker_zh_cuts_all.jsonl.gz
|
|
||||||
touch data/spectrogram/.baker-validated.done
|
|
||||||
fi
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
|
|
||||||
log "Stage 4: Prepare tokens for baker-zh (may take 20 seconds)"
|
|
||||||
if [ ! -e data/spectrogram/.baker_zh_with_token.done ]; then
|
|
||||||
|
|
||||||
./local/prepare_tokens_baker_zh.py
|
|
||||||
|
|
||||||
mv -v data/spectrogram/baker_zh_cuts_with_tokens_all.jsonl.gz \
|
|
||||||
data/spectrogram/baker_zh_cuts_all.jsonl.gz
|
|
||||||
|
|
||||||
touch data/spectrogram/.baker_zh_with_token.done
|
|
||||||
fi
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
|
|
||||||
log "Stage 5: Split the baker-zh cuts into train, valid and test sets (may take 25 seconds)"
|
|
||||||
if [ ! -e data/spectrogram/.baker_zh_split.done ]; then
|
|
||||||
lhotse subset --last 600 \
|
|
||||||
data/spectrogram/baker_zh_cuts_all.jsonl.gz \
|
|
||||||
data/spectrogram/baker_zh_cuts_validtest.jsonl.gz
|
|
||||||
lhotse subset --first 100 \
|
|
||||||
data/spectrogram/baker_zh_cuts_validtest.jsonl.gz \
|
|
||||||
data/spectrogram/baker_zh_cuts_valid.jsonl.gz
|
|
||||||
lhotse subset --last 500 \
|
|
||||||
data/spectrogram/baker_zh_cuts_validtest.jsonl.gz \
|
|
||||||
data/spectrogram/baker_zh_cuts_test.jsonl.gz
|
|
||||||
|
|
||||||
rm data/spectrogram/baker_zh_cuts_validtest.jsonl.gz
|
|
||||||
|
|
||||||
n=$(( $(gunzip -c data/spectrogram/baker_zh_cuts_all.jsonl.gz | wc -l) - 600 ))
|
|
||||||
lhotse subset --first $n \
|
|
||||||
data/spectrogram/baker_zh_cuts_all.jsonl.gz \
|
|
||||||
data/spectrogram/baker_zh_cuts_train.jsonl.gz
|
|
||||||
touch data/spectrogram/.baker_zh_split.done
|
|
||||||
fi
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then
|
|
||||||
log "Stage 6: Generate token file"
|
|
||||||
if [ ! -e data/tokens.txt ]; then
|
|
||||||
./local/prepare_token_file.py --tokens data/tokens.txt
|
|
||||||
fi
|
|
||||||
fi
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../icefall/shared
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/duration_predictor.py
|
|
||||||
@ -1,414 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
#
|
|
||||||
# Copyright 2023 Xiaomi Corporation (Author: Zengwei Yao)
|
|
||||||
#
|
|
||||||
# See ../../../../LICENSE for clarification regarding multiple authors
|
|
||||||
#
|
|
||||||
# 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.
|
|
||||||
|
|
||||||
"""
|
|
||||||
This script exports a VITS model from PyTorch to ONNX.
|
|
||||||
|
|
||||||
Export the model to ONNX:
|
|
||||||
./vits/export-onnx.py \
|
|
||||||
--epoch 1000 \
|
|
||||||
--exp-dir vits/exp \
|
|
||||||
--tokens data/tokens.txt
|
|
||||||
|
|
||||||
It will generate one file inside vits/exp:
|
|
||||||
- vits-epoch-1000.onnx
|
|
||||||
|
|
||||||
See ./test_onnx.py for how to use the exported ONNX models.
|
|
||||||
"""
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import logging
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Dict, Tuple
|
|
||||||
|
|
||||||
import onnx
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
from tokenizer import Tokenizer
|
|
||||||
from train import get_model, get_params
|
|
||||||
|
|
||||||
from icefall.checkpoint import load_checkpoint
|
|
||||||
|
|
||||||
|
|
||||||
def get_parser():
|
|
||||||
parser = argparse.ArgumentParser(
|
|
||||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--epoch",
|
|
||||||
type=int,
|
|
||||||
default=1000,
|
|
||||||
help="""It specifies the checkpoint to use for decoding.
|
|
||||||
Note: Epoch counts from 1.
|
|
||||||
""",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--exp-dir",
|
|
||||||
type=str,
|
|
||||||
default="vits/exp",
|
|
||||||
help="The experiment dir",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--tokens",
|
|
||||||
type=str,
|
|
||||||
default="data/tokens.txt",
|
|
||||||
help="""Path to vocabulary.""",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--model-type",
|
|
||||||
type=str,
|
|
||||||
default="high",
|
|
||||||
choices=["low", "medium", "high"],
|
|
||||||
help="""If not empty, valid values are: low, medium, high.
|
|
||||||
It controls the model size. low -> runs faster.
|
|
||||||
""",
|
|
||||||
)
|
|
||||||
|
|
||||||
return parser
|
|
||||||
|
|
||||||
|
|
||||||
def add_meta_data(filename: str, meta_data: Dict[str, str]):
|
|
||||||
"""Add meta data to an ONNX model. It is changed in-place.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
filename:
|
|
||||||
Filename of the ONNX model to be changed.
|
|
||||||
meta_data:
|
|
||||||
Key-value pairs.
|
|
||||||
"""
|
|
||||||
model = onnx.load(filename)
|
|
||||||
for key, value in meta_data.items():
|
|
||||||
meta = model.metadata_props.add()
|
|
||||||
meta.key = key
|
|
||||||
meta.value = str(value)
|
|
||||||
|
|
||||||
onnx.save(model, filename)
|
|
||||||
|
|
||||||
|
|
||||||
class OnnxModel(nn.Module):
|
|
||||||
"""A wrapper for VITS generator."""
|
|
||||||
|
|
||||||
def __init__(self, model: nn.Module):
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
model:
|
|
||||||
A VITS generator.
|
|
||||||
frame_shift:
|
|
||||||
The frame shift in samples.
|
|
||||||
"""
|
|
||||||
super().__init__()
|
|
||||||
self.model = model
|
|
||||||
|
|
||||||
def forward(
|
|
||||||
self,
|
|
||||||
tokens: torch.Tensor,
|
|
||||||
tokens_lens: torch.Tensor,
|
|
||||||
noise_scale: float = 0.667,
|
|
||||||
alpha: float = 1.0,
|
|
||||||
noise_scale_dur: float = 0.8,
|
|
||||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
||||||
"""Please see the help information of VITS.inference_batch
|
|
||||||
|
|
||||||
Args:
|
|
||||||
tokens:
|
|
||||||
Input text token indexes (1, T_text)
|
|
||||||
tokens_lens:
|
|
||||||
Number of tokens of shape (1,)
|
|
||||||
noise_scale (float):
|
|
||||||
Noise scale parameter for flow.
|
|
||||||
noise_scale_dur (float):
|
|
||||||
Noise scale parameter for duration predictor.
|
|
||||||
alpha (float):
|
|
||||||
Alpha parameter to control the speed of generated speech.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Return a tuple containing:
|
|
||||||
- audio, generated wavform tensor, (B, T_wav)
|
|
||||||
"""
|
|
||||||
audio, _, _ = self.model.generator.inference(
|
|
||||||
text=tokens,
|
|
||||||
text_lengths=tokens_lens,
|
|
||||||
noise_scale=noise_scale,
|
|
||||||
noise_scale_dur=noise_scale_dur,
|
|
||||||
alpha=alpha,
|
|
||||||
)
|
|
||||||
return audio
|
|
||||||
|
|
||||||
|
|
||||||
def export_model_onnx(
|
|
||||||
model: nn.Module,
|
|
||||||
model_filename: str,
|
|
||||||
vocab_size: int,
|
|
||||||
opset_version: int = 11,
|
|
||||||
) -> None:
|
|
||||||
"""Export the given generator model to ONNX format.
|
|
||||||
The exported model has one input:
|
|
||||||
|
|
||||||
- tokens, a tensor of shape (1, T_text); dtype is torch.int64
|
|
||||||
|
|
||||||
and it has one output:
|
|
||||||
|
|
||||||
- audio, a tensor of shape (1, T'); dtype is torch.float32
|
|
||||||
|
|
||||||
Args:
|
|
||||||
model:
|
|
||||||
The VITS generator.
|
|
||||||
model_filename:
|
|
||||||
The filename to save the exported ONNX model.
|
|
||||||
vocab_size:
|
|
||||||
Number of tokens used in training.
|
|
||||||
opset_version:
|
|
||||||
The opset version to use.
|
|
||||||
"""
|
|
||||||
tokens = torch.randint(low=0, high=vocab_size, size=(1, 13), dtype=torch.int64)
|
|
||||||
tokens_lens = torch.tensor([tokens.shape[1]], dtype=torch.int64)
|
|
||||||
noise_scale = torch.tensor([1], dtype=torch.float32)
|
|
||||||
noise_scale_dur = torch.tensor([1], dtype=torch.float32)
|
|
||||||
alpha = torch.tensor([1], dtype=torch.float32)
|
|
||||||
|
|
||||||
torch.onnx.export(
|
|
||||||
model,
|
|
||||||
(tokens, tokens_lens, noise_scale, alpha, noise_scale_dur),
|
|
||||||
model_filename,
|
|
||||||
verbose=False,
|
|
||||||
opset_version=opset_version,
|
|
||||||
input_names=[
|
|
||||||
"tokens",
|
|
||||||
"tokens_lens",
|
|
||||||
"noise_scale",
|
|
||||||
"alpha",
|
|
||||||
"noise_scale_dur",
|
|
||||||
],
|
|
||||||
output_names=["audio"],
|
|
||||||
dynamic_axes={
|
|
||||||
"tokens": {0: "N", 1: "T"},
|
|
||||||
"tokens_lens": {0: "N"},
|
|
||||||
"audio": {0: "N", 1: "T"},
|
|
||||||
},
|
|
||||||
)
|
|
||||||
|
|
||||||
if model.model.spks is None:
|
|
||||||
num_speakers = 1
|
|
||||||
else:
|
|
||||||
num_speakers = model.model.spks
|
|
||||||
|
|
||||||
meta_data = {
|
|
||||||
"model_type": "vits",
|
|
||||||
"version": "1",
|
|
||||||
"model_author": "k2-fsa",
|
|
||||||
"comment": "icefall", # must be icefall for models from icefall
|
|
||||||
"language": "Chinese",
|
|
||||||
"n_speakers": num_speakers,
|
|
||||||
"sample_rate": model.model.sampling_rate, # Must match the real sample rate
|
|
||||||
}
|
|
||||||
logging.info(f"meta_data: {meta_data}")
|
|
||||||
|
|
||||||
add_meta_data(filename=model_filename, meta_data=meta_data)
|
|
||||||
|
|
||||||
|
|
||||||
@torch.no_grad()
|
|
||||||
def main():
|
|
||||||
args = get_parser().parse_args()
|
|
||||||
args.exp_dir = Path(args.exp_dir)
|
|
||||||
|
|
||||||
params = get_params()
|
|
||||||
params.update(vars(args))
|
|
||||||
|
|
||||||
tokenizer = Tokenizer(params.tokens)
|
|
||||||
params.blank_id = tokenizer.pad_id
|
|
||||||
params.vocab_size = tokenizer.vocab_size
|
|
||||||
|
|
||||||
logging.info(params)
|
|
||||||
|
|
||||||
logging.info("About to create model")
|
|
||||||
model = get_model(params)
|
|
||||||
|
|
||||||
load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model)
|
|
||||||
|
|
||||||
model.to("cpu")
|
|
||||||
model.eval()
|
|
||||||
|
|
||||||
model = OnnxModel(model=model)
|
|
||||||
|
|
||||||
num_param = sum([p.numel() for p in model.parameters()])
|
|
||||||
logging.info(f"generator parameters: {num_param}, or {num_param/1000/1000} M")
|
|
||||||
|
|
||||||
suffix = f"epoch-{params.epoch}"
|
|
||||||
|
|
||||||
opset_version = 13
|
|
||||||
|
|
||||||
logging.info("Exporting encoder")
|
|
||||||
model_filename = params.exp_dir / f"vits-{suffix}.onnx"
|
|
||||||
export_model_onnx(
|
|
||||||
model,
|
|
||||||
model_filename,
|
|
||||||
params.vocab_size,
|
|
||||||
opset_version=opset_version,
|
|
||||||
)
|
|
||||||
logging.info(f"Exported generator to {model_filename}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
|
|
||||||
logging.basicConfig(format=formatter, level=logging.INFO)
|
|
||||||
main()
|
|
||||||
|
|
||||||
"""
|
|
||||||
Supported languages.
|
|
||||||
|
|
||||||
LJSpeech is using "en-us" from the second column.
|
|
||||||
|
|
||||||
Pty Language Age/Gender VoiceName File Other Languages
|
|
||||||
5 af --/M Afrikaans gmw/af
|
|
||||||
5 am --/M Amharic sem/am
|
|
||||||
5 an --/M Aragonese roa/an
|
|
||||||
5 ar --/M Arabic sem/ar
|
|
||||||
5 as --/M Assamese inc/as
|
|
||||||
5 az --/M Azerbaijani trk/az
|
|
||||||
5 ba --/M Bashkir trk/ba
|
|
||||||
5 be --/M Belarusian zle/be
|
|
||||||
5 bg --/M Bulgarian zls/bg
|
|
||||||
5 bn --/M Bengali inc/bn
|
|
||||||
5 bpy --/M Bishnupriya_Manipuri inc/bpy
|
|
||||||
5 bs --/M Bosnian zls/bs
|
|
||||||
5 ca --/M Catalan roa/ca
|
|
||||||
5 chr-US-Qaaa-x-west --/M Cherokee_ iro/chr
|
|
||||||
5 cmn --/M Chinese_(Mandarin,_latin_as_English) sit/cmn (zh-cmn 5)(zh 5)
|
|
||||||
5 cmn-latn-pinyin --/M Chinese_(Mandarin,_latin_as_Pinyin) sit/cmn-Latn-pinyin (zh-cmn 5)(zh 5)
|
|
||||||
5 cs --/M Czech zlw/cs
|
|
||||||
5 cv --/M Chuvash trk/cv
|
|
||||||
5 cy --/M Welsh cel/cy
|
|
||||||
5 da --/M Danish gmq/da
|
|
||||||
5 de --/M German gmw/de
|
|
||||||
5 el --/M Greek grk/el
|
|
||||||
5 en-029 --/M English_(Caribbean) gmw/en-029 (en 10)
|
|
||||||
2 en-gb --/M English_(Great_Britain) gmw/en (en 2)
|
|
||||||
5 en-gb-scotland --/M English_(Scotland) gmw/en-GB-scotland (en 4)
|
|
||||||
5 en-gb-x-gbclan --/M English_(Lancaster) gmw/en-GB-x-gbclan (en-gb 3)(en 5)
|
|
||||||
5 en-gb-x-gbcwmd --/M English_(West_Midlands) gmw/en-GB-x-gbcwmd (en-gb 9)(en 9)
|
|
||||||
5 en-gb-x-rp --/M English_(Received_Pronunciation) gmw/en-GB-x-rp (en-gb 4)(en 5)
|
|
||||||
2 en-us --/M English_(America) gmw/en-US (en 3)
|
|
||||||
5 en-us-nyc --/M English_(America,_New_York_City) gmw/en-US-nyc
|
|
||||||
5 eo --/M Esperanto art/eo
|
|
||||||
5 es --/M Spanish_(Spain) roa/es
|
|
||||||
5 es-419 --/M Spanish_(Latin_America) roa/es-419 (es-mx 6)
|
|
||||||
5 et --/M Estonian urj/et
|
|
||||||
5 eu --/M Basque eu
|
|
||||||
5 fa --/M Persian ira/fa
|
|
||||||
5 fa-latn --/M Persian_(Pinglish) ira/fa-Latn
|
|
||||||
5 fi --/M Finnish urj/fi
|
|
||||||
5 fr-be --/M French_(Belgium) roa/fr-BE (fr 8)
|
|
||||||
5 fr-ch --/M French_(Switzerland) roa/fr-CH (fr 8)
|
|
||||||
5 fr-fr --/M French_(France) roa/fr (fr 5)
|
|
||||||
5 ga --/M Gaelic_(Irish) cel/ga
|
|
||||||
5 gd --/M Gaelic_(Scottish) cel/gd
|
|
||||||
5 gn --/M Guarani sai/gn
|
|
||||||
5 grc --/M Greek_(Ancient) grk/grc
|
|
||||||
5 gu --/M Gujarati inc/gu
|
|
||||||
5 hak --/M Hakka_Chinese sit/hak
|
|
||||||
5 haw --/M Hawaiian map/haw
|
|
||||||
5 he --/M Hebrew sem/he
|
|
||||||
5 hi --/M Hindi inc/hi
|
|
||||||
5 hr --/M Croatian zls/hr (hbs 5)
|
|
||||||
5 ht --/M Haitian_Creole roa/ht
|
|
||||||
5 hu --/M Hungarian urj/hu
|
|
||||||
5 hy --/M Armenian_(East_Armenia) ine/hy (hy-arevela 5)
|
|
||||||
5 hyw --/M Armenian_(West_Armenia) ine/hyw (hy-arevmda 5)(hy 8)
|
|
||||||
5 ia --/M Interlingua art/ia
|
|
||||||
5 id --/M Indonesian poz/id
|
|
||||||
5 io --/M Ido art/io
|
|
||||||
5 is --/M Icelandic gmq/is
|
|
||||||
5 it --/M Italian roa/it
|
|
||||||
5 ja --/M Japanese jpx/ja
|
|
||||||
5 jbo --/M Lojban art/jbo
|
|
||||||
5 ka --/M Georgian ccs/ka
|
|
||||||
5 kk --/M Kazakh trk/kk
|
|
||||||
5 kl --/M Greenlandic esx/kl
|
|
||||||
5 kn --/M Kannada dra/kn
|
|
||||||
5 ko --/M Korean ko
|
|
||||||
5 kok --/M Konkani inc/kok
|
|
||||||
5 ku --/M Kurdish ira/ku
|
|
||||||
5 ky --/M Kyrgyz trk/ky
|
|
||||||
5 la --/M Latin itc/la
|
|
||||||
5 lb --/M Luxembourgish gmw/lb
|
|
||||||
5 lfn --/M Lingua_Franca_Nova art/lfn
|
|
||||||
5 lt --/M Lithuanian bat/lt
|
|
||||||
5 ltg --/M Latgalian bat/ltg
|
|
||||||
5 lv --/M Latvian bat/lv
|
|
||||||
5 mi --/M Māori poz/mi
|
|
||||||
5 mk --/M Macedonian zls/mk
|
|
||||||
5 ml --/M Malayalam dra/ml
|
|
||||||
5 mr --/M Marathi inc/mr
|
|
||||||
5 ms --/M Malay poz/ms
|
|
||||||
5 mt --/M Maltese sem/mt
|
|
||||||
5 mto --/M Totontepec_Mixe miz/mto
|
|
||||||
5 my --/M Myanmar_(Burmese) sit/my
|
|
||||||
5 nb --/M Norwegian_Bokmål gmq/nb (no 5)
|
|
||||||
5 nci --/M Nahuatl_(Classical) azc/nci
|
|
||||||
5 ne --/M Nepali inc/ne
|
|
||||||
5 nl --/M Dutch gmw/nl
|
|
||||||
5 nog --/M Nogai trk/nog
|
|
||||||
5 om --/M Oromo cus/om
|
|
||||||
5 or --/M Oriya inc/or
|
|
||||||
5 pa --/M Punjabi inc/pa
|
|
||||||
5 pap --/M Papiamento roa/pap
|
|
||||||
5 piqd --/M Klingon art/piqd
|
|
||||||
5 pl --/M Polish zlw/pl
|
|
||||||
5 pt --/M Portuguese_(Portugal) roa/pt (pt-pt 5)
|
|
||||||
5 pt-br --/M Portuguese_(Brazil) roa/pt-BR (pt 6)
|
|
||||||
5 py --/M Pyash art/py
|
|
||||||
5 qdb --/M Lang_Belta art/qdb
|
|
||||||
5 qu --/M Quechua qu
|
|
||||||
5 quc --/M K'iche' myn/quc
|
|
||||||
5 qya --/M Quenya art/qya
|
|
||||||
5 ro --/M Romanian roa/ro
|
|
||||||
5 ru --/M Russian zle/ru
|
|
||||||
5 ru-cl --/M Russian_(Classic) zle/ru-cl
|
|
||||||
2 ru-lv --/M Russian_(Latvia) zle/ru-LV
|
|
||||||
5 sd --/M Sindhi inc/sd
|
|
||||||
5 shn --/M Shan_(Tai_Yai) tai/shn
|
|
||||||
5 si --/M Sinhala inc/si
|
|
||||||
5 sjn --/M Sindarin art/sjn
|
|
||||||
5 sk --/M Slovak zlw/sk
|
|
||||||
5 sl --/M Slovenian zls/sl
|
|
||||||
5 smj --/M Lule_Saami urj/smj
|
|
||||||
5 sq --/M Albanian ine/sq
|
|
||||||
5 sr --/M Serbian zls/sr
|
|
||||||
5 sv --/M Swedish gmq/sv
|
|
||||||
5 sw --/M Swahili bnt/sw
|
|
||||||
5 ta --/M Tamil dra/ta
|
|
||||||
5 te --/M Telugu dra/te
|
|
||||||
5 th --/M Thai tai/th
|
|
||||||
5 tk --/M Turkmen trk/tk
|
|
||||||
5 tn --/M Setswana bnt/tn
|
|
||||||
5 tr --/M Turkish trk/tr
|
|
||||||
5 tt --/M Tatar trk/tt
|
|
||||||
5 ug --/M Uyghur trk/ug
|
|
||||||
5 uk --/M Ukrainian zle/uk
|
|
||||||
5 ur --/M Urdu inc/ur
|
|
||||||
5 uz --/M Uzbek trk/uz
|
|
||||||
5 vi --/M Vietnamese_(Northern) aav/vi
|
|
||||||
5 vi-vn-x-central --/M Vietnamese_(Central) aav/vi-VN-x-central
|
|
||||||
5 vi-vn-x-south --/M Vietnamese_(Southern) aav/vi-VN-x-south
|
|
||||||
5 yue --/M Chinese_(Cantonese) sit/yue (zh-yue 5)(zh 8)
|
|
||||||
5 yue --/M Chinese_(Cantonese,_latin_as_Jyutping) sit/yue-Latn-jyutping (zh-yue 5)(zh 8)
|
|
||||||
"""
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/flow.py
|
|
||||||
@ -1,39 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
|
|
||||||
from pypinyin import phrases_dict, pinyin_dict
|
|
||||||
from tokenizer import Tokenizer
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
filename = "lexicon.txt"
|
|
||||||
tokens = "./data/tokens.txt"
|
|
||||||
tokenizer = Tokenizer(tokens)
|
|
||||||
|
|
||||||
word_dict = pinyin_dict.pinyin_dict
|
|
||||||
phrases = phrases_dict.phrases_dict
|
|
||||||
|
|
||||||
i = 0
|
|
||||||
with open(filename, "w", encoding="utf-8") as f:
|
|
||||||
for key in word_dict:
|
|
||||||
if not (0x4E00 <= key <= 0x9FFF):
|
|
||||||
continue
|
|
||||||
|
|
||||||
w = chr(key)
|
|
||||||
|
|
||||||
# 1 to remove the initial sil
|
|
||||||
# :-1 to remove the final eos
|
|
||||||
tokens = tokenizer.text_to_tokens(w)[1:-1]
|
|
||||||
|
|
||||||
tokens = " ".join(tokens)
|
|
||||||
f.write(f"{w} {tokens}\n")
|
|
||||||
|
|
||||||
for key in phrases:
|
|
||||||
# 1 to remove the initial sil
|
|
||||||
# :-1 to remove the final eos
|
|
||||||
tokens = tokenizer.text_to_tokens(key)[1:-1]
|
|
||||||
tokens = " ".join(tokens)
|
|
||||||
f.write(f"{key} {tokens}\n")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/generator.py
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/hifigan.py
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/loss.py
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/monotonic_align
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../local/pinyin_dict.py
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/posterior_encoder.py
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../local/pypinyin-local.dict
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/residual_coupling.py
|
|
||||||
@ -1,142 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
#
|
|
||||||
# Copyright 2023 Xiaomi Corporation (Author: Zengwei Yao)
|
|
||||||
#
|
|
||||||
# See ../../../../LICENSE for clarification regarding multiple authors
|
|
||||||
#
|
|
||||||
# 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.
|
|
||||||
|
|
||||||
"""
|
|
||||||
This script is used to test the exported onnx model by vits/export-onnx.py
|
|
||||||
|
|
||||||
Use the onnx model to generate a wav:
|
|
||||||
./vits/test_onnx.py \
|
|
||||||
--model-filename vits/exp/vits-epoch-1000.onnx \
|
|
||||||
--tokens data/tokens.txt
|
|
||||||
"""
|
|
||||||
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import logging
|
|
||||||
|
|
||||||
import onnxruntime as ort
|
|
||||||
import torch
|
|
||||||
import torchaudio
|
|
||||||
from tokenizer import Tokenizer
|
|
||||||
|
|
||||||
|
|
||||||
def get_parser():
|
|
||||||
parser = argparse.ArgumentParser(
|
|
||||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--model-filename",
|
|
||||||
type=str,
|
|
||||||
required=True,
|
|
||||||
help="Path to the onnx model.",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--tokens",
|
|
||||||
type=str,
|
|
||||||
default="data/tokens.txt",
|
|
||||||
help="""Path to vocabulary.""",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--text",
|
|
||||||
type=str,
|
|
||||||
default="Ask not what your country can do for you; ask what you can do for your country.",
|
|
||||||
help="Text to generate speech for",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--output-filename",
|
|
||||||
type=str,
|
|
||||||
default="test_onnx.wav",
|
|
||||||
help="Filename to save the generated wave file.",
|
|
||||||
)
|
|
||||||
|
|
||||||
return parser
|
|
||||||
|
|
||||||
|
|
||||||
class OnnxModel:
|
|
||||||
def __init__(self, model_filename: str):
|
|
||||||
session_opts = ort.SessionOptions()
|
|
||||||
session_opts.inter_op_num_threads = 1
|
|
||||||
session_opts.intra_op_num_threads = 1
|
|
||||||
|
|
||||||
self.session_opts = session_opts
|
|
||||||
|
|
||||||
self.model = ort.InferenceSession(
|
|
||||||
model_filename,
|
|
||||||
sess_options=self.session_opts,
|
|
||||||
providers=["CPUExecutionProvider"],
|
|
||||||
)
|
|
||||||
logging.info(f"{self.model.get_modelmeta().custom_metadata_map}")
|
|
||||||
|
|
||||||
metadata = self.model.get_modelmeta().custom_metadata_map
|
|
||||||
self.sample_rate = int(metadata["sample_rate"])
|
|
||||||
|
|
||||||
def __call__(self, tokens: torch.Tensor, tokens_lens: torch.Tensor) -> torch.Tensor:
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
tokens:
|
|
||||||
A 1-D tensor of shape (1, T)
|
|
||||||
Returns:
|
|
||||||
A tensor of shape (1, T')
|
|
||||||
"""
|
|
||||||
noise_scale = torch.tensor([0.667], dtype=torch.float32)
|
|
||||||
noise_scale_dur = torch.tensor([0.8], dtype=torch.float32)
|
|
||||||
alpha = torch.tensor([1.0], dtype=torch.float32)
|
|
||||||
|
|
||||||
out = self.model.run(
|
|
||||||
[
|
|
||||||
self.model.get_outputs()[0].name,
|
|
||||||
],
|
|
||||||
{
|
|
||||||
self.model.get_inputs()[0].name: tokens.numpy(),
|
|
||||||
self.model.get_inputs()[1].name: tokens_lens.numpy(),
|
|
||||||
self.model.get_inputs()[2].name: noise_scale.numpy(),
|
|
||||||
self.model.get_inputs()[3].name: alpha.numpy(),
|
|
||||||
self.model.get_inputs()[4].name: noise_scale_dur.numpy(),
|
|
||||||
},
|
|
||||||
)[0]
|
|
||||||
return torch.from_numpy(out)
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
args = get_parser().parse_args()
|
|
||||||
logging.info(vars(args))
|
|
||||||
|
|
||||||
tokenizer = Tokenizer(args.tokens)
|
|
||||||
|
|
||||||
logging.info("About to create onnx model")
|
|
||||||
model = OnnxModel(args.model_filename)
|
|
||||||
|
|
||||||
text = args.text
|
|
||||||
tokens = tokenizer.texts_to_token_ids([text])
|
|
||||||
tokens = torch.tensor(tokens) # (1, T)
|
|
||||||
tokens_lens = torch.tensor([tokens.shape[1]], dtype=torch.int64) # (1, T)
|
|
||||||
audio = model(tokens, tokens_lens) # (1, T')
|
|
||||||
|
|
||||||
output_filename = args.output_filename
|
|
||||||
torchaudio.save(output_filename, audio, sample_rate=model.sample_rate)
|
|
||||||
logging.info(f"Saved to {output_filename}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
|
|
||||||
logging.basicConfig(format=formatter, level=logging.INFO)
|
|
||||||
main()
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/text_encoder.py
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../local/tokenizer.py
|
|
||||||
@ -1,927 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
# Copyright 2023 Xiaomi Corp. (authors: Zengwei Yao)
|
|
||||||
#
|
|
||||||
# See ../../../../LICENSE for clarification regarding multiple authors
|
|
||||||
#
|
|
||||||
# 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.
|
|
||||||
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import logging
|
|
||||||
from pathlib import Path
|
|
||||||
from shutil import copyfile
|
|
||||||
from typing import Any, Dict, Optional, Tuple, Union
|
|
||||||
|
|
||||||
import k2
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
import torch.multiprocessing as mp
|
|
||||||
import torch.nn as nn
|
|
||||||
from lhotse.cut import Cut
|
|
||||||
from lhotse.utils import fix_random_seed
|
|
||||||
from tokenizer import Tokenizer
|
|
||||||
from torch.cuda.amp import GradScaler, autocast
|
|
||||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
||||||
from torch.optim import Optimizer
|
|
||||||
from torch.utils.tensorboard import SummaryWriter
|
|
||||||
from tts_datamodule import BakerZhSpeechTtsDataModule
|
|
||||||
from utils import MetricsTracker, plot_feature, save_checkpoint
|
|
||||||
from vits import VITS
|
|
||||||
|
|
||||||
from icefall import diagnostics
|
|
||||||
from icefall.checkpoint import load_checkpoint
|
|
||||||
from icefall.dist import cleanup_dist, setup_dist
|
|
||||||
from icefall.env import get_env_info
|
|
||||||
from icefall.hooks import register_inf_check_hooks
|
|
||||||
from icefall.utils import AttributeDict, setup_logger, str2bool
|
|
||||||
|
|
||||||
LRSchedulerType = torch.optim.lr_scheduler._LRScheduler
|
|
||||||
|
|
||||||
|
|
||||||
def get_parser():
|
|
||||||
parser = argparse.ArgumentParser(
|
|
||||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--world-size",
|
|
||||||
type=int,
|
|
||||||
default=1,
|
|
||||||
help="Number of GPUs for DDP training.",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--master-port",
|
|
||||||
type=int,
|
|
||||||
default=12354,
|
|
||||||
help="Master port to use for DDP training.",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--tensorboard",
|
|
||||||
type=str2bool,
|
|
||||||
default=True,
|
|
||||||
help="Should various information be logged in tensorboard.",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--num-epochs",
|
|
||||||
type=int,
|
|
||||||
default=1000,
|
|
||||||
help="Number of epochs to train.",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--start-epoch",
|
|
||||||
type=int,
|
|
||||||
default=1,
|
|
||||||
help="""Resume training from this epoch. It should be positive.
|
|
||||||
If larger than 1, it will load checkpoint from
|
|
||||||
exp-dir/epoch-{start_epoch-1}.pt
|
|
||||||
""",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--exp-dir",
|
|
||||||
type=str,
|
|
||||||
default="vits/exp",
|
|
||||||
help="""The experiment dir.
|
|
||||||
It specifies the directory where all training related
|
|
||||||
files, e.g., checkpoints, log, etc, are saved
|
|
||||||
""",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--tokens",
|
|
||||||
type=str,
|
|
||||||
default="data/tokens.txt",
|
|
||||||
help="""Path to vocabulary.""",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--lr", type=float, default=2.0e-4, help="The base learning rate."
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--seed",
|
|
||||||
type=int,
|
|
||||||
default=42,
|
|
||||||
help="The seed for random generators intended for reproducibility",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--print-diagnostics",
|
|
||||||
type=str2bool,
|
|
||||||
default=False,
|
|
||||||
help="Accumulate stats on activations, print them and exit.",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--inf-check",
|
|
||||||
type=str2bool,
|
|
||||||
default=False,
|
|
||||||
help="Add hooks to check for infinite module outputs and gradients.",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--save-every-n",
|
|
||||||
type=int,
|
|
||||||
default=20,
|
|
||||||
help="""Save checkpoint after processing this number of epochs"
|
|
||||||
periodically. We save checkpoint to exp-dir/ whenever
|
|
||||||
params.cur_epoch % save_every_n == 0. The checkpoint filename
|
|
||||||
has the form: f'exp-dir/epoch-{params.cur_epoch}.pt'.
|
|
||||||
Since it will take around 1000 epochs, we suggest using a large
|
|
||||||
save_every_n to save disk space.
|
|
||||||
""",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--use-fp16",
|
|
||||||
type=str2bool,
|
|
||||||
default=False,
|
|
||||||
help="Whether to use half precision training.",
|
|
||||||
)
|
|
||||||
|
|
||||||
parser.add_argument(
|
|
||||||
"--model-type",
|
|
||||||
type=str,
|
|
||||||
default="high",
|
|
||||||
choices=["low", "medium", "high"],
|
|
||||||
help="""If not empty, valid values are: low, medium, high.
|
|
||||||
It controls the model size. low -> runs faster.
|
|
||||||
""",
|
|
||||||
)
|
|
||||||
|
|
||||||
return parser
|
|
||||||
|
|
||||||
|
|
||||||
def get_params() -> AttributeDict:
|
|
||||||
"""Return a dict containing training parameters.
|
|
||||||
|
|
||||||
All training related parameters that are not passed from the commandline
|
|
||||||
are saved in the variable `params`.
|
|
||||||
|
|
||||||
Commandline options are merged into `params` after they are parsed, so
|
|
||||||
you can also access them via `params`.
|
|
||||||
|
|
||||||
Explanation of options saved in `params`:
|
|
||||||
|
|
||||||
- best_train_loss: Best training loss so far. It is used to select
|
|
||||||
the model that has the lowest training loss. It is
|
|
||||||
updated during the training.
|
|
||||||
|
|
||||||
- best_valid_loss: Best validation loss so far. It is used to select
|
|
||||||
the model that has the lowest validation loss. It is
|
|
||||||
updated during the training.
|
|
||||||
|
|
||||||
- best_train_epoch: It is the epoch that has the best training loss.
|
|
||||||
|
|
||||||
- best_valid_epoch: It is the epoch that has the best validation loss.
|
|
||||||
|
|
||||||
- batch_idx_train: Used to writing statistics to tensorboard. It
|
|
||||||
contains number of batches trained so far across
|
|
||||||
epochs.
|
|
||||||
|
|
||||||
- log_interval: Print training loss if batch_idx % log_interval` is 0
|
|
||||||
|
|
||||||
- valid_interval: Run validation if batch_idx % valid_interval is 0
|
|
||||||
|
|
||||||
- feature_dim: The model input dim. It has to match the one used
|
|
||||||
in computing features.
|
|
||||||
"""
|
|
||||||
params = AttributeDict(
|
|
||||||
{
|
|
||||||
# training params
|
|
||||||
"best_train_loss": float("inf"),
|
|
||||||
"best_valid_loss": float("inf"),
|
|
||||||
"best_train_epoch": -1,
|
|
||||||
"best_valid_epoch": -1,
|
|
||||||
"batch_idx_train": -1, # 0
|
|
||||||
"log_interval": 50,
|
|
||||||
"valid_interval": 200,
|
|
||||||
"env_info": get_env_info(),
|
|
||||||
"sampling_rate": 48000,
|
|
||||||
"frame_shift": 256,
|
|
||||||
"frame_length": 1024,
|
|
||||||
"feature_dim": 513, # 1024 // 2 + 1, 1024 is fft_length
|
|
||||||
"n_mels": 80,
|
|
||||||
"lambda_adv": 1.0, # loss scaling coefficient for adversarial loss
|
|
||||||
"lambda_mel": 45.0, # loss scaling coefficient for Mel loss
|
|
||||||
"lambda_feat_match": 2.0, # loss scaling coefficient for feat match loss
|
|
||||||
"lambda_dur": 1.0, # loss scaling coefficient for duration loss
|
|
||||||
"lambda_kl": 1.0, # loss scaling coefficient for KL divergence loss
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
return params
|
|
||||||
|
|
||||||
|
|
||||||
def load_checkpoint_if_available(
|
|
||||||
params: AttributeDict, model: nn.Module
|
|
||||||
) -> Optional[Dict[str, Any]]:
|
|
||||||
"""Load checkpoint from file.
|
|
||||||
|
|
||||||
If params.start_epoch is larger than 1, it will load the checkpoint from
|
|
||||||
`params.start_epoch - 1`.
|
|
||||||
|
|
||||||
Apart from loading state dict for `model` and `optimizer` it also updates
|
|
||||||
`best_train_epoch`, `best_train_loss`, `best_valid_epoch`,
|
|
||||||
and `best_valid_loss` in `params`.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
params:
|
|
||||||
The return value of :func:`get_params`.
|
|
||||||
model:
|
|
||||||
The training model.
|
|
||||||
Returns:
|
|
||||||
Return a dict containing previously saved training info.
|
|
||||||
"""
|
|
||||||
if params.start_epoch > 1:
|
|
||||||
filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt"
|
|
||||||
else:
|
|
||||||
return None
|
|
||||||
|
|
||||||
assert filename.is_file(), f"{filename} does not exist!"
|
|
||||||
|
|
||||||
saved_params = load_checkpoint(filename, model=model)
|
|
||||||
|
|
||||||
keys = [
|
|
||||||
"best_train_epoch",
|
|
||||||
"best_valid_epoch",
|
|
||||||
"batch_idx_train",
|
|
||||||
"best_train_loss",
|
|
||||||
"best_valid_loss",
|
|
||||||
]
|
|
||||||
for k in keys:
|
|
||||||
params[k] = saved_params[k]
|
|
||||||
|
|
||||||
return saved_params
|
|
||||||
|
|
||||||
|
|
||||||
def get_model(params: AttributeDict) -> nn.Module:
|
|
||||||
mel_loss_params = {
|
|
||||||
"n_mels": params.n_mels,
|
|
||||||
"frame_length": params.frame_length,
|
|
||||||
"frame_shift": params.frame_shift,
|
|
||||||
}
|
|
||||||
model = VITS(
|
|
||||||
vocab_size=params.vocab_size,
|
|
||||||
feature_dim=params.feature_dim,
|
|
||||||
sampling_rate=params.sampling_rate,
|
|
||||||
model_type=params.model_type,
|
|
||||||
mel_loss_params=mel_loss_params,
|
|
||||||
lambda_adv=params.lambda_adv,
|
|
||||||
lambda_mel=params.lambda_mel,
|
|
||||||
lambda_feat_match=params.lambda_feat_match,
|
|
||||||
lambda_dur=params.lambda_dur,
|
|
||||||
lambda_kl=params.lambda_kl,
|
|
||||||
)
|
|
||||||
return model
|
|
||||||
|
|
||||||
|
|
||||||
def prepare_input(batch: dict, tokenizer: Tokenizer, device: torch.device):
|
|
||||||
"""Parse batch data"""
|
|
||||||
audio = batch["audio"].to(device)
|
|
||||||
features = batch["features"].to(device)
|
|
||||||
audio_lens = batch["audio_lens"].to(device)
|
|
||||||
features_lens = batch["features_lens"].to(device)
|
|
||||||
tokens = batch["tokens"]
|
|
||||||
|
|
||||||
tokens = tokenizer.tokens_to_token_ids(tokens)
|
|
||||||
tokens = k2.RaggedTensor(tokens)
|
|
||||||
row_splits = tokens.shape.row_splits(1)
|
|
||||||
tokens_lens = row_splits[1:] - row_splits[:-1]
|
|
||||||
tokens = tokens.to(device)
|
|
||||||
tokens_lens = tokens_lens.to(device)
|
|
||||||
# a tensor of shape (B, T)
|
|
||||||
tokens = tokens.pad(mode="constant", padding_value=tokenizer.pad_id)
|
|
||||||
|
|
||||||
return audio, audio_lens, features, features_lens, tokens, tokens_lens
|
|
||||||
|
|
||||||
|
|
||||||
def train_one_epoch(
|
|
||||||
params: AttributeDict,
|
|
||||||
model: Union[nn.Module, DDP],
|
|
||||||
tokenizer: Tokenizer,
|
|
||||||
optimizer_g: Optimizer,
|
|
||||||
optimizer_d: Optimizer,
|
|
||||||
scheduler_g: LRSchedulerType,
|
|
||||||
scheduler_d: LRSchedulerType,
|
|
||||||
train_dl: torch.utils.data.DataLoader,
|
|
||||||
valid_dl: torch.utils.data.DataLoader,
|
|
||||||
scaler: GradScaler,
|
|
||||||
tb_writer: Optional[SummaryWriter] = None,
|
|
||||||
world_size: int = 1,
|
|
||||||
rank: int = 0,
|
|
||||||
) -> None:
|
|
||||||
"""Train the model for one epoch.
|
|
||||||
|
|
||||||
The training loss from the mean of all frames is saved in
|
|
||||||
`params.train_loss`. It runs the validation process every
|
|
||||||
`params.valid_interval` batches.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
params:
|
|
||||||
It is returned by :func:`get_params`.
|
|
||||||
model:
|
|
||||||
The model for training.
|
|
||||||
tokenizer:
|
|
||||||
Used to convert text to phonemes.
|
|
||||||
optimizer_g:
|
|
||||||
The optimizer for generator.
|
|
||||||
optimizer_d:
|
|
||||||
The optimizer for discriminator.
|
|
||||||
scheduler_g:
|
|
||||||
The learning rate scheduler for generator, we call step() every epoch.
|
|
||||||
scheduler_d:
|
|
||||||
The learning rate scheduler for discriminator, we call step() every epoch.
|
|
||||||
train_dl:
|
|
||||||
Dataloader for the training dataset.
|
|
||||||
valid_dl:
|
|
||||||
Dataloader for the validation dataset.
|
|
||||||
scaler:
|
|
||||||
The scaler used for mix precision training.
|
|
||||||
tb_writer:
|
|
||||||
Writer to write log messages to tensorboard.
|
|
||||||
world_size:
|
|
||||||
Number of nodes in DDP training. If it is 1, DDP is disabled.
|
|
||||||
rank:
|
|
||||||
The rank of the node in DDP training. If no DDP is used, it should
|
|
||||||
be set to 0.
|
|
||||||
"""
|
|
||||||
model.train()
|
|
||||||
device = model.device if isinstance(model, DDP) else next(model.parameters()).device
|
|
||||||
|
|
||||||
# used to track the stats over iterations in one epoch
|
|
||||||
tot_loss = MetricsTracker()
|
|
||||||
|
|
||||||
saved_bad_model = False
|
|
||||||
|
|
||||||
def save_bad_model(suffix: str = ""):
|
|
||||||
save_checkpoint(
|
|
||||||
filename=params.exp_dir / f"bad-model{suffix}-{rank}.pt",
|
|
||||||
model=model,
|
|
||||||
params=params,
|
|
||||||
optimizer_g=optimizer_g,
|
|
||||||
optimizer_d=optimizer_d,
|
|
||||||
scheduler_g=scheduler_g,
|
|
||||||
scheduler_d=scheduler_d,
|
|
||||||
sampler=train_dl.sampler,
|
|
||||||
scaler=scaler,
|
|
||||||
rank=0,
|
|
||||||
)
|
|
||||||
|
|
||||||
for batch_idx, batch in enumerate(train_dl):
|
|
||||||
params.batch_idx_train += 1
|
|
||||||
|
|
||||||
batch_size = len(batch["tokens"])
|
|
||||||
audio, audio_lens, features, features_lens, tokens, tokens_lens = prepare_input(
|
|
||||||
batch, tokenizer, device
|
|
||||||
)
|
|
||||||
|
|
||||||
loss_info = MetricsTracker()
|
|
||||||
loss_info["samples"] = batch_size
|
|
||||||
|
|
||||||
try:
|
|
||||||
with autocast(enabled=params.use_fp16):
|
|
||||||
# forward discriminator
|
|
||||||
loss_d, stats_d = model(
|
|
||||||
text=tokens,
|
|
||||||
text_lengths=tokens_lens,
|
|
||||||
feats=features,
|
|
||||||
feats_lengths=features_lens,
|
|
||||||
speech=audio,
|
|
||||||
speech_lengths=audio_lens,
|
|
||||||
forward_generator=False,
|
|
||||||
)
|
|
||||||
for k, v in stats_d.items():
|
|
||||||
loss_info[k] = v * batch_size
|
|
||||||
# update discriminator
|
|
||||||
optimizer_d.zero_grad()
|
|
||||||
scaler.scale(loss_d).backward()
|
|
||||||
scaler.step(optimizer_d)
|
|
||||||
|
|
||||||
with autocast(enabled=params.use_fp16):
|
|
||||||
# forward generator
|
|
||||||
loss_g, stats_g = model(
|
|
||||||
text=tokens,
|
|
||||||
text_lengths=tokens_lens,
|
|
||||||
feats=features,
|
|
||||||
feats_lengths=features_lens,
|
|
||||||
speech=audio,
|
|
||||||
speech_lengths=audio_lens,
|
|
||||||
forward_generator=True,
|
|
||||||
return_sample=params.batch_idx_train % params.log_interval == 0,
|
|
||||||
)
|
|
||||||
for k, v in stats_g.items():
|
|
||||||
if "returned_sample" not in k:
|
|
||||||
loss_info[k] = v * batch_size
|
|
||||||
# update generator
|
|
||||||
optimizer_g.zero_grad()
|
|
||||||
scaler.scale(loss_g).backward()
|
|
||||||
scaler.step(optimizer_g)
|
|
||||||
scaler.update()
|
|
||||||
|
|
||||||
# summary stats
|
|
||||||
tot_loss = tot_loss + loss_info
|
|
||||||
except: # noqa
|
|
||||||
save_bad_model()
|
|
||||||
raise
|
|
||||||
|
|
||||||
if params.print_diagnostics and batch_idx == 5:
|
|
||||||
return
|
|
||||||
|
|
||||||
if params.batch_idx_train % 100 == 0 and params.use_fp16:
|
|
||||||
# If the grad scale was less than 1, try increasing it. The _growth_interval
|
|
||||||
# of the grad scaler is configurable, but we can't configure it to have different
|
|
||||||
# behavior depending on the current grad scale.
|
|
||||||
cur_grad_scale = scaler._scale.item()
|
|
||||||
|
|
||||||
if cur_grad_scale < 8.0 or (
|
|
||||||
cur_grad_scale < 32.0 and params.batch_idx_train % 400 == 0
|
|
||||||
):
|
|
||||||
scaler.update(cur_grad_scale * 2.0)
|
|
||||||
if cur_grad_scale < 0.01:
|
|
||||||
if not saved_bad_model:
|
|
||||||
save_bad_model(suffix="-first-warning")
|
|
||||||
saved_bad_model = True
|
|
||||||
logging.warning(f"Grad scale is small: {cur_grad_scale}")
|
|
||||||
if cur_grad_scale < 1.0e-05:
|
|
||||||
save_bad_model()
|
|
||||||
raise RuntimeError(
|
|
||||||
f"grad_scale is too small, exiting: {cur_grad_scale}"
|
|
||||||
)
|
|
||||||
|
|
||||||
if params.batch_idx_train % params.log_interval == 0:
|
|
||||||
cur_lr_g = max(scheduler_g.get_last_lr())
|
|
||||||
cur_lr_d = max(scheduler_d.get_last_lr())
|
|
||||||
cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0
|
|
||||||
|
|
||||||
logging.info(
|
|
||||||
f"Epoch {params.cur_epoch}, batch {batch_idx}, "
|
|
||||||
f"global_batch_idx: {params.batch_idx_train}, batch size: {batch_size}, "
|
|
||||||
f"loss[{loss_info}], tot_loss[{tot_loss}], "
|
|
||||||
f"cur_lr_g: {cur_lr_g:.2e}, cur_lr_d: {cur_lr_d:.2e}, "
|
|
||||||
+ (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "")
|
|
||||||
)
|
|
||||||
|
|
||||||
if tb_writer is not None:
|
|
||||||
tb_writer.add_scalar(
|
|
||||||
"train/learning_rate_g", cur_lr_g, params.batch_idx_train
|
|
||||||
)
|
|
||||||
tb_writer.add_scalar(
|
|
||||||
"train/learning_rate_d", cur_lr_d, params.batch_idx_train
|
|
||||||
)
|
|
||||||
loss_info.write_summary(
|
|
||||||
tb_writer, "train/current_", params.batch_idx_train
|
|
||||||
)
|
|
||||||
tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train)
|
|
||||||
if params.use_fp16:
|
|
||||||
tb_writer.add_scalar(
|
|
||||||
"train/grad_scale", cur_grad_scale, params.batch_idx_train
|
|
||||||
)
|
|
||||||
if "returned_sample" in stats_g:
|
|
||||||
speech_hat_, speech_, mel_hat_, mel_ = stats_g["returned_sample"]
|
|
||||||
tb_writer.add_audio(
|
|
||||||
"train/speech_hat_",
|
|
||||||
speech_hat_,
|
|
||||||
params.batch_idx_train,
|
|
||||||
params.sampling_rate,
|
|
||||||
)
|
|
||||||
tb_writer.add_audio(
|
|
||||||
"train/speech_",
|
|
||||||
speech_,
|
|
||||||
params.batch_idx_train,
|
|
||||||
params.sampling_rate,
|
|
||||||
)
|
|
||||||
tb_writer.add_image(
|
|
||||||
"train/mel_hat_",
|
|
||||||
plot_feature(mel_hat_),
|
|
||||||
params.batch_idx_train,
|
|
||||||
dataformats="HWC",
|
|
||||||
)
|
|
||||||
tb_writer.add_image(
|
|
||||||
"train/mel_",
|
|
||||||
plot_feature(mel_),
|
|
||||||
params.batch_idx_train,
|
|
||||||
dataformats="HWC",
|
|
||||||
)
|
|
||||||
|
|
||||||
if (
|
|
||||||
params.batch_idx_train % params.valid_interval == 0
|
|
||||||
and not params.print_diagnostics
|
|
||||||
):
|
|
||||||
logging.info("Computing validation loss")
|
|
||||||
valid_info, (speech_hat, speech) = compute_validation_loss(
|
|
||||||
params=params,
|
|
||||||
model=model,
|
|
||||||
tokenizer=tokenizer,
|
|
||||||
valid_dl=valid_dl,
|
|
||||||
world_size=world_size,
|
|
||||||
)
|
|
||||||
model.train()
|
|
||||||
logging.info(f"Epoch {params.cur_epoch}, validation: {valid_info}")
|
|
||||||
logging.info(
|
|
||||||
f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
|
|
||||||
)
|
|
||||||
if tb_writer is not None:
|
|
||||||
valid_info.write_summary(
|
|
||||||
tb_writer, "train/valid_", params.batch_idx_train
|
|
||||||
)
|
|
||||||
tb_writer.add_audio(
|
|
||||||
"train/valdi_speech_hat",
|
|
||||||
speech_hat,
|
|
||||||
params.batch_idx_train,
|
|
||||||
params.sampling_rate,
|
|
||||||
)
|
|
||||||
tb_writer.add_audio(
|
|
||||||
"train/valdi_speech",
|
|
||||||
speech,
|
|
||||||
params.batch_idx_train,
|
|
||||||
params.sampling_rate,
|
|
||||||
)
|
|
||||||
|
|
||||||
loss_value = tot_loss["generator_loss"] / tot_loss["samples"]
|
|
||||||
params.train_loss = loss_value
|
|
||||||
if params.train_loss < params.best_train_loss:
|
|
||||||
params.best_train_epoch = params.cur_epoch
|
|
||||||
params.best_train_loss = params.train_loss
|
|
||||||
|
|
||||||
|
|
||||||
def compute_validation_loss(
|
|
||||||
params: AttributeDict,
|
|
||||||
model: Union[nn.Module, DDP],
|
|
||||||
tokenizer: Tokenizer,
|
|
||||||
valid_dl: torch.utils.data.DataLoader,
|
|
||||||
world_size: int = 1,
|
|
||||||
rank: int = 0,
|
|
||||||
) -> Tuple[MetricsTracker, Tuple[np.ndarray, np.ndarray]]:
|
|
||||||
"""Run the validation process."""
|
|
||||||
model.eval()
|
|
||||||
device = model.device if isinstance(model, DDP) else next(model.parameters()).device
|
|
||||||
|
|
||||||
# used to summary the stats over iterations
|
|
||||||
tot_loss = MetricsTracker()
|
|
||||||
returned_sample = None
|
|
||||||
|
|
||||||
with torch.no_grad():
|
|
||||||
for batch_idx, batch in enumerate(valid_dl):
|
|
||||||
batch_size = len(batch["tokens"])
|
|
||||||
(
|
|
||||||
audio,
|
|
||||||
audio_lens,
|
|
||||||
features,
|
|
||||||
features_lens,
|
|
||||||
tokens,
|
|
||||||
tokens_lens,
|
|
||||||
) = prepare_input(batch, tokenizer, device)
|
|
||||||
|
|
||||||
loss_info = MetricsTracker()
|
|
||||||
loss_info["samples"] = batch_size
|
|
||||||
|
|
||||||
# forward discriminator
|
|
||||||
loss_d, stats_d = model(
|
|
||||||
text=tokens,
|
|
||||||
text_lengths=tokens_lens,
|
|
||||||
feats=features,
|
|
||||||
feats_lengths=features_lens,
|
|
||||||
speech=audio,
|
|
||||||
speech_lengths=audio_lens,
|
|
||||||
forward_generator=False,
|
|
||||||
)
|
|
||||||
assert loss_d.requires_grad is False
|
|
||||||
for k, v in stats_d.items():
|
|
||||||
loss_info[k] = v * batch_size
|
|
||||||
|
|
||||||
# forward generator
|
|
||||||
loss_g, stats_g = model(
|
|
||||||
text=tokens,
|
|
||||||
text_lengths=tokens_lens,
|
|
||||||
feats=features,
|
|
||||||
feats_lengths=features_lens,
|
|
||||||
speech=audio,
|
|
||||||
speech_lengths=audio_lens,
|
|
||||||
forward_generator=True,
|
|
||||||
)
|
|
||||||
assert loss_g.requires_grad is False
|
|
||||||
for k, v in stats_g.items():
|
|
||||||
loss_info[k] = v * batch_size
|
|
||||||
|
|
||||||
# summary stats
|
|
||||||
tot_loss = tot_loss + loss_info
|
|
||||||
|
|
||||||
# infer for first batch:
|
|
||||||
if batch_idx == 0 and rank == 0:
|
|
||||||
inner_model = model.module if isinstance(model, DDP) else model
|
|
||||||
audio_pred, _, duration = inner_model.inference(
|
|
||||||
text=tokens[0, : tokens_lens[0].item()]
|
|
||||||
)
|
|
||||||
audio_pred = audio_pred.data.cpu().numpy()
|
|
||||||
audio_len_pred = (
|
|
||||||
(duration.sum(0) * params.frame_shift).to(dtype=torch.int64).item()
|
|
||||||
)
|
|
||||||
assert audio_len_pred == len(audio_pred), (
|
|
||||||
audio_len_pred,
|
|
||||||
len(audio_pred),
|
|
||||||
)
|
|
||||||
audio_gt = audio[0, : audio_lens[0].item()].data.cpu().numpy()
|
|
||||||
returned_sample = (audio_pred, audio_gt)
|
|
||||||
|
|
||||||
if world_size > 1:
|
|
||||||
tot_loss.reduce(device)
|
|
||||||
|
|
||||||
loss_value = tot_loss["generator_loss"] / tot_loss["samples"]
|
|
||||||
if loss_value < params.best_valid_loss:
|
|
||||||
params.best_valid_epoch = params.cur_epoch
|
|
||||||
params.best_valid_loss = loss_value
|
|
||||||
|
|
||||||
return tot_loss, returned_sample
|
|
||||||
|
|
||||||
|
|
||||||
def scan_pessimistic_batches_for_oom(
|
|
||||||
model: Union[nn.Module, DDP],
|
|
||||||
train_dl: torch.utils.data.DataLoader,
|
|
||||||
tokenizer: Tokenizer,
|
|
||||||
optimizer_g: torch.optim.Optimizer,
|
|
||||||
optimizer_d: torch.optim.Optimizer,
|
|
||||||
params: AttributeDict,
|
|
||||||
):
|
|
||||||
from lhotse.dataset import find_pessimistic_batches
|
|
||||||
|
|
||||||
logging.info(
|
|
||||||
"Sanity check -- see if any of the batches in epoch 1 would cause OOM."
|
|
||||||
)
|
|
||||||
device = model.device if isinstance(model, DDP) else next(model.parameters()).device
|
|
||||||
batches, crit_values = find_pessimistic_batches(train_dl.sampler)
|
|
||||||
for criterion, cuts in batches.items():
|
|
||||||
batch = train_dl.dataset[cuts]
|
|
||||||
audio, audio_lens, features, features_lens, tokens, tokens_lens = prepare_input(
|
|
||||||
batch, tokenizer, device
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
# for discriminator
|
|
||||||
with autocast(enabled=params.use_fp16):
|
|
||||||
loss_d, stats_d = model(
|
|
||||||
text=tokens,
|
|
||||||
text_lengths=tokens_lens,
|
|
||||||
feats=features,
|
|
||||||
feats_lengths=features_lens,
|
|
||||||
speech=audio,
|
|
||||||
speech_lengths=audio_lens,
|
|
||||||
forward_generator=False,
|
|
||||||
)
|
|
||||||
optimizer_d.zero_grad()
|
|
||||||
loss_d.backward()
|
|
||||||
# for generator
|
|
||||||
with autocast(enabled=params.use_fp16):
|
|
||||||
loss_g, stats_g = model(
|
|
||||||
text=tokens,
|
|
||||||
text_lengths=tokens_lens,
|
|
||||||
feats=features,
|
|
||||||
feats_lengths=features_lens,
|
|
||||||
speech=audio,
|
|
||||||
speech_lengths=audio_lens,
|
|
||||||
forward_generator=True,
|
|
||||||
)
|
|
||||||
optimizer_g.zero_grad()
|
|
||||||
loss_g.backward()
|
|
||||||
except Exception as e:
|
|
||||||
if "CUDA out of memory" in str(e):
|
|
||||||
logging.error(
|
|
||||||
"Your GPU ran out of memory with the current "
|
|
||||||
"max_duration setting. We recommend decreasing "
|
|
||||||
"max_duration and trying again.\n"
|
|
||||||
f"Failing criterion: {criterion} "
|
|
||||||
f"(={crit_values[criterion]}) ..."
|
|
||||||
)
|
|
||||||
raise
|
|
||||||
logging.info(
|
|
||||||
f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def run(rank, world_size, args):
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
rank:
|
|
||||||
It is a value between 0 and `world_size-1`, which is
|
|
||||||
passed automatically by `mp.spawn()` in :func:`main`.
|
|
||||||
The node with rank 0 is responsible for saving checkpoint.
|
|
||||||
world_size:
|
|
||||||
Number of GPUs for DDP training.
|
|
||||||
args:
|
|
||||||
The return value of get_parser().parse_args()
|
|
||||||
"""
|
|
||||||
params = get_params()
|
|
||||||
params.update(vars(args))
|
|
||||||
|
|
||||||
fix_random_seed(params.seed)
|
|
||||||
if world_size > 1:
|
|
||||||
setup_dist(rank, world_size, params.master_port)
|
|
||||||
|
|
||||||
setup_logger(f"{params.exp_dir}/log/log-train")
|
|
||||||
logging.info("Training started")
|
|
||||||
|
|
||||||
if args.tensorboard and rank == 0:
|
|
||||||
tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard")
|
|
||||||
else:
|
|
||||||
tb_writer = None
|
|
||||||
|
|
||||||
device = torch.device("cpu")
|
|
||||||
if torch.cuda.is_available():
|
|
||||||
device = torch.device("cuda", rank)
|
|
||||||
logging.info(f"Device: {device}")
|
|
||||||
|
|
||||||
tokenizer = Tokenizer(params.tokens)
|
|
||||||
params.blank_id = tokenizer.pad_id
|
|
||||||
params.vocab_size = tokenizer.vocab_size
|
|
||||||
|
|
||||||
logging.info(params)
|
|
||||||
|
|
||||||
logging.info("About to create model")
|
|
||||||
model = get_model(params)
|
|
||||||
generator = model.generator
|
|
||||||
discriminator = model.discriminator
|
|
||||||
|
|
||||||
num_param_g = sum([p.numel() for p in generator.parameters()])
|
|
||||||
logging.info(f"Number of parameters in generator: {num_param_g}")
|
|
||||||
num_param_d = sum([p.numel() for p in discriminator.parameters()])
|
|
||||||
logging.info(f"Number of parameters in discriminator: {num_param_d}")
|
|
||||||
logging.info(f"Total number of parameters: {num_param_g + num_param_d}")
|
|
||||||
|
|
||||||
assert params.start_epoch > 0, params.start_epoch
|
|
||||||
checkpoints = load_checkpoint_if_available(params=params, model=model)
|
|
||||||
|
|
||||||
model.to(device)
|
|
||||||
if world_size > 1:
|
|
||||||
logging.info("Using DDP")
|
|
||||||
model = DDP(model, device_ids=[rank], find_unused_parameters=True)
|
|
||||||
|
|
||||||
optimizer_g = torch.optim.AdamW(
|
|
||||||
generator.parameters(), lr=params.lr, betas=(0.8, 0.99), eps=1e-9
|
|
||||||
)
|
|
||||||
optimizer_d = torch.optim.AdamW(
|
|
||||||
discriminator.parameters(), lr=params.lr, betas=(0.8, 0.99), eps=1e-9
|
|
||||||
)
|
|
||||||
|
|
||||||
scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optimizer_g, gamma=0.999875)
|
|
||||||
scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optimizer_d, gamma=0.999875)
|
|
||||||
|
|
||||||
if checkpoints is not None:
|
|
||||||
# load state_dict for optimizers
|
|
||||||
if "optimizer_g" in checkpoints:
|
|
||||||
logging.info("Loading optimizer_g state dict")
|
|
||||||
optimizer_g.load_state_dict(checkpoints["optimizer_g"])
|
|
||||||
if "optimizer_d" in checkpoints:
|
|
||||||
logging.info("Loading optimizer_d state dict")
|
|
||||||
optimizer_d.load_state_dict(checkpoints["optimizer_d"])
|
|
||||||
|
|
||||||
# load state_dict for schedulers
|
|
||||||
if "scheduler_g" in checkpoints:
|
|
||||||
logging.info("Loading scheduler_g state dict")
|
|
||||||
scheduler_g.load_state_dict(checkpoints["scheduler_g"])
|
|
||||||
if "scheduler_d" in checkpoints:
|
|
||||||
logging.info("Loading scheduler_d state dict")
|
|
||||||
scheduler_d.load_state_dict(checkpoints["scheduler_d"])
|
|
||||||
|
|
||||||
if params.print_diagnostics:
|
|
||||||
opts = diagnostics.TensorDiagnosticOptions(
|
|
||||||
512
|
|
||||||
) # allow 4 megabytes per sub-module
|
|
||||||
diagnostic = diagnostics.attach_diagnostics(model, opts)
|
|
||||||
|
|
||||||
if params.inf_check:
|
|
||||||
register_inf_check_hooks(model)
|
|
||||||
|
|
||||||
baker_zh = BakerZhSpeechTtsDataModule(args)
|
|
||||||
|
|
||||||
train_cuts = baker_zh.train_cuts()
|
|
||||||
|
|
||||||
def remove_short_and_long_utt(c: Cut):
|
|
||||||
# Keep only utterances with duration between 1 second and 20 seconds
|
|
||||||
# You should use ../local/display_manifest_statistics.py to get
|
|
||||||
# an utterance duration distribution for your dataset to select
|
|
||||||
# the threshold
|
|
||||||
if c.duration < 1.0 or c.duration > 20.0:
|
|
||||||
# logging.warning(
|
|
||||||
# f"Exclude cut with ID {c.id} from training. Duration: {c.duration}"
|
|
||||||
# )
|
|
||||||
return False
|
|
||||||
return True
|
|
||||||
|
|
||||||
train_cuts = train_cuts.filter(remove_short_and_long_utt)
|
|
||||||
train_dl = baker_zh.train_dataloaders(train_cuts)
|
|
||||||
|
|
||||||
valid_cuts = baker_zh.valid_cuts()
|
|
||||||
valid_dl = baker_zh.valid_dataloaders(valid_cuts)
|
|
||||||
|
|
||||||
if not params.print_diagnostics:
|
|
||||||
scan_pessimistic_batches_for_oom(
|
|
||||||
model=model,
|
|
||||||
train_dl=train_dl,
|
|
||||||
tokenizer=tokenizer,
|
|
||||||
optimizer_g=optimizer_g,
|
|
||||||
optimizer_d=optimizer_d,
|
|
||||||
params=params,
|
|
||||||
)
|
|
||||||
|
|
||||||
scaler = GradScaler(enabled=params.use_fp16, init_scale=1.0)
|
|
||||||
if checkpoints and "grad_scaler" in checkpoints:
|
|
||||||
logging.info("Loading grad scaler state dict")
|
|
||||||
scaler.load_state_dict(checkpoints["grad_scaler"])
|
|
||||||
|
|
||||||
for epoch in range(params.start_epoch, params.num_epochs + 1):
|
|
||||||
logging.info(f"Start epoch {epoch}")
|
|
||||||
|
|
||||||
fix_random_seed(params.seed + epoch - 1)
|
|
||||||
train_dl.sampler.set_epoch(epoch - 1)
|
|
||||||
|
|
||||||
params.cur_epoch = epoch
|
|
||||||
|
|
||||||
if tb_writer is not None:
|
|
||||||
tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train)
|
|
||||||
|
|
||||||
train_one_epoch(
|
|
||||||
params=params,
|
|
||||||
model=model,
|
|
||||||
tokenizer=tokenizer,
|
|
||||||
optimizer_g=optimizer_g,
|
|
||||||
optimizer_d=optimizer_d,
|
|
||||||
scheduler_g=scheduler_g,
|
|
||||||
scheduler_d=scheduler_d,
|
|
||||||
train_dl=train_dl,
|
|
||||||
valid_dl=valid_dl,
|
|
||||||
scaler=scaler,
|
|
||||||
tb_writer=tb_writer,
|
|
||||||
world_size=world_size,
|
|
||||||
rank=rank,
|
|
||||||
)
|
|
||||||
|
|
||||||
if params.print_diagnostics:
|
|
||||||
diagnostic.print_diagnostics()
|
|
||||||
break
|
|
||||||
|
|
||||||
if epoch % params.save_every_n == 0 or epoch == params.num_epochs:
|
|
||||||
filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt"
|
|
||||||
save_checkpoint(
|
|
||||||
filename=filename,
|
|
||||||
params=params,
|
|
||||||
model=model,
|
|
||||||
optimizer_g=optimizer_g,
|
|
||||||
optimizer_d=optimizer_d,
|
|
||||||
scheduler_g=scheduler_g,
|
|
||||||
scheduler_d=scheduler_d,
|
|
||||||
sampler=train_dl.sampler,
|
|
||||||
scaler=scaler,
|
|
||||||
rank=rank,
|
|
||||||
)
|
|
||||||
if rank == 0:
|
|
||||||
if params.best_train_epoch == params.cur_epoch:
|
|
||||||
best_train_filename = params.exp_dir / "best-train-loss.pt"
|
|
||||||
copyfile(src=filename, dst=best_train_filename)
|
|
||||||
|
|
||||||
if params.best_valid_epoch == params.cur_epoch:
|
|
||||||
best_valid_filename = params.exp_dir / "best-valid-loss.pt"
|
|
||||||
copyfile(src=filename, dst=best_valid_filename)
|
|
||||||
|
|
||||||
# step per epoch
|
|
||||||
scheduler_g.step()
|
|
||||||
scheduler_d.step()
|
|
||||||
|
|
||||||
logging.info("Done!")
|
|
||||||
|
|
||||||
if world_size > 1:
|
|
||||||
torch.distributed.barrier()
|
|
||||||
cleanup_dist()
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
parser = get_parser()
|
|
||||||
BakerZhSpeechTtsDataModule.add_arguments(parser)
|
|
||||||
args = parser.parse_args()
|
|
||||||
args.exp_dir = Path(args.exp_dir)
|
|
||||||
|
|
||||||
world_size = args.world_size
|
|
||||||
assert world_size >= 1
|
|
||||||
if world_size > 1:
|
|
||||||
mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True)
|
|
||||||
else:
|
|
||||||
run(rank=0, world_size=1, args=args)
|
|
||||||
|
|
||||||
|
|
||||||
torch.set_num_threads(1)
|
|
||||||
torch.set_num_interop_threads(1)
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/transform.py
|
|
||||||
@ -1,330 +0,0 @@
|
|||||||
# Copyright 2021 Piotr Żelasko
|
|
||||||
# Copyright 2022-2023 Xiaomi Corporation (Authors: Mingshuang Luo,
|
|
||||||
# Zengwei Yao)
|
|
||||||
#
|
|
||||||
# See ../../../../LICENSE for clarification regarding multiple authors
|
|
||||||
#
|
|
||||||
# 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.
|
|
||||||
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import logging
|
|
||||||
from functools import lru_cache
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Any, Dict, Optional
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from lhotse import CutSet, Spectrogram, SpectrogramConfig, load_manifest_lazy
|
|
||||||
from lhotse.dataset import ( # noqa F401 for PrecomputedFeatures
|
|
||||||
CutConcatenate,
|
|
||||||
CutMix,
|
|
||||||
DynamicBucketingSampler,
|
|
||||||
PrecomputedFeatures,
|
|
||||||
SimpleCutSampler,
|
|
||||||
SpecAugment,
|
|
||||||
SpeechSynthesisDataset,
|
|
||||||
)
|
|
||||||
from lhotse.dataset.input_strategies import ( # noqa F401 For AudioSamples
|
|
||||||
AudioSamples,
|
|
||||||
OnTheFlyFeatures,
|
|
||||||
)
|
|
||||||
from lhotse.utils import fix_random_seed
|
|
||||||
from torch.utils.data import DataLoader
|
|
||||||
|
|
||||||
from icefall.utils import str2bool
|
|
||||||
|
|
||||||
|
|
||||||
class _SeedWorkers:
|
|
||||||
def __init__(self, seed: int):
|
|
||||||
self.seed = seed
|
|
||||||
|
|
||||||
def __call__(self, worker_id: int):
|
|
||||||
fix_random_seed(self.seed + worker_id)
|
|
||||||
|
|
||||||
|
|
||||||
class BakerZhSpeechTtsDataModule:
|
|
||||||
"""
|
|
||||||
DataModule for tts experiments.
|
|
||||||
It assumes there is always one train and valid dataloader,
|
|
||||||
but there can be multiple test dataloaders (e.g. LibriSpeech test-clean
|
|
||||||
and test-other).
|
|
||||||
|
|
||||||
It contains all the common data pipeline modules used in ASR
|
|
||||||
experiments, e.g.:
|
|
||||||
- dynamic batch size,
|
|
||||||
- bucketing samplers,
|
|
||||||
- cut concatenation,
|
|
||||||
- on-the-fly feature extraction
|
|
||||||
|
|
||||||
This class should be derived for specific corpora used in TTS tasks.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, args: argparse.Namespace):
|
|
||||||
self.args = args
|
|
||||||
self.sampling_rate = 48000
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def add_arguments(cls, parser: argparse.ArgumentParser):
|
|
||||||
group = parser.add_argument_group(
|
|
||||||
title="TTS data related options",
|
|
||||||
description="These options are used for the preparation of "
|
|
||||||
"PyTorch DataLoaders from Lhotse CutSet's -- they control the "
|
|
||||||
"effective batch sizes, sampling strategies, applied data "
|
|
||||||
"augmentations, etc.",
|
|
||||||
)
|
|
||||||
|
|
||||||
group.add_argument(
|
|
||||||
"--manifest-dir",
|
|
||||||
type=Path,
|
|
||||||
default=Path("data/spectrogram"),
|
|
||||||
help="Path to directory with train/valid/test cuts.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--max-duration",
|
|
||||||
type=int,
|
|
||||||
default=200.0,
|
|
||||||
help="Maximum pooled recordings duration (seconds) in a "
|
|
||||||
"single batch. You can reduce it if it causes CUDA OOM.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--bucketing-sampler",
|
|
||||||
type=str2bool,
|
|
||||||
default=True,
|
|
||||||
help="When enabled, the batches will come from buckets of "
|
|
||||||
"similar duration (saves padding frames).",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--num-buckets",
|
|
||||||
type=int,
|
|
||||||
default=30,
|
|
||||||
help="The number of buckets for the DynamicBucketingSampler"
|
|
||||||
"(you might want to increase it for larger datasets).",
|
|
||||||
)
|
|
||||||
|
|
||||||
group.add_argument(
|
|
||||||
"--on-the-fly-feats",
|
|
||||||
type=str2bool,
|
|
||||||
default=False,
|
|
||||||
help="When enabled, use on-the-fly cut mixing and feature "
|
|
||||||
"extraction. Will drop existing precomputed feature manifests "
|
|
||||||
"if available.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--shuffle",
|
|
||||||
type=str2bool,
|
|
||||||
default=True,
|
|
||||||
help="When enabled (=default), the examples will be "
|
|
||||||
"shuffled for each epoch.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--drop-last",
|
|
||||||
type=str2bool,
|
|
||||||
default=True,
|
|
||||||
help="Whether to drop last batch. Used by sampler.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--return-cuts",
|
|
||||||
type=str2bool,
|
|
||||||
default=False,
|
|
||||||
help="When enabled, each batch will have the "
|
|
||||||
"field: batch['cut'] with the cuts that "
|
|
||||||
"were used to construct it.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--num-workers",
|
|
||||||
type=int,
|
|
||||||
default=2,
|
|
||||||
help="The number of training dataloader workers that "
|
|
||||||
"collect the batches.",
|
|
||||||
)
|
|
||||||
|
|
||||||
group.add_argument(
|
|
||||||
"--input-strategy",
|
|
||||||
type=str,
|
|
||||||
default="PrecomputedFeatures",
|
|
||||||
help="AudioSamples or PrecomputedFeatures",
|
|
||||||
)
|
|
||||||
|
|
||||||
def train_dataloaders(
|
|
||||||
self,
|
|
||||||
cuts_train: CutSet,
|
|
||||||
sampler_state_dict: Optional[Dict[str, Any]] = None,
|
|
||||||
) -> DataLoader:
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
cuts_train:
|
|
||||||
CutSet for training.
|
|
||||||
sampler_state_dict:
|
|
||||||
The state dict for the training sampler.
|
|
||||||
"""
|
|
||||||
logging.info("About to create train dataset")
|
|
||||||
train = SpeechSynthesisDataset(
|
|
||||||
return_text=False,
|
|
||||||
return_tokens=True,
|
|
||||||
feature_input_strategy=eval(self.args.input_strategy)(),
|
|
||||||
return_cuts=self.args.return_cuts,
|
|
||||||
)
|
|
||||||
|
|
||||||
if self.args.on_the_fly_feats:
|
|
||||||
sampling_rate = self.sampling_rate
|
|
||||||
config = SpectrogramConfig(
|
|
||||||
sampling_rate=sampling_rate,
|
|
||||||
frame_length=1024 / sampling_rate, # (in second),
|
|
||||||
frame_shift=256 / sampling_rate, # (in second)
|
|
||||||
use_fft_mag=True,
|
|
||||||
)
|
|
||||||
train = SpeechSynthesisDataset(
|
|
||||||
return_text=False,
|
|
||||||
return_tokens=True,
|
|
||||||
feature_input_strategy=OnTheFlyFeatures(Spectrogram(config)),
|
|
||||||
return_cuts=self.args.return_cuts,
|
|
||||||
)
|
|
||||||
|
|
||||||
if self.args.bucketing_sampler:
|
|
||||||
logging.info("Using DynamicBucketingSampler.")
|
|
||||||
train_sampler = DynamicBucketingSampler(
|
|
||||||
cuts_train,
|
|
||||||
max_duration=self.args.max_duration,
|
|
||||||
shuffle=self.args.shuffle,
|
|
||||||
num_buckets=self.args.num_buckets,
|
|
||||||
buffer_size=self.args.num_buckets * 2000,
|
|
||||||
shuffle_buffer_size=self.args.num_buckets * 5000,
|
|
||||||
drop_last=self.args.drop_last,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
logging.info("Using SimpleCutSampler.")
|
|
||||||
train_sampler = SimpleCutSampler(
|
|
||||||
cuts_train,
|
|
||||||
max_duration=self.args.max_duration,
|
|
||||||
shuffle=self.args.shuffle,
|
|
||||||
)
|
|
||||||
logging.info("About to create train dataloader")
|
|
||||||
|
|
||||||
if sampler_state_dict is not None:
|
|
||||||
logging.info("Loading sampler state dict")
|
|
||||||
train_sampler.load_state_dict(sampler_state_dict)
|
|
||||||
|
|
||||||
# 'seed' is derived from the current random state, which will have
|
|
||||||
# previously been set in the main process.
|
|
||||||
seed = torch.randint(0, 100000, ()).item()
|
|
||||||
worker_init_fn = _SeedWorkers(seed)
|
|
||||||
|
|
||||||
train_dl = DataLoader(
|
|
||||||
train,
|
|
||||||
sampler=train_sampler,
|
|
||||||
batch_size=None,
|
|
||||||
num_workers=self.args.num_workers,
|
|
||||||
persistent_workers=False,
|
|
||||||
worker_init_fn=worker_init_fn,
|
|
||||||
)
|
|
||||||
|
|
||||||
return train_dl
|
|
||||||
|
|
||||||
def valid_dataloaders(self, cuts_valid: CutSet) -> DataLoader:
|
|
||||||
logging.info("About to create dev dataset")
|
|
||||||
if self.args.on_the_fly_feats:
|
|
||||||
sampling_rate = self.sampling_rate
|
|
||||||
config = SpectrogramConfig(
|
|
||||||
sampling_rate=sampling_rate,
|
|
||||||
frame_length=1024 / sampling_rate, # (in second),
|
|
||||||
frame_shift=256 / sampling_rate, # (in second)
|
|
||||||
use_fft_mag=True,
|
|
||||||
)
|
|
||||||
validate = SpeechSynthesisDataset(
|
|
||||||
return_text=False,
|
|
||||||
return_tokens=True,
|
|
||||||
feature_input_strategy=OnTheFlyFeatures(Spectrogram(config)),
|
|
||||||
return_cuts=self.args.return_cuts,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
validate = SpeechSynthesisDataset(
|
|
||||||
return_text=False,
|
|
||||||
return_tokens=True,
|
|
||||||
feature_input_strategy=eval(self.args.input_strategy)(),
|
|
||||||
return_cuts=self.args.return_cuts,
|
|
||||||
)
|
|
||||||
valid_sampler = DynamicBucketingSampler(
|
|
||||||
cuts_valid,
|
|
||||||
max_duration=self.args.max_duration,
|
|
||||||
num_buckets=self.args.num_buckets,
|
|
||||||
shuffle=False,
|
|
||||||
)
|
|
||||||
logging.info("About to create valid dataloader")
|
|
||||||
valid_dl = DataLoader(
|
|
||||||
validate,
|
|
||||||
sampler=valid_sampler,
|
|
||||||
batch_size=None,
|
|
||||||
num_workers=2,
|
|
||||||
persistent_workers=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
return valid_dl
|
|
||||||
|
|
||||||
def test_dataloaders(self, cuts: CutSet) -> DataLoader:
|
|
||||||
logging.info("About to create test dataset")
|
|
||||||
if self.args.on_the_fly_feats:
|
|
||||||
sampling_rate = self.sampling_rate
|
|
||||||
config = SpectrogramConfig(
|
|
||||||
sampling_rate=sampling_rate,
|
|
||||||
frame_length=1024 / sampling_rate, # (in second),
|
|
||||||
frame_shift=256 / sampling_rate, # (in second)
|
|
||||||
use_fft_mag=True,
|
|
||||||
)
|
|
||||||
test = SpeechSynthesisDataset(
|
|
||||||
return_text=False,
|
|
||||||
return_tokens=True,
|
|
||||||
feature_input_strategy=OnTheFlyFeatures(Spectrogram(config)),
|
|
||||||
return_cuts=self.args.return_cuts,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
test = SpeechSynthesisDataset(
|
|
||||||
return_text=False,
|
|
||||||
return_tokens=True,
|
|
||||||
feature_input_strategy=eval(self.args.input_strategy)(),
|
|
||||||
return_cuts=self.args.return_cuts,
|
|
||||||
)
|
|
||||||
test_sampler = DynamicBucketingSampler(
|
|
||||||
cuts,
|
|
||||||
max_duration=self.args.max_duration,
|
|
||||||
num_buckets=self.args.num_buckets,
|
|
||||||
shuffle=False,
|
|
||||||
)
|
|
||||||
logging.info("About to create test dataloader")
|
|
||||||
test_dl = DataLoader(
|
|
||||||
test,
|
|
||||||
batch_size=None,
|
|
||||||
sampler=test_sampler,
|
|
||||||
num_workers=self.args.num_workers,
|
|
||||||
)
|
|
||||||
return test_dl
|
|
||||||
|
|
||||||
@lru_cache()
|
|
||||||
def train_cuts(self) -> CutSet:
|
|
||||||
logging.info("About to get train cuts")
|
|
||||||
return load_manifest_lazy(
|
|
||||||
self.args.manifest_dir / "baker_zh_cuts_train.jsonl.gz"
|
|
||||||
)
|
|
||||||
|
|
||||||
@lru_cache()
|
|
||||||
def valid_cuts(self) -> CutSet:
|
|
||||||
logging.info("About to get validation cuts")
|
|
||||||
return load_manifest_lazy(
|
|
||||||
self.args.manifest_dir / "baker_zh_cuts_valid.jsonl.gz"
|
|
||||||
)
|
|
||||||
|
|
||||||
@lru_cache()
|
|
||||||
def test_cuts(self) -> CutSet:
|
|
||||||
logging.info("About to get test cuts")
|
|
||||||
return load_manifest_lazy(
|
|
||||||
self.args.manifest_dir / "baker_zh_cuts_test.jsonl.gz"
|
|
||||||
)
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/utils.py
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/vits.py
|
|
||||||
@ -1 +0,0 @@
|
|||||||
../../../ljspeech/TTS/vits/wavenet.py
|
|
||||||
Loading…
x
Reference in New Issue
Block a user