tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Refactor: separate module for utilities related to loading/saving checkpoints and safetensors

Browse Source
This commit is contained in:
tsukumi
2024-03-09 15:58:57 +00:00
parent c594f7ea7a
commit 98ab8e7978
7 changed files with 510 additions and 457 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
style_bert_vits2/models/infer.py
View File

@@ -80,9 +80,9 @@ def get_net_g(model_path: str, version: str, device: str, hps: HyperParameters):
net_g.state_dict()
_ = net_g.eval()
if model_path.endswith(".pth") or model_path.endswith(".pt"):
_ = utils.load_checkpoint(model_path, net_g, None, skip_optimizer=True)
_ = utils.checkpoints.load_checkpoint(model_path, net_g, None, skip_optimizer=True)
elif model_path.endswith(".safetensors"):
_ = utils.load_safetensors(model_path, net_g, True)
_ = utils.safetensors.load_safetensors(model_path, net_g, True)
else:
raise ValueError(f"Unknown model format: {model_path}")
return net_g

357
style_bert_vits2/models/utils.py
View File

@@ -1,357 +0,0 @@
import argparse
import glob
import json
import logging
import os
import re
import subprocess
import numpy as np
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from scipy.io.wavfile import read
from style_bert_vits2.logging import logger
MATPLOTLIB_FLAG = False
def load_checkpoint(
checkpoint_path, model, optimizer=None, skip_optimizer=False, for_infer=False
):
assert os.path.isfile(checkpoint_path)
checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")
iteration = checkpoint_dict["iteration"]
learning_rate = checkpoint_dict["learning_rate"]
logger.info(
f"Loading model and optimizer at iteration {iteration} from {checkpoint_path}"
)
if (
optimizer is not None
and not skip_optimizer
and checkpoint_dict["optimizer"] is not None
):
optimizer.load_state_dict(checkpoint_dict["optimizer"])
elif optimizer is None and not skip_optimizer:
# else: Disable this line if Infer and resume checkpoint,then enable the line upper
new_opt_dict = optimizer.state_dict()
new_opt_dict_params = new_opt_dict["param_groups"][0]["params"]
new_opt_dict["param_groups"] = checkpoint_dict["optimizer"]["param_groups"]
new_opt_dict["param_groups"][0]["params"] = new_opt_dict_params
optimizer.load_state_dict(new_opt_dict)
saved_state_dict = checkpoint_dict["model"]
if hasattr(model, "module"):
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
new_state_dict = {}
for k, v in state_dict.items():
try:
# assert "emb_g" not in k
new_state_dict[k] = saved_state_dict[k]
assert saved_state_dict[k].shape == v.shape, (
saved_state_dict[k].shape,
v.shape,
)
except:
# For upgrading from the old version
if "ja_bert_proj" in k:
v = torch.zeros_like(v)
logger.warning(
f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility"
)
elif "enc_q" in k and for_infer:
continue
else:
logger.error(f"{k} is not in the checkpoint {checkpoint_path}")
new_state_dict[k] = v
if hasattr(model, "module"):
model.module.load_state_dict(new_state_dict, strict=False)
else:
model.load_state_dict(new_state_dict, strict=False)
logger.info("Loaded '{}' (iteration {})".format(checkpoint_path, iteration))
return model, optimizer, learning_rate, iteration
def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
logger.info(
"Saving model and optimizer state at iteration {} to {}".format(
iteration, checkpoint_path
)
)
if hasattr(model, "module"):
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
torch.save(
{
"model": state_dict,
"iteration": iteration,
"optimizer": optimizer.state_dict(),
"learning_rate": learning_rate,
},
checkpoint_path,
)
def clean_checkpoints(path_to_models="logs/44k/", n_ckpts_to_keep=2, sort_by_time=True):
"""Freeing up space by deleting saved ckpts
Arguments:
path_to_models -- Path to the model directory
n_ckpts_to_keep -- Number of ckpts to keep, excluding G_0.pth and D_0.pth
sort_by_time -- True -> chronologically delete ckpts
False -> lexicographically delete ckpts
"""
import re
ckpts_files = [
f
for f in os.listdir(path_to_models)
if os.path.isfile(os.path.join(path_to_models, f))
]
def name_key(_f):
return int(re.compile("._(\\d+)\\.pth").match(_f).group(1))
def time_key(_f):
return os.path.getmtime(os.path.join(path_to_models, _f))
sort_key = time_key if sort_by_time else name_key
def x_sorted(_x):
return sorted(
[f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")],
key=sort_key,
)
to_del = [
os.path.join(path_to_models, fn)
for fn in (
x_sorted("G_")[:-n_ckpts_to_keep]
+ x_sorted("D_")[:-n_ckpts_to_keep]
+ x_sorted("WD_")[:-n_ckpts_to_keep]
+ x_sorted("DUR_")[:-n_ckpts_to_keep]
)
]
def del_info(fn):
return logger.info(f"Free up space by deleting ckpt {fn}")
def del_routine(x):
return [os.remove(x), del_info(x)]
[del_routine(fn) for fn in to_del]
def load_safetensors(checkpoint_path, model, for_infer=False):
"""
Load safetensors model.
"""
tensors = {}
iteration = None
with safe_open(checkpoint_path, framework="pt", device="cpu") as f:
for key in f.keys():
if key == "iteration":
iteration = f.get_tensor(key).item()
tensors[key] = f.get_tensor(key)
if hasattr(model, "module"):
result = model.module.load_state_dict(tensors, strict=False)
else:
result = model.load_state_dict(tensors, strict=False)
for key in result.missing_keys:
if key.startswith("enc_q") and for_infer:
continue
logger.warning(f"Missing key: {key}")
for key in result.unexpected_keys:
if key == "iteration":
continue
logger.warning(f"Unexpected key: {key}")
if iteration is None:
logger.info(f"Loaded '{checkpoint_path}'")
else:
logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})")
return model, iteration
def save_safetensors(model, iteration, checkpoint_path, is_half=False, for_infer=False):
"""
Save model with safetensors.
"""
if hasattr(model, "module"):
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
keys = []
for k in state_dict:
if "enc_q" in k and for_infer:
continue # noqa: E701
keys.append(k)
new_dict = (
{k: state_dict[k].half() for k in keys}
if is_half
else {k: state_dict[k] for k in keys}
)
new_dict["iteration"] = torch.LongTensor([iteration])
logger.info(f"Saved safetensors to {checkpoint_path}")
save_file(new_dict, checkpoint_path)
def summarize(
writer,
global_step,
scalars={},
histograms={},
images={},
audios={},
audio_sampling_rate=22050,
):
for k, v in scalars.items():
writer.add_scalar(k, v, global_step)
for k, v in histograms.items():
writer.add_histogram(k, v, global_step)
for k, v in images.items():
writer.add_image(k, v, global_step, dataformats="HWC")
for k, v in audios.items():
writer.add_audio(k, v, global_step, audio_sampling_rate)
def is_resuming(dir_path):
# JP-ExtraバージョンではDURがなくWDがあったり変わるため、Gのみで判断する
g_list = glob.glob(os.path.join(dir_path, "G_*.pth"))
# d_list = glob.glob(os.path.join(dir_path, "D_*.pth"))
# dur_list = glob.glob(os.path.join(dir_path, "DUR_*.pth"))
return len(g_list) > 0
def latest_checkpoint_path(dir_path, regex="G_*.pth"):
f_list = glob.glob(os.path.join(dir_path, regex))
f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
try:
x = f_list[-1]
except IndexError:
raise ValueError(f"No checkpoint found in {dir_path} with regex {regex}")
return x
def plot_spectrogram_to_numpy(spectrogram):
global MATPLOTLIB_FLAG
if not MATPLOTLIB_FLAG:
import matplotlib
matplotlib.use("Agg")
MATPLOTLIB_FLAG = True
mpl_logger = logging.getLogger("matplotlib")
mpl_logger.setLevel(logging.WARNING)
import matplotlib.pylab as plt
import numpy as np
fig, ax = plt.subplots(figsize=(10, 2))
im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
plt.colorbar(im, ax=ax)
plt.xlabel("Frames")
plt.ylabel("Channels")
plt.tight_layout()
fig.canvas.draw()
data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
plt.close()
return data
def plot_alignment_to_numpy(alignment, info=None):
global MATPLOTLIB_FLAG
if not MATPLOTLIB_FLAG:
import matplotlib
matplotlib.use("Agg")
MATPLOTLIB_FLAG = True
mpl_logger = logging.getLogger("matplotlib")
mpl_logger.setLevel(logging.WARNING)
import matplotlib.pylab as plt
import numpy as np
fig, ax = plt.subplots(figsize=(6, 4))
im = ax.imshow(
alignment.transpose(), aspect="auto", origin="lower", interpolation="none"
)
fig.colorbar(im, ax=ax)
xlabel = "Decoder timestep"
if info is not None:
xlabel += "\n\n" + info
plt.xlabel(xlabel)
plt.ylabel("Encoder timestep")
plt.tight_layout()
fig.canvas.draw()
data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
plt.close()
return data
def load_wav_to_torch(full_path):
sampling_rate, data = read(full_path)
return torch.FloatTensor(data.astype(np.float32)), sampling_rate
def load_filepaths_and_text(filename, split="|"):
with open(filename, encoding="utf-8") as f:
filepaths_and_text = [line.strip().split(split) for line in f]
return filepaths_and_text
def get_logger(model_dir, filename="train.log"):
global logger
logger = logging.getLogger(os.path.basename(model_dir))
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
if not os.path.exists(model_dir):
os.makedirs(model_dir)
h = logging.FileHandler(os.path.join(model_dir, filename))
h.setLevel(logging.DEBUG)
h.setFormatter(formatter)
logger.addHandler(h)
return logger
def get_steps(model_path):
matches = re.findall(r"\d+", model_path)
return matches[-1] if matches else None
def check_git_hash(model_dir):
source_dir = os.path.dirname(os.path.realpath(__file__))
if not os.path.exists(os.path.join(source_dir, ".git")):
logger.warning(
"{} is not a git repository, therefore hash value comparison will be ignored.".format(
source_dir
)
)
return
cur_hash = subprocess.getoutput("git rev-parse HEAD")
path = os.path.join(model_dir, "githash")
if os.path.exists(path):
saved_hash = open(path).read()
if saved_hash != cur_hash:
logger.warning(
"git hash values are different. {}(saved) != {}(current)".format(
saved_hash[:8], cur_hash[:8]
)
)
else:
open(path, "w").write(cur_hash)

156
style_bert_vits2/models/utils/__init__.py Normal file
View File

@@ -0,0 +1,156 @@
import glob
import logging
import os
import re
import subprocess
import numpy as np
import torch
from scipy.io.wavfile import read
from style_bert_vits2.logging import logger
from style_bert_vits2.models.utils import checkpoints # type: ignore
from style_bert_vits2.models.utils import safetensors # type: ignore
MATPLOTLIB_FLAG = False
def summarize(
writer,
global_step,
scalars={},
histograms={},
images={},
audios={},
audio_sampling_rate=22050,
):
for k, v in scalars.items():
writer.add_scalar(k, v, global_step)
for k, v in histograms.items():
writer.add_histogram(k, v, global_step)
for k, v in images.items():
writer.add_image(k, v, global_step, dataformats="HWC")
for k, v in audios.items():
writer.add_audio(k, v, global_step, audio_sampling_rate)
def is_resuming(dir_path):
# JP-ExtraバージョンではDURがなくWDがあったり変わるため、Gのみで判断する
g_list = glob.glob(os.path.join(dir_path, "G_*.pth"))
# d_list = glob.glob(os.path.join(dir_path, "D_*.pth"))
# dur_list = glob.glob(os.path.join(dir_path, "DUR_*.pth"))
return len(g_list) > 0
def plot_spectrogram_to_numpy(spectrogram):
global MATPLOTLIB_FLAG
if not MATPLOTLIB_FLAG:
import matplotlib
matplotlib.use("Agg")
MATPLOTLIB_FLAG = True
mpl_logger = logging.getLogger("matplotlib")
mpl_logger.setLevel(logging.WARNING)
import matplotlib.pylab as plt
import numpy as np
fig, ax = plt.subplots(figsize=(10, 2))
im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
plt.colorbar(im, ax=ax)
plt.xlabel("Frames")
plt.ylabel("Channels")
plt.tight_layout()
fig.canvas.draw()
data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
plt.close()
return data
def plot_alignment_to_numpy(alignment, info=None):
global MATPLOTLIB_FLAG
if not MATPLOTLIB_FLAG:
import matplotlib
matplotlib.use("Agg")
MATPLOTLIB_FLAG = True
mpl_logger = logging.getLogger("matplotlib")
mpl_logger.setLevel(logging.WARNING)
import matplotlib.pylab as plt
import numpy as np
fig, ax = plt.subplots(figsize=(6, 4))
im = ax.imshow(
alignment.transpose(), aspect="auto", origin="lower", interpolation="none"
)
fig.colorbar(im, ax=ax)
xlabel = "Decoder timestep"
if info is not None:
xlabel += "\n\n" + info
plt.xlabel(xlabel)
plt.ylabel("Encoder timestep")
plt.tight_layout()
fig.canvas.draw()
data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
plt.close()
return data
def load_wav_to_torch(full_path):
sampling_rate, data = read(full_path)
return torch.FloatTensor(data.astype(np.float32)), sampling_rate
def load_filepaths_and_text(filename, split="|"):
with open(filename, encoding="utf-8") as f:
filepaths_and_text = [line.strip().split(split) for line in f]
return filepaths_and_text
def get_logger(model_dir, filename="train.log"):
global logger
logger = logging.getLogger(os.path.basename(model_dir))
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
if not os.path.exists(model_dir):
os.makedirs(model_dir)
h = logging.FileHandler(os.path.join(model_dir, filename))
h.setLevel(logging.DEBUG)
h.setFormatter(formatter)
logger.addHandler(h)
return logger
def get_steps(model_path):
matches = re.findall(r"\d+", model_path)
return matches[-1] if matches else None
def check_git_hash(model_dir):
source_dir = os.path.dirname(os.path.realpath(__file__))
if not os.path.exists(os.path.join(source_dir, ".git")):
logger.warning(
"{} is not a git repository, therefore hash value comparison will be ignored.".format(
source_dir
)
)
return
cur_hash = subprocess.getoutput("git rev-parse HEAD")
path = os.path.join(model_dir, "githash")
if os.path.exists(path):
saved_hash = open(path).read()
if saved_hash != cur_hash:
logger.warning(
"git hash values are different. {}(saved) != {}(current)".format(
saved_hash[:8], cur_hash[:8]
)
)
else:
open(path, "w").write(cur_hash)

194
style_bert_vits2/models/utils/checkpoints.py Normal file
View File

@@ -0,0 +1,194 @@
import glob
import os
import re
from pathlib import Path
from typing import Any, Optional, Union
import torch
from style_bert_vits2.logging import logger
def load_checkpoint(
checkpoint_path: Union[str, Path],
model: torch.nn.Module,
optimizer: Optional[torch.optim.Optimizer] = None,
skip_optimizer: bool = False,
for_infer: bool = False
) -> tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]:
"""
指定されたパスからチェックポイントを読み込み、モデルとオプティマイザーを更新する。
Args:
checkpoint_path (Union[str, Path]): チェックポイントファイルのパス
model (torch.nn.Module): 更新するモデル
optimizer (Optional[torch.optim.Optimizer]): 更新するオプティマイザー。None の場合は更新しない
skip_optimizer (bool): オプティマイザーの更新をスキップするかどうかのフラグ
for_infer (bool): 推論用に読み込むかどうかのフラグ
Returns:
tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]: 更新されたモデルとオプティマイザー、学習率、イテレーション番号
"""
assert os.path.isfile(checkpoint_path)
checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")
iteration = checkpoint_dict["iteration"]
learning_rate = checkpoint_dict["learning_rate"]
logger.info(
f"Loading model and optimizer at iteration {iteration} from {checkpoint_path}"
)
if (
optimizer is not None
and not skip_optimizer
and checkpoint_dict["optimizer"] is not None
):
optimizer.load_state_dict(checkpoint_dict["optimizer"])
elif optimizer is None and not skip_optimizer:
# else: Disable this line if Infer and resume checkpoint,then enable the line upper
new_opt_dict = optimizer.state_dict() # type: ignore
new_opt_dict_params = new_opt_dict["param_groups"][0]["params"]
new_opt_dict["param_groups"] = checkpoint_dict["optimizer"]["param_groups"]
new_opt_dict["param_groups"][0]["params"] = new_opt_dict_params
optimizer.load_state_dict(new_opt_dict) # type: ignore
saved_state_dict = checkpoint_dict["model"]
if hasattr(model, "module"):
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
new_state_dict = {}
for k, v in state_dict.items():
try:
# assert "emb_g" not in k
new_state_dict[k] = saved_state_dict[k]
assert saved_state_dict[k].shape == v.shape, (
saved_state_dict[k].shape,
v.shape,
)
except:
# For upgrading from the old version
if "ja_bert_proj" in k:
v = torch.zeros_like(v)
logger.warning(
f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility"
)
elif "enc_q" in k and for_infer:
continue
else:
logger.error(f"{k} is not in the checkpoint {checkpoint_path}")
new_state_dict[k] = v
if hasattr(model, "module"):
model.module.load_state_dict(new_state_dict, strict=False)
else:
model.load_state_dict(new_state_dict, strict=False)
logger.info("Loaded '{}' (iteration {})".format(checkpoint_path, iteration))
return model, optimizer, learning_rate, iteration
def save_checkpoint(
model: torch.nn.Module,
optimizer: Union[torch.optim.Optimizer, torch.optim.AdamW],
learning_rate: float,
iteration: int,
checkpoint_path: Union[str, Path],
) -> None:
"""
モデルとオプティマイザーの状態を指定されたパスに保存する。
Args:
model (torch.nn.Module): 保存するモデル
optimizer (Union[torch.optim.Optimizer, torch.optim.AdamW]): 保存するオプティマイザー
learning_rate (float): 学習率
iteration (int): イテレーション数
checkpoint_path (Union[str, Path]): 保存先のパス
"""
logger.info(f"Saving model and optimizer state at iteration {iteration} to {checkpoint_path}")
if hasattr(model, "module"):
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
torch.save(
{
"model": state_dict,
"iteration": iteration,
"optimizer": optimizer.state_dict(),
"learning_rate": learning_rate,
},
checkpoint_path,
)
def clean_checkpoints(model_dir_path: Union[str, Path] = "logs/44k/", n_ckpts_to_keep: int = 2, sort_by_time: bool = True) -> None:
"""
指定されたディレクトリから古いチェックポイントを削除して空き容量を確保する
Args:
model_dir_path (Union[str, Path]): モデルが保存されているディレクトリのパス
n_ckpts_to_keep (int): 保持するチェックポイントの数（G_0.pth と D_0.pth を除く）
sort_by_time (bool): True の場合、時間順に削除。False の場合、名前順に削除
"""
ckpts_files = [
f
for f in os.listdir(model_dir_path)
if os.path.isfile(os.path.join(model_dir_path, f))
]
def name_key(_f: str) -> int:
return int(re.compile("._(\\d+)\\.pth").match(_f).group(1)) # type: ignore
def time_key(_f: str) -> float:
return os.path.getmtime(os.path.join(model_dir_path, _f))
sort_key = time_key if sort_by_time else name_key
def x_sorted(_x: str) -> list[str]:
return sorted(
[f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")],
key=sort_key,
)
to_del = [
os.path.join(model_dir_path, fn)
for fn in (
x_sorted("G_")[:-n_ckpts_to_keep]
+ x_sorted("D_")[:-n_ckpts_to_keep]
+ x_sorted("WD_")[:-n_ckpts_to_keep]
+ x_sorted("DUR_")[:-n_ckpts_to_keep]
)
]
def del_info(fn: str) -> None:
return logger.info(f"Free up space by deleting ckpt {fn}")
def del_routine(x: str) -> list[Any]:
return [os.remove(x), del_info(x)]
[del_routine(fn) for fn in to_del]
def get_latest_checkpoint_path(model_dir_path: Union[str, Path], regex: str = "G_*.pth") -> str:
"""
指定されたディレクトリから最新のチェックポイントのパスを取得する
Args:
model_dir_path (Union[str, Path]): モデルが保存されているディレクトリのパス
regex (str): チェックポイントのファイル名の正規表現
Returns:
str: 最新のチェックポイントのパス
"""
f_list = glob.glob(os.path.join(str(model_dir_path), regex))
f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
try:
x = f_list[-1]
except IndexError:
raise ValueError(f"No checkpoint found in {model_dir_path} with regex {regex}")
return x

91
style_bert_vits2/models/utils/safetensors.py Normal file
View File

@@ -0,0 +1,91 @@
from pathlib import Path
from typing import Any, Optional, Union
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from style_bert_vits2.logging import logger
def load_safetensors(
checkpoint_path: Union[str, Path],
model: torch.nn.Module,
for_infer: bool = False,
) -> tuple[torch.nn.Module, Optional[int]]:
"""
指定されたパスから safetensors モデルを読み込み、モデルとイテレーションを返す。
Args:
checkpoint_path (Union[str, Path]): モデルのチェックポイントファイルのパス
model (torch.nn.Module): 読み込む対象のモデル
for_infer (bool): 推論用に読み込むかどうかのフラグ
Returns:
tuple[torch.nn.Module, Optional[int]]: 読み込まれたモデルとイテレーション番号（存在する場合）
"""
tensors: dict[str, Any] = {}
iteration: Optional[int] = None
with safe_open(str(checkpoint_path), framework="pt", device="cpu") as f: # type: ignore
for key in f.keys():
if key == "iteration":
iteration = f.get_tensor(key).item()
tensors[key] = f.get_tensor(key)
if hasattr(model, "module"):
result = model.module.load_state_dict(tensors, strict=False)
else:
result = model.load_state_dict(tensors, strict=False)
for key in result.missing_keys:
if key.startswith("enc_q") and for_infer:
continue
logger.warning(f"Missing key: {key}")
for key in result.unexpected_keys:
if key == "iteration":
continue
logger.warning(f"Unexpected key: {key}")
if iteration is None:
logger.info(f"Loaded '{checkpoint_path}'")
else:
logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})")
return model, iteration
def save_safetensors(
model: torch.nn.Module,
iteration: int,
checkpoint_path: Union[str, Path],
is_half: bool = False,
for_infer: bool = False,
) -> None:
"""
モデルを safetensors 形式で保存する。
Args:
model (torch.nn.Module): 保存するモデル
iteration (int): イテレーション番号
checkpoint_path (Union[str, Path]): 保存先のパス
is_half (bool): モデルを半精度で保存するかどうかのフラグ
for_infer (bool): 推論用に保存するかどうかのフラグ
"""
if hasattr(model, "module"):
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
keys = []
for k in state_dict:
if "enc_q" in k and for_infer:
continue # noqa: E701
keys.append(k)
new_dict = (
{k: state_dict[k].half() for k in keys}
if is_half
else {k: state_dict[k] for k in keys}
)
new_dict["iteration"] = torch.LongTensor([iteration])
logger.info(f"Saved safetensors to {checkpoint_path}")
save_file(new_dict, checkpoint_path)

71
train_ms.py
View File

@@ -248,10 +248,7 @@ def run():
drop_last=False,
collate_fn=collate_fn,
)
if (
"use_noise_scaled_mas" in hps.model.keys()
and hps.model.use_noise_scaled_mas is True
):
if hps.model.use_noise_scaled_mas is True:
logger.info("Using noise scaled MAS for VITS2")
mas_noise_scale_initial = 0.01
noise_scale_delta = 2e-6
@@ -259,10 +256,7 @@ def run():
logger.info("Using normal MAS for VITS1")
mas_noise_scale_initial = 0.0
noise_scale_delta = 0.0
if (
"use_duration_discriminator" in hps.model.keys()
and hps.model.use_duration_discriminator is True
):
if hps.model.use_duration_discriminator is True:
logger.info("Using duration discriminator for VITS2")
net_dur_disc = DurationDiscriminator(
hps.model.hidden_channels,
@@ -271,10 +265,7 @@ def run():
0.1,
gin_channels=hps.model.gin_channels if hps.data.n_speakers != 0 else 0,
).cuda(local_rank)
if (
"use_spk_conditioned_encoder" in hps.model.keys()
and hps.model.use_spk_conditioned_encoder is True
):
if hps.model.use_spk_conditioned_encoder is True:
if hps.data.n_speakers == 0:
raise ValueError(
"n_speakers must be > 0 when using spk conditioned encoder to train multi-speaker model"
@@ -370,31 +361,25 @@ def run():
if utils.is_resuming(model_dir):
if net_dur_disc is not None:
_, _, dur_resume_lr, epoch_str = utils.load_checkpoint(
utils.latest_checkpoint_path(model_dir, "DUR_*.pth"),
_, _, dur_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
utils.checkpoints.get_latest_checkpoint_path(model_dir, "DUR_*.pth"),
net_dur_disc,
optim_dur_disc,
skip_optimizer=(
hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
),
skip_optimizer=hps.train.skip_optimizer,
)
if not optim_dur_disc.param_groups[0].get("initial_lr"):
optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
_, optim_g, g_resume_lr, epoch_str = utils.load_checkpoint(
utils.latest_checkpoint_path(model_dir, "G_*.pth"),
_, optim_g, g_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"),
net_g,
optim_g,
skip_optimizer=(
hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
),
skip_optimizer=hps.train.skip_optimizer,
)
_, optim_d, d_resume_lr, epoch_str = utils.load_checkpoint(
utils.latest_checkpoint_path(model_dir, "D_*.pth"),
_, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"),
net_d,
optim_d,
skip_optimizer=(
hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
),
skip_optimizer=hps.train.skip_optimizer,
)
if not optim_g.param_groups[0].get("initial_lr"):
optim_g.param_groups[0]["initial_lr"] = g_resume_lr
@@ -404,21 +389,21 @@ def run():
epoch_str = max(epoch_str, 1)
# global_step = (epoch_str - 1) * len(train_loader)
global_step = int(
utils.get_steps(utils.latest_checkpoint_path(model_dir, "G_*.pth"))
utils.get_steps(utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"))
)
logger.info(
f"******************Found the model. Current epoch is {epoch_str}, gloabl step is {global_step}*********************"
)
else:
try:
_ = utils.load_safetensors(
_ = utils.safetensors.load_safetensors(
os.path.join(model_dir, "G_0.safetensors"), net_g
)
_ = utils.load_safetensors(
_ = utils.safetensors.load_safetensors(
os.path.join(model_dir, "D_0.safetensors"), net_d
)
if net_dur_disc is not None:
_ = utils.load_safetensors(
_ = utils.safetensors.load_safetensors(
os.path.join(model_dir, "DUR_0.safetensors"), net_dur_disc
)
logger.info("Loaded the pretrained models.")
@@ -511,14 +496,16 @@ def run():
if epoch == hps.train.epochs:
# Save the final models
utils.save_checkpoint(
assert optim_g is not None
utils.checkpoints.save_checkpoint(
net_g,
optim_g,
hps.train.learning_rate,
epoch,
os.path.join(model_dir, "G_{}.pth".format(global_step)),
)
utils.save_checkpoint(
assert optim_d is not None
utils.checkpoints.save_checkpoint(
net_d,
optim_d,
hps.train.learning_rate,
@@ -526,14 +513,15 @@ def run():
os.path.join(model_dir, "D_{}.pth".format(global_step)),
)
if net_dur_disc is not None:
utils.save_checkpoint(
assert optim_dur_disc is not None
utils.checkpoints.save_checkpoint(
net_dur_disc,
optim_dur_disc,
hps.train.learning_rate,
epoch,
os.path.join(model_dir, "DUR_{}.pth".format(global_step)),
)
utils.save_safetensors(
utils.safetensors.save_safetensors(
net_g,
epoch,
os.path.join(
@@ -804,14 +792,15 @@ def train_and_evaluate(
):
if not hps.speedup:
evaluate(hps, net_g, eval_loader, writer_eval)
utils.save_checkpoint(
assert hps.model_dir is not None
utils.checkpoints.save_checkpoint(
net_g,
optim_g,
hps.train.learning_rate,
epoch,
os.path.join(hps.model_dir, "G_{}.pth".format(global_step)),
)
utils.save_checkpoint(
utils.checkpoints.save_checkpoint(
net_d,
optim_d,
hps.train.learning_rate,
@@ -819,7 +808,7 @@ def train_and_evaluate(
os.path.join(hps.model_dir, "D_{}.pth".format(global_step)),
)
if net_dur_disc is not None:
utils.save_checkpoint(
utils.checkpoints.save_checkpoint(
net_dur_disc,
optim_dur_disc,
hps.train.learning_rate,
@@ -828,13 +817,13 @@ def train_and_evaluate(
)
keep_ckpts = config.train_ms_config.keep_ckpts
if keep_ckpts > 0:
utils.clean_checkpoints(
path_to_models=hps.model_dir,
utils.checkpoints.clean_checkpoints(
model_dir_path=hps.model_dir,
n_ckpts_to_keep=keep_ckpts,
sort_by_time=True,
)
# Save safetensors (for inference) to `model_assets/{model_name}`
utils.save_safetensors(
utils.safetensors.save_safetensors(
net_g,
epoch,
os.path.join(

94
train_ms_jp_extra.py
View File

@@ -247,10 +247,7 @@ def run():
drop_last=False,
collate_fn=collate_fn,
)
if (
"use_noise_scaled_mas" in hps.model.keys()
and hps.model.use_noise_scaled_mas is True
):
if hps.model.use_noise_scaled_mas is True:
logger.info("Using noise scaled MAS for VITS2")
mas_noise_scale_initial = 0.01
noise_scale_delta = 2e-6
@@ -258,10 +255,7 @@ def run():
logger.info("Using normal MAS for VITS1")
mas_noise_scale_initial = 0.0
noise_scale_delta = 0.0
if (
"use_duration_discriminator" in hps.model.keys()
and hps.model.use_duration_discriminator is True
):
if hps.model.use_duration_discriminator is True:
logger.info("Using duration discriminator for VITS2")
net_dur_disc = DurationDiscriminator(
hps.model.hidden_channels,
@@ -272,19 +266,13 @@ def run():
).cuda(local_rank)
else:
net_dur_disc = None
if (
"use_wavlm_discriminator" in hps.model.keys()
and hps.model.use_wavlm_discriminator is True
):
if hps.model.use_wavlm_discriminator is True:
net_wd = WavLMDiscriminator(
hps.model.slm.hidden, hps.model.slm.nlayers, hps.model.slm.initial_channel
).cuda(local_rank)
else:
net_wd = None
if (
"use_spk_conditioned_encoder" in hps.model.keys()
and hps.model.use_spk_conditioned_encoder is True
):
if hps.model.use_spk_conditioned_encoder is True:
if hps.data.n_speakers == 0:
raise ValueError(
"n_speakers must be > 0 when using spk conditioned encoder to train multi-speaker model"
@@ -394,15 +382,11 @@ def run():
if utils.is_resuming(model_dir):
if net_dur_disc is not None:
try:
_, _, dur_resume_lr, epoch_str = utils.load_checkpoint(
utils.latest_checkpoint_path(model_dir, "DUR_*.pth"),
_, _, dur_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
utils.checkpoints.get_latest_checkpoint_path(model_dir, "DUR_*.pth"),
net_dur_disc,
optim_dur_disc,
skip_optimizer=(
hps.train.skip_optimizer
if "skip_optimizer" in hps.train
else True
),
skip_optimizer=hps.train.skip_optimizer,
)
if not optim_dur_disc.param_groups[0].get("initial_lr"):
optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
@@ -412,15 +396,11 @@ def run():
print("Initialize dur_disc")
if net_wd is not None:
try:
_, optim_wd, wd_resume_lr, epoch_str = utils.load_checkpoint(
utils.latest_checkpoint_path(model_dir, "WD_*.pth"),
_, optim_wd, wd_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
utils.checkpoints.get_latest_checkpoint_path(model_dir, "WD_*.pth"),
net_wd,
optim_wd,
skip_optimizer=(
hps.train.skip_optimizer
if "skip_optimizer" in hps.train
else True
),
skip_optimizer=hps.train.skip_optimizer,
)
if not optim_wd.param_groups[0].get("initial_lr"):
optim_wd.param_groups[0]["initial_lr"] = wd_resume_lr
@@ -430,21 +410,17 @@ def run():
logger.info("Initialize wavlm")
try:
_, optim_g, g_resume_lr, epoch_str = utils.load_checkpoint(
utils.latest_checkpoint_path(model_dir, "G_*.pth"),
_, optim_g, g_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"),
net_g,
optim_g,
skip_optimizer=(
hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
),
skip_optimizer=hps.train.skip_optimizer,
)
_, optim_d, d_resume_lr, epoch_str = utils.load_checkpoint(
utils.latest_checkpoint_path(model_dir, "D_*.pth"),
_, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"),
net_d,
optim_d,
skip_optimizer=(
hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
),
skip_optimizer=hps.train.skip_optimizer,
)
if not optim_g.param_groups[0].get("initial_lr"):
optim_g.param_groups[0]["initial_lr"] = g_resume_lr
@@ -454,7 +430,7 @@ def run():
epoch_str = max(epoch_str, 1)
# global_step = (epoch_str - 1) * len(train_loader)
global_step = int(
utils.get_steps(utils.latest_checkpoint_path(model_dir, "G_*.pth"))
utils.get_steps(utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"))
)
logger.info(
f"******************Found the model. Current epoch is {epoch_str}, gloabl step is {global_step}*********************"
@@ -468,18 +444,18 @@ def run():
global_step = 0
else:
try:
_ = utils.load_safetensors(
_ = utils.safetensors.load_safetensors(
os.path.join(model_dir, "G_0.safetensors"), net_g
)
_ = utils.load_safetensors(
_ = utils.safetensors.load_safetensors(
os.path.join(model_dir, "D_0.safetensors"), net_d
)
if net_dur_disc is not None:
_ = utils.load_safetensors(
_ = utils.safetensors.load_safetensors(
os.path.join(model_dir, "DUR_0.safetensors"), net_dur_disc
)
if net_wd is not None:
_ = utils.load_safetensors(
_ = utils.safetensors.load_safetensors(
os.path.join(model_dir, "WD_0.safetensors"), net_wd
)
logger.info("Loaded the pretrained models.")
@@ -586,14 +562,16 @@ def run():
scheduler_wd.step()
if epoch == hps.train.epochs:
# Save the final models
utils.save_checkpoint(
assert optim_g is not None
utils.checkpoints.save_checkpoint(
net_g,
optim_g,
hps.train.learning_rate,
epoch,
os.path.join(model_dir, "G_{}.pth".format(global_step)),
)
utils.save_checkpoint(
assert optim_d is not None
utils.checkpoints.save_checkpoint(
net_d,
optim_d,
hps.train.learning_rate,
@@ -601,7 +579,8 @@ def run():
os.path.join(model_dir, "D_{}.pth".format(global_step)),
)
if net_dur_disc is not None:
utils.save_checkpoint(
assert optim_dur_disc is not None
utils.checkpoints.save_checkpoint(
net_dur_disc,
optim_dur_disc,
hps.train.learning_rate,
@@ -609,14 +588,15 @@ def run():
os.path.join(model_dir, "DUR_{}.pth".format(global_step)),
)
if net_wd is not None:
utils.save_checkpoint(
assert optim_wd is not None
utils.checkpoints.save_checkpoint(
net_wd,
optim_wd,
hps.train.learning_rate,
epoch,
os.path.join(model_dir, "WD_{}.pth".format(global_step)),
)
utils.save_safetensors(
utils.safetensors.save_safetensors(
net_g,
epoch,
os.path.join(
@@ -949,14 +929,14 @@ def train_and_evaluate(
):
if not hps.speedup:
evaluate(hps, net_g, eval_loader, writer_eval)
utils.save_checkpoint(
utils.checkpoints.save_checkpoint(
net_g,
optim_g,
hps.train.learning_rate,
epoch,
os.path.join(hps.model_dir, "G_{}.pth".format(global_step)),
)
utils.save_checkpoint(
utils.checkpoints.save_checkpoint(
net_d,
optim_d,
hps.train.learning_rate,
@@ -964,7 +944,7 @@ def train_and_evaluate(
os.path.join(hps.model_dir, "D_{}.pth".format(global_step)),
)
if net_dur_disc is not None:
utils.save_checkpoint(
utils.checkpoints.save_checkpoint(
net_dur_disc,
optim_dur_disc,
hps.train.learning_rate,
@@ -972,7 +952,7 @@ def train_and_evaluate(
os.path.join(hps.model_dir, "DUR_{}.pth".format(global_step)),
)
if net_wd is not None:
utils.save_checkpoint(
utils.checkpoints.save_checkpoint(
net_wd,
optim_wd,
hps.train.learning_rate,
@@ -981,13 +961,13 @@ def train_and_evaluate(
)
keep_ckpts = config.train_ms_config.keep_ckpts
if keep_ckpts > 0:
utils.clean_checkpoints(
path_to_models=hps.model_dir,
utils.checkpoints.clean_checkpoints(
model_dir_path=hps.model_dir,
n_ckpts_to_keep=keep_ckpts,
sort_by_time=True,
)
# Save safetensors (for inference) to `model_assets/{model_name}`
utils.save_safetensors(
utils.safetensors.save_safetensors(
net_g,
epoch,
os.path.join(