2880 lines
140 KiB
Python
2880 lines
140 KiB
Python
"""Voice training pipeline: warm start/resume, piper compatibility fixes, reliability, data filters."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import errno
|
|
import functools
|
|
import io
|
|
import json
|
|
import math
|
|
import os
|
|
import platform
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
import warnings
|
|
from datetime import timedelta
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
|
|
from javis.voices.training_control import EXIT_BUSY, EXIT_DIVERGED, EXIT_STOPPED, process_alive
|
|
|
|
HIGH_ARCH = {
|
|
"resblock": "1", "resblock_kernel_sizes": (3, 7, 11),
|
|
"resblock_dilation_sizes": ((1, 3, 5), (1, 3, 5), (1, 3, 5)),
|
|
"upsample_rates": (8, 8, 2, 2), "upsample_initial_channel": 512,
|
|
"upsample_kernel_sizes": (16, 16, 4, 4),
|
|
}
|
|
|
|
|
|
def _voice_dir(root: Path) -> Path:
|
|
training = root / "v" / "training"
|
|
(training / "wavs").mkdir(parents=True)
|
|
(training / "metadata.csv").write_text("clip_00000|hallo|hallo\n", encoding="utf-8")
|
|
return root / "v"
|
|
|
|
|
|
def _ckpt_dir(voice: Path, version: int = 0) -> Path:
|
|
return (voice / "training" / "piper_lightning" / "lightning_logs" / f"version_{version}"
|
|
/ "checkpoints")
|
|
|
|
|
|
def _checkpoint(path: Path, *, step: int = 0, nan: bool = False) -> Path:
|
|
"""A minimal Lightning checkpoint at *path* (with NaN weights if *nan*)."""
|
|
import torch
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
weights = torch.full((2,), float("nan")) if nan else torch.zeros(2)
|
|
torch.save({"state_dict": {"model_g.w": weights}, "global_step": step,
|
|
"hyper_parameters": HIGH_ARCH}, path)
|
|
return path
|
|
|
|
|
|
def _truncate(path: Path, size: int = 200) -> Path:
|
|
path.write_bytes(path.read_bytes()[:size])
|
|
return path
|
|
|
|
|
|
def _exited_process() -> tuple[int, int]:
|
|
"""``(pid, start time)`` of a process that has exited."""
|
|
from javis.voices.training_control import process_start_time
|
|
proc = subprocess.Popen([sys.executable, "-c", "pass"])
|
|
start = process_start_time(proc.pid) or 0
|
|
proc.wait()
|
|
return proc.pid, start
|
|
|
|
|
|
def _captured(run):
|
|
"""``(run(), stdout, stderr)``, with warnings suppressed."""
|
|
out, err = io.StringIO(), io.StringIO()
|
|
with warnings.catch_warnings(), contextlib.redirect_stdout(out), \
|
|
contextlib.redirect_stderr(err):
|
|
warnings.simplefilter("ignore")
|
|
result = run()
|
|
return result, out.getvalue(), err.getvalue()
|
|
|
|
|
|
def _toy_training(root: Path, *, epochs: int, callbacks=(), on_step=None, resumable=False):
|
|
"""``(model, trainer)``: a tiny run with piper's checkpoint callbacks, saving to *root*.
|
|
|
|
val_mel gets worse every epoch, so only epoch 0 makes the top 1; val_mos is
|
|
never logged (as when the MOS predictor can't be loaded). Four clips, one
|
|
per batch; ``model.items`` lists the clips trained. *resumable*: batched
|
|
like piper_compat's training, whose position within an epoch is saved.
|
|
"""
|
|
import lightning as L
|
|
import torch
|
|
from lightning.pytorch.callbacks import ModelCheckpoint
|
|
from torch.utils.data import DataLoader
|
|
|
|
class Toy(L.LightningModule):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.layer = torch.nn.Linear(2, 1)
|
|
self.steps: list[tuple[int, int]] = [] # (epoch, batch)
|
|
self.items: list[int] = []
|
|
|
|
def training_step(self, batch, batch_idx):
|
|
self.steps.append((self.current_epoch, batch_idx))
|
|
self.items += [int(clip) for clip in batch[:, 0]]
|
|
if on_step is not None:
|
|
on_step(self, batch_idx)
|
|
return self.layer(batch).pow(2).mean()
|
|
|
|
def validation_step(self, batch, batch_idx):
|
|
self.log("val_mel", 1.0 + self.current_epoch)
|
|
|
|
def configure_optimizers(self):
|
|
return torch.optim.SGD(self.parameters(), lr=0.01)
|
|
|
|
def train_dataloader(self):
|
|
clips = torch.arange(4.0).unsqueeze(1).repeat(1, 2) # clip i: [i, i]
|
|
if not resumable:
|
|
return DataLoader(clips, batch_size=1)
|
|
from javis.voices import piper_compat as pc
|
|
return pc.ResumableDataLoader(
|
|
clips, batch_sampler=pc.LengthBucketBatchSampler([3, 1, 4, 2], 1, drop_last=False),
|
|
batches_done=functools.partial(pc._batches_done, self))
|
|
|
|
def val_dataloader(self):
|
|
return DataLoader(torch.ones(1, 2), batch_size=1)
|
|
|
|
checkpoints = [ # as in piper.train.__main__
|
|
ModelCheckpoint(dirpath=root, monitor="val_mel", mode="min", save_top_k=1, save_last=True,
|
|
filename="epoch={epoch}-val_mel={val_mel:.4f}",
|
|
auto_insert_metric_name=False),
|
|
ModelCheckpoint(dirpath=root, monitor="val_mos", mode="max", save_top_k=1,
|
|
save_last=False, filename="epoch={epoch}-val_mos={val_mos:.4f}",
|
|
auto_insert_metric_name=False),
|
|
]
|
|
trainer = L.Trainer(accelerator="cpu", max_epochs=epochs, callbacks=[*checkpoints, *callbacks],
|
|
logger=False, num_sanity_val_steps=0, enable_progress_bar=False,
|
|
enable_model_summary=False)
|
|
return Toy(), trainer
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _checkpoint_patches(pc):
|
|
"""piper_compat's checkpoint patches, undone afterwards: they patch Lightning itself."""
|
|
from lightning.fabric.plugins.io import torch_io
|
|
from lightning.pytorch.callbacks import ModelCheckpoint
|
|
with contextlib.ExitStack() as stack:
|
|
for owner, name in ((torch_io, "_atomic_save"),
|
|
(torch_io.TorchCheckpointIO, "remove_checkpoint"),
|
|
(ModelCheckpoint, "on_train_epoch_end"),
|
|
(ModelCheckpoint, "on_validation_end"),
|
|
(ModelCheckpoint, "_save_topk_checkpoint")):
|
|
stack.enter_context(patch.object(owner, name, getattr(owner, name)))
|
|
stack.enter_context(patch.object(pc, "_missing_monitors", set()))
|
|
pc._safe_checkpoint_files()
|
|
pc._checkpoint_every_epoch()
|
|
yield
|
|
|
|
|
|
class TrainingStartTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
try:
|
|
import torch # noqa: F401
|
|
except ImportError:
|
|
self.skipTest("torch not installed")
|
|
from javis.voices import piper_train
|
|
self.pt = piper_train
|
|
self.cmds: list[list[str]] = []
|
|
# What the next training processes do: an exit code, or a callable that
|
|
# gets _run's keyword arguments and returns one. After those: crash (1).
|
|
self.outcomes: list = []
|
|
|
|
def fake_run(cmd, **kw):
|
|
self.cmds.append(cmd)
|
|
if "export" in cmd:
|
|
Path(cmd[cmd.index("--output-file") + 1]).write_bytes(b"onnx")
|
|
if "--info-file" in cmd: # the record piper_compat's export writes
|
|
ckpt = Path(cmd[cmd.index("--checkpoint") + 1])
|
|
Path(cmd[cmd.index("--info-file") + 1]).write_text(json.dumps({
|
|
"checkpoint": os.path.abspath(ckpt),
|
|
"checkpoint_mtime_ns": ckpt.stat().st_mtime_ns, "epoch": 0}),
|
|
encoding="utf-8")
|
|
return 0
|
|
outcome = self.outcomes.pop(0) if self.outcomes else 1
|
|
return outcome(**kw) if callable(outcome) else outcome
|
|
|
|
patches = [patch.object(piper_train, "_run", fake_run),
|
|
patch.object(piper_train, "MIN_FREE_BYTES", 0),
|
|
patch.object(piper_train, "RESTART_DELAY_SECONDS", 0)]
|
|
for p in patches:
|
|
p.start()
|
|
self.addCleanup(p.stop)
|
|
self.addCleanup(piper_train._STOP.clear)
|
|
|
|
def _train(self, voice: Path, **kw) -> tuple[int, str, str]:
|
|
return _captured(lambda: self.pt.run_pipeline(voice, **kw))
|
|
|
|
def _fits(self) -> list[list[str]]:
|
|
return [cmd for cmd in self.cmds if "fit" in cmd]
|
|
|
|
def _ckpt(self, path: Path, **hp) -> Path:
|
|
import torch
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save({"state_dict": {"w": torch.zeros(1)}, "hyper_parameters": hp}, path)
|
|
return path
|
|
|
|
def test_first_run_warm_starts_from_base_with_its_architecture(self) -> None:
|
|
import yaml
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
self._ckpt(voice / "training" / "piper_base" / "epoch=2665-step=1.ckpt",
|
|
dataset=[Path("x")], **HIGH_ARCH)
|
|
self.pt.run_pipeline(voice)
|
|
cmd = self.cmds[0]
|
|
self.assertIn("javis.voices.piper_compat", cmd)
|
|
self.assertIn("--model.warmstart_ckpt", cmd)
|
|
self.assertNotIn("--ckpt_path", cmd)
|
|
arch = yaml.safe_load(Path(cmd[cmd.index("--config") + 1]).read_text(encoding="utf-8"))
|
|
self.assertEqual(arch["model"]["resblock"], "1")
|
|
self.assertEqual(arch["model"]["upsample_rates"], [8, 8, 2, 2])
|
|
self.assertFalse((voice / "training" / "warmstart.pt").exists()) # removed after fit
|
|
|
|
def test_later_runs_resume_own_checkpoint_and_ignore_base(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
base = self._ckpt(voice / "training" / "piper_base" / "epoch=1.ckpt", **HIGH_ARCH)
|
|
last = self._ckpt(voice / "training" / "piper_lightning" / "lightning_logs"
|
|
/ "version_0" / "checkpoints" / "last.ckpt", **HIGH_ARCH)
|
|
self.pt.run_pipeline(voice, base_checkpoint=base)
|
|
cmd = self.cmds[0]
|
|
self.assertEqual(cmd[cmd.index("--ckpt_path") + 1], str(last))
|
|
self.assertNotIn("--model.warmstart_ckpt", cmd)
|
|
self.assertEqual(cmd[cmd.index("--trainer.precision") + 1], "bf16-mixed")
|
|
|
|
def test_refuses_to_resume_a_diverged_checkpoint(self) -> None:
|
|
import torch
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
last = (voice / "training" / "piper_lightning" / "lightning_logs" / "version_0"
|
|
/ "checkpoints" / "last.ckpt")
|
|
last.parent.mkdir(parents=True)
|
|
torch.save({"state_dict": {"model_g.w": torch.zeros(2),
|
|
"model_d.w": torch.full((2,), float("nan"))},
|
|
"hyper_parameters": HIGH_ARCH}, last)
|
|
err = io.StringIO()
|
|
with contextlib.redirect_stderr(err), contextlib.redirect_stdout(io.StringIO()):
|
|
self.assertEqual(self.pt.run_pipeline(voice), 1)
|
|
self.assertEqual(self.cmds, []) # never started training
|
|
self.assertIn("model_d", err.getvalue())
|
|
|
|
def test_passes_over_damaged_checkpoints_to_the_newest_usable_one(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
ckpts = _ckpt_dir(voice)
|
|
usable = _checkpoint(ckpts / "epoch=1-step=100.ckpt", step=100)
|
|
_checkpoint(ckpts / "epoch=2-step=200.ckpt", step=200, nan=True)
|
|
_truncate(_checkpoint(ckpts / "last.ckpt", step=300)) # cut short by a power cut
|
|
self.outcomes = [EXIT_STOPPED]
|
|
_, _, err = self._train(voice)
|
|
cmd = self.cmds[0]
|
|
self.assertEqual(cmd[cmd.index("--ckpt_path") + 1], str(usable))
|
|
self.assertIn("passing over", err)
|
|
self.assertIn("last.ckpt: can't be loaded", err)
|
|
self.assertIn("epoch=2-step=200.ckpt: 1 weight tensors contain NaN/inf", err)
|
|
|
|
def test_crashed_training_restarts_from_its_newest_checkpoint(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
(voice / "training" / "piper_config.json").write_text(
|
|
json.dumps({"phoneme_id_map": {"a": [1]}}), encoding="utf-8")
|
|
|
|
def crash_after_saving(step):
|
|
def training(**kw):
|
|
path = _checkpoint(_ckpt_dir(voice) / f"epoch={step}-step={step}.ckpt",
|
|
step=step)
|
|
later = time.time() + step # saved after everything before it
|
|
os.utime(path, (later, later))
|
|
return 1
|
|
return training
|
|
|
|
self.outcomes = [crash_after_saving(1), crash_after_saving(2), 0]
|
|
rc, _, err = self._train(voice)
|
|
self.assertEqual(rc, 0)
|
|
fits = self._fits()
|
|
self.assertEqual(len(fits), 3)
|
|
self.assertNotIn("--ckpt_path", fits[0])
|
|
for fit, step in zip(fits[1:], (1, 2)):
|
|
self.assertTrue(fit[fit.index("--ckpt_path") + 1].endswith(
|
|
f"epoch={step}-step={step}.ckpt"))
|
|
self.assertEqual(err.count("Training crashed (exit 1)"), 2)
|
|
export = self.cmds[-1]
|
|
self.assertTrue(export[export.index("--checkpoint") + 1].endswith("epoch=2-step=2.ckpt"))
|
|
self.assertEqual((voice / "piper_model.onnx").read_bytes(), b"onnx")
|
|
|
|
def test_gives_up_on_crashes_that_save_nothing(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
self.outcomes = [7] * 5
|
|
rc, _, err = self._train(voice)
|
|
self.assertEqual(rc, 7)
|
|
self.assertEqual(len(self.cmds), self.pt.MAX_FAILURES_WITHOUT_PROGRESS) # no export
|
|
self.assertIn("not restarting it again", err)
|
|
self.assertFalse((voice / "training" / ".training.lock").exists()) # released
|
|
|
|
def test_diverged_or_stopped_training_is_not_restarted(self) -> None:
|
|
def stop_requested(**kw):
|
|
self.pt.request_stop() # e.g. Javis quits during a scheduled run
|
|
return 1
|
|
|
|
for outcome, expected in ((EXIT_DIVERGED, EXIT_DIVERGED), (EXIT_STOPPED, EXIT_STOPPED),
|
|
(stop_requested, EXIT_STOPPED)):
|
|
with self.subTest(outcome=outcome), tempfile.TemporaryDirectory() as d:
|
|
self.cmds.clear()
|
|
self.outcomes = [outcome, 0]
|
|
rc, _, _ = self._train(_voice_dir(Path(d)))
|
|
self.assertEqual(rc, expected)
|
|
self.assertEqual(len(self.cmds), 1) # neither restarted nor exported
|
|
|
|
def test_a_stopped_run_lets_the_voice_speak_as_trained_so_far(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
(voice / "training" / "piper_config.json").write_text(
|
|
json.dumps({"phoneme_id_map": {"a": [1]}}), encoding="utf-8")
|
|
onnx, last = voice / "piper_model.onnx", _ckpt_dir(voice) / "last.ckpt"
|
|
|
|
def stopped(**kw): # Ctrl+C: training saved its progress
|
|
_checkpoint(last, step=7)
|
|
return EXIT_STOPPED
|
|
|
|
self.outcomes = [stopped]
|
|
rc, out, _ = self._train(voice)
|
|
self.assertEqual((rc, len(self.cmds)), (EXIT_STOPPED, 2))
|
|
export = self.cmds[1]
|
|
self.assertIn("export", export)
|
|
self.assertEqual(export[export.index("--checkpoint") + 1], str(last))
|
|
self.assertEqual(onnx.read_bytes(), b"onnx")
|
|
self.assertTrue(onnx.with_suffix(".onnx.json").is_file())
|
|
self.assertIn("Voice updated: it speaks as trained up to epoch 0.", out)
|
|
|
|
self.cmds.clear() # stopped again before saving anything new: nothing to export
|
|
self.outcomes = [EXIT_STOPPED]
|
|
rc, out, _ = self._train(voice)
|
|
self.assertEqual((rc, len(self.cmds)), (EXIT_STOPPED, 1))
|
|
self.assertNotIn("Voice updated", out)
|
|
|
|
def javis_quits(**kw): # no time for an export then
|
|
_checkpoint(last, step=8)
|
|
later = last.stat().st_mtime_ns + 10**9
|
|
os.utime(last, ns=(later, later))
|
|
self.pt.request_stop()
|
|
return EXIT_STOPPED
|
|
|
|
self.cmds.clear()
|
|
self.outcomes = [javis_quits]
|
|
rc, out, _ = self._train(voice)
|
|
self.assertEqual((rc, len(self.cmds)), (EXIT_STOPPED, 1))
|
|
self.assertFalse(self.pt._exported_from(last, onnx))
|
|
|
|
def test_training_keeps_the_voice_close_behind_its_progress(self) -> None:
|
|
polls = []
|
|
|
|
def training(rc):
|
|
def run(on_poll, **kw):
|
|
polls.append(on_poll)
|
|
return rc
|
|
return run
|
|
|
|
with tempfile.TemporaryDirectory() as d, \
|
|
patch.object(self.pt, "SnapshotExporter") as exporter_class:
|
|
voice = _voice_dir(Path(d))
|
|
self.outcomes = [training(1), training(EXIT_STOPPED)] # crashed, restarted, stopped
|
|
rc, _, _ = self._train(voice, snapshot_minutes=5)
|
|
self.assertEqual(rc, EXIT_STOPPED)
|
|
exporter_class.assert_called_once_with(voice, 5)
|
|
exporter = exporter_class.return_value
|
|
self.assertEqual(polls, [exporter.poll] * 2)
|
|
exporter.stop.assert_called_once_with() # an export still running is abandoned
|
|
|
|
def test_stop_during_the_restart_pause(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d, \
|
|
patch.object(self.pt, "_pause", return_value=False) as pause:
|
|
rc, out, _ = self._train(_voice_dir(Path(d)))
|
|
self.assertEqual(rc, EXIT_STOPPED)
|
|
pause.assert_called_once_with(0)
|
|
self.assertEqual(len(self.cmds), 1)
|
|
self.assertIn("Training stopped", out)
|
|
|
|
def test_a_voice_is_trained_by_one_run_at_a_time(self) -> None:
|
|
from javis.voices.training_control import TrainingLock
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
other_run = TrainingLock(voice) # held by a live process: this one
|
|
self.assertTrue(other_run.acquire())
|
|
try:
|
|
rc, _, err = self._train(voice)
|
|
finally:
|
|
other_run.release()
|
|
self.assertEqual(rc, EXIT_BUSY)
|
|
self.assertEqual(self.cmds, [])
|
|
self.assertIn("being trained already", err)
|
|
|
|
locks = []
|
|
|
|
def training(on_start, **kw):
|
|
on_start(os.getpid()) # stands in for the training process
|
|
locks.append(json.loads(other_run.path.read_text(encoding="utf-8")))
|
|
return EXIT_STOPPED
|
|
|
|
self.outcomes = [training]
|
|
self.assertEqual(self._train(voice)[0], EXIT_STOPPED)
|
|
self.assertEqual(len(locks[0]["processes"]), 2) # the launcher and its trainer
|
|
self.assertFalse(other_run.path.exists()) # released
|
|
|
|
def test_leftovers_of_interrupted_saves_are_removed(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
last = _checkpoint(_ckpt_dir(voice) / "last.ckpt", step=5)
|
|
partial = last.with_name("last.ckpt.partial")
|
|
partial.write_bytes(b"cut short")
|
|
present_at_start = []
|
|
|
|
def killed_while_saving(**kw):
|
|
present_at_start.append(partial.exists())
|
|
partial.write_bytes(b"cut short again")
|
|
return 1
|
|
|
|
def training(**kw):
|
|
present_at_start.append(partial.exists())
|
|
return EXIT_STOPPED
|
|
|
|
self.outcomes = [killed_while_saving, training]
|
|
self._train(voice)
|
|
self.assertEqual(present_at_start, [False, False]) # at the start and the restart
|
|
self.assertFalse(partial.exists())
|
|
self.assertTrue(last.exists())
|
|
|
|
def test_checkpoint_problem_spots_divergence(self) -> None:
|
|
import torch
|
|
problem = self.pt._checkpoint_problem
|
|
healthy = {"state_dict": {"model_g.w": torch.zeros(2), "steps": torch.tensor(3)},
|
|
"optimizer_states": [{"state": {0: {"exp_avg": torch.zeros(2)}}}],
|
|
"MixedPrecision": {"scale": 1024.0}}
|
|
self.assertIsNone(problem(healthy, "16-mixed"))
|
|
self.assertIn("model_g", problem(
|
|
{"state_dict": {"model_g.w": torch.tensor([1.0, float("inf")])}}, "32"))
|
|
self.assertIn("optimizer 1", problem(
|
|
{"state_dict": {}, "optimizer_states": [
|
|
{"state": {}}, {"state": {0: {"exp_avg_sq": torch.tensor([float("nan")])}}}]},
|
|
"bf16-mixed"))
|
|
collapsed = {"state_dict": {"model_g.w": torch.zeros(2)},
|
|
"MixedPrecision": {"scale": torch.tensor(0.0)}}
|
|
self.assertIn("loss scale", problem(collapsed, "16-mixed"))
|
|
self.assertIsNone(problem(collapsed, "bf16-mixed")) # no scaler is used then
|
|
|
|
def test_linux_checkpoint_with_posixpath_loads(self) -> None:
|
|
import pathlib
|
|
import torch
|
|
|
|
class FakePosixPath: # pickled as pathlib.PosixPath, like a Linux checkpoint
|
|
def __init__(self, p):
|
|
self.p = p
|
|
|
|
def __reduce__(self):
|
|
return (FakePosixPath, (self.p,))
|
|
|
|
FakePosixPath.__module__ = "pathlib"
|
|
FakePosixPath.__name__ = FakePosixPath.__qualname__ = "PosixPath"
|
|
with tempfile.TemporaryDirectory() as d:
|
|
p = Path(d) / "linux.ckpt"
|
|
with patch.object(pathlib, "PosixPath", FakePosixPath):
|
|
torch.save({"hyper_parameters": {"dataset": [FakePosixPath("/ssd/x")]}}, p)
|
|
ckpt = self.pt._load_checkpoint(p)
|
|
self.assertEqual(str(ckpt["hyper_parameters"]["dataset"][0]), "/ssd/x")
|
|
|
|
|
|
class CheckpointPickerTests(unittest.TestCase):
|
|
"""Training continues from the most advanced checkpoint that is usable."""
|
|
|
|
def setUp(self) -> None:
|
|
try:
|
|
import torch # noqa: F401
|
|
except ImportError:
|
|
self.skipTest("torch not installed")
|
|
from javis.voices import piper_train
|
|
self.pick = piper_train._newest_usable_checkpoint
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.root = Path(d.name)
|
|
|
|
def test_training_progress_beats_a_stale_last_checkpoint(self) -> None:
|
|
newest = _checkpoint(self.root / "version_0" / "epoch=9-step=900.ckpt", step=900)
|
|
# last.ckpt of a later run that crashed early: a newer file, but less trained
|
|
_checkpoint(self.root / "version_1" / "last.ckpt", step=500)
|
|
path, arch, problems = self.pick(self.root, "bf16-mixed")
|
|
self.assertEqual((path, problems), (newest, []))
|
|
self.assertEqual(arch["resblock"], "1")
|
|
last = _checkpoint(self.root / "version_0" / "last.ckpt", step=900)
|
|
self.assertEqual(self.pick(self.root, "bf16-mixed")[0], last) # equal progress
|
|
|
|
def test_unusable_checkpoints_are_passed_over_and_reported(self) -> None:
|
|
import torch
|
|
usable = _checkpoint(self.root / "a.ckpt", step=1)
|
|
_checkpoint(self.root / "b.ckpt", step=3, nan=True)
|
|
_truncate(_checkpoint(self.root / "c.ckpt", step=4))
|
|
(self.root / "d.ckpt").write_bytes(b"")
|
|
torch.save({"weights": torch.zeros(1)}, self.root / "e.ckpt") # not from Lightning
|
|
path, _, problems = self.pick(self.root, "32")
|
|
self.assertEqual(path, usable)
|
|
self.assertEqual(len(problems), 4)
|
|
self.assertIn("b.ckpt: 1 weight tensors contain NaN/inf", problems[-1])
|
|
usable.unlink() # not left mapped (Windows can't delete a mapped file)
|
|
path, _, problems = self.pick(self.root, "32")
|
|
self.assertIsNone(path)
|
|
self.assertEqual(len(problems), 4)
|
|
|
|
|
|
class TrainingControlTests(unittest.TestCase):
|
|
"""training_control: one training run per voice, and its process helpers."""
|
|
|
|
def setUp(self) -> None:
|
|
from javis.voices import training_control
|
|
self.tc = training_control
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.tmp = Path(d.name)
|
|
self.voice = self.tmp / "voice"
|
|
self.lock_file = training_control.lock_path(self.voice)
|
|
|
|
def _write_lock(self, content, *, age: float = 0.0) -> None:
|
|
self.lock_file.parent.mkdir(parents=True, exist_ok=True)
|
|
text = content if isinstance(content, str) else json.dumps(content)
|
|
self.lock_file.write_text(text, encoding="utf-8")
|
|
if age:
|
|
then = time.time() - age
|
|
os.utime(self.lock_file, (then, then))
|
|
|
|
def test_one_run_per_voice(self) -> None:
|
|
first, second = self.tc.TrainingLock(self.voice), self.tc.TrainingLock(self.voice)
|
|
self.assertTrue(first.acquire())
|
|
self.assertTrue(self.tc.training_in_progress(self.voice))
|
|
self.assertFalse(second.acquire())
|
|
second.release() # not its lock
|
|
self.assertTrue(self.lock_file.exists())
|
|
first.release()
|
|
self.assertFalse(self.lock_file.exists())
|
|
self.assertFalse(self.tc.training_in_progress(self.voice))
|
|
|
|
def test_the_lock_of_a_run_that_has_exited_is_replaced(self) -> None:
|
|
pid, start = _exited_process()
|
|
self._write_lock({"host": platform.node(), "processes": [{"pid": pid, "start": start}]})
|
|
self.assertFalse(self.tc.training_in_progress(self.voice))
|
|
lock = self.tc.TrainingLock(self.voice)
|
|
self.assertTrue(lock.acquire())
|
|
info = json.loads(self.lock_file.read_text(encoding="utf-8"))
|
|
self.assertEqual(info["processes"], [self.tc.this_process()])
|
|
lock.release()
|
|
|
|
def test_the_lock_lasts_while_the_training_process_saves(self) -> None:
|
|
pid, start = _exited_process() # the launcher was killed; its trainer still runs
|
|
self._write_lock({"host": platform.node(), "processes": [
|
|
{"pid": pid, "start": start}, self.tc.this_process()]})
|
|
self.assertTrue(self.tc.training_in_progress(self.voice))
|
|
self.assertFalse(self.tc.TrainingLock(self.voice).acquire())
|
|
|
|
def test_set_trainer_and_release(self) -> None:
|
|
lock = self.tc.TrainingLock(self.voice)
|
|
self.assertTrue(lock.acquire())
|
|
lock.set_trainer(os.getpid()) # stands in for the training process
|
|
info = json.loads(self.lock_file.read_text(encoding="utf-8"))
|
|
self.assertEqual([p["pid"] for p in info["processes"]], [os.getpid()] * 2)
|
|
lock.set_trainer(None)
|
|
info = json.loads(self.lock_file.read_text(encoding="utf-8"))
|
|
self.assertEqual(len(info["processes"]), 1)
|
|
# Taken over by another run meanwhile (this one's lock looked stale): left alone.
|
|
self._write_lock({"host": platform.node(), "processes": [{"pid": 4, "start": 1}]})
|
|
lock.release()
|
|
self.assertTrue(self.lock_file.exists())
|
|
|
|
def test_locks_without_running_processes(self) -> None:
|
|
now = time.time()
|
|
cases = [ # (name, lock file content, its age in seconds, held?)
|
|
("older Javis, recent", {"started_epoch": now - 3600}, 0, True),
|
|
("older Javis, two days old", {"started_epoch": now - 49 * 3600}, 0, False),
|
|
("another computer", {"host": "elsewhere", "started_epoch": now,
|
|
"processes": [{"pid": 4, "start": 1}]}, 0, True),
|
|
("being written", "{", 0, True),
|
|
("garbled", "{", 120, False),
|
|
]
|
|
for name, content, age, held in cases:
|
|
with self.subTest(name):
|
|
self._write_lock(content, age=age)
|
|
self.assertEqual(self.tc.training_in_progress(self.voice), held)
|
|
self.assertEqual(self.tc.TrainingLock(self.voice).acquire(), not held)
|
|
self.lock_file.unlink()
|
|
|
|
def test_process_alive(self) -> None:
|
|
me = self.tc.this_process()
|
|
self.assertTrue(self.tc.process_alive(me["pid"], me["start"]))
|
|
self.assertFalse(self.tc.process_alive(*_exited_process()))
|
|
if sys.platform == "win32": # IDs are soon reused there: the start time tells them apart
|
|
self.assertFalse(self.tc.process_alive(me["pid"], me["start"] + 1))
|
|
|
|
def test_launcher_from_environment(self) -> None:
|
|
for value, expected in (("12:345", (12, 345)), ("", None), ("12", None), ("a:b", None)):
|
|
with self.subTest(value), patch.dict(os.environ, {self.tc.LAUNCHER_ENV: value}):
|
|
self.assertEqual(self.tc.launcher_from_environment(), expected)
|
|
|
|
def test_replace_file_waits_for_the_target_to_be_closed(self) -> None:
|
|
src, dst = self.tmp / "new", self.tmp / "current"
|
|
src.write_text("new")
|
|
dst.write_text("old")
|
|
real_replace, calls = os.replace, []
|
|
|
|
def busy_twice(a, b):
|
|
calls.append(a)
|
|
if len(calls) <= 2:
|
|
raise PermissionError(errno.EACCES, "in use")
|
|
real_replace(a, b)
|
|
|
|
with patch.object(self.tc.os, "replace", busy_twice):
|
|
self.tc.replace_file(src, dst)
|
|
self.assertEqual((dst.read_text(), len(calls)), ("new", 3))
|
|
src.write_text("newer")
|
|
with patch.object(self.tc.os, "replace", side_effect=PermissionError(errno.EACCES, "in use")), \
|
|
self.assertRaises(PermissionError):
|
|
self.tc.replace_file(src, dst, seconds=0)
|
|
with self.assertRaises(FileNotFoundError): # other errors at once
|
|
self.tc.replace_file(self.tmp / "missing", dst)
|
|
|
|
def test_windows_helpers_are_safe_anywhere(self) -> None:
|
|
with self.tc.keep_system_awake(), self.tc.console_selection_cannot_pause():
|
|
pass
|
|
|
|
def test_python_command_starts_this_environment(self) -> None:
|
|
exe, env = self.tc.python_command()
|
|
prefix = subprocess.run(
|
|
[exe, "-c", "import sys, javis, lightning; print(sys.prefix)"],
|
|
env={**os.environ, **env}, cwd=tempfile.gettempdir(), capture_output=True,
|
|
text=True, timeout=120, check=True).stdout.strip()
|
|
self.assertEqual(os.path.normcase(prefix), os.path.normcase(sys.prefix))
|
|
if sys.platform == "win32":
|
|
# torch's Intel Fortran runtime must not abort it when its console window closes.
|
|
self.assertEqual(env.get("FOR_DISABLE_CONSOLE_CTRL_HANDLER"), "1")
|
|
if sys.platform == "win32" and sys.prefix != sys.base_prefix:
|
|
# Not the virtual environment's launcher: closing the console would end it at once.
|
|
self.assertEqual(os.path.normcase(exe), os.path.normcase(sys._base_executable))
|
|
self.assertEqual(env.get("__PYVENV_LAUNCHER__"), sys.executable)
|
|
with patch.object(sys, "platform", "linux"):
|
|
self.assertEqual(self.tc.python_command(), (sys.executable, {}))
|
|
|
|
def test_closing_console_holds_the_process_until_the_progress_is_saved(self) -> None:
|
|
console = self.tc.ConsoleClose(timeout=30)
|
|
self.assertEqual((console.handle(0), console.handle(1)), (False, False)) # Ctrl+C/Break
|
|
self.assertFalse(console.closing.is_set())
|
|
handled = []
|
|
handler = threading.Thread(target=lambda: handled.append(console.handle(5)))
|
|
handler.start()
|
|
self.assertTrue(console.closing.wait(10))
|
|
self.assertEqual(console.reason, "the user is logging off")
|
|
handler.join(0.3)
|
|
self.assertTrue(handler.is_alive()) # Windows waits for the handler to return
|
|
console.done()
|
|
handler.join(10)
|
|
self.assertEqual(handled, [True])
|
|
unsaved = self.tc.ConsoleClose(timeout=0.05) # nothing saves: let Windows go anyway
|
|
self.assertTrue(unsaved.handle(6))
|
|
self.assertEqual(unsaved.reason, "the system is shutting down")
|
|
|
|
@unittest.skipUnless(sys.platform == "win32", "a Windows console handler")
|
|
def test_console_close_handler_is_installed_once(self) -> None:
|
|
register, shutdown_order = MagicMock(return_value=0), MagicMock(return_value=1)
|
|
console = self.tc.ConsoleClose(timeout=0.05)
|
|
with patch.object(self.tc, "_kernel32", lambda: SimpleNamespace(
|
|
SetConsoleCtrlHandler=register, SetProcessShutdownParameters=shutdown_order)), \
|
|
patch.object(console, "_watch_session_end") as watch_session_end:
|
|
console.install() # refused: tried again next time
|
|
register.return_value = 1
|
|
console.install()
|
|
console.install()
|
|
self.assertEqual(register.call_count, 2)
|
|
watch_session_end.assert_called_once_with()
|
|
# Ends before the rest of the session: its data-loader workers are still there.
|
|
level = shutdown_order.call_args.args[0]
|
|
self.assertGreater(level, 0x280) # the default
|
|
self.assertLessEqual(level, 0x3FF) # the applications' range
|
|
handler, add = register.call_args.args
|
|
self.assertTrue(add)
|
|
self.assertEqual(handler(0), 0) # Ctrl+C: left to Python
|
|
self.assertFalse(console.closing.is_set())
|
|
self.assertEqual(handler(2), 1)
|
|
self.assertEqual(console.reason, "the console window was closed")
|
|
with patch.object(sys, "platform", "linux"):
|
|
other = self.tc.ConsoleClose()
|
|
other.install()
|
|
self.assertIsNone(other._handler)
|
|
|
|
@unittest.skipUnless(sys.platform == "win32", "a Windows window")
|
|
def test_logging_off_or_shutting_down_holds_the_process_until_the_progress_is_saved(
|
|
self) -> None:
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
|
|
user32 = ctypes.WinDLL("user32", use_last_error=True)
|
|
user32.SendMessageTimeoutW.argtypes = (
|
|
wintypes.HWND, wintypes.UINT, wintypes.WPARAM, wintypes.LPARAM, wintypes.UINT,
|
|
wintypes.UINT, ctypes.POINTER(ctypes.c_size_t))
|
|
user32.SendMessageTimeoutW.restype = wintypes.LPARAM
|
|
user32.PostMessageW.argtypes = (wintypes.HWND, wintypes.UINT, wintypes.WPARAM,
|
|
wintypes.LPARAM)
|
|
user32.GetAncestor.argtypes = (wintypes.HWND, wintypes.UINT)
|
|
user32.GetAncestor.restype = wintypes.HWND
|
|
user32.GetDesktopWindow.restype = wintypes.HWND
|
|
user32.IsWindowVisible.argtypes = (wintypes.HWND,)
|
|
query_end, end, logoff = 0x0011, 0x0016, 0x80000000
|
|
|
|
def session_window(console):
|
|
threads = set(threading.enumerate())
|
|
console._watch_session_end()
|
|
self.assertTrue(console._window)
|
|
[thread] = set(threading.enumerate()) - threads
|
|
self.addCleanup(thread.join, 10)
|
|
self.addCleanup(user32.PostMessageW, console._window, 0x0010, 0, 0) # WM_CLOSE
|
|
# Top-level, as the session's messages need, but never shown.
|
|
self.assertEqual(user32.GetAncestor(console._window, 1), user32.GetDesktopWindow())
|
|
self.assertFalse(user32.IsWindowVisible(console._window))
|
|
return thread
|
|
|
|
def send(console, message, wparam=0, lparam=0):
|
|
"""As Windows sends it: from another thread, waiting for the answer."""
|
|
result = ctypes.c_size_t()
|
|
if not user32.SendMessageTimeoutW(console._window, message, wparam, lparam, 0,
|
|
20_000, ctypes.byref(result)):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
return result.value
|
|
|
|
console = self.tc.ConsoleClose(timeout=30)
|
|
thread = session_window(console)
|
|
self.assertTrue(thread.daemon) # never keeps the process
|
|
self.assertEqual(send(console, query_end, 0, logoff), 1) # may end
|
|
send(console, end, 0, logoff) # cancelled after all
|
|
self.assertFalse(console.closing.is_set())
|
|
answers = []
|
|
windows = threading.Thread(target=lambda: answers.append(send(console, end, 1, logoff)))
|
|
windows.start()
|
|
self.assertTrue(console.closing.wait(10))
|
|
self.assertEqual(console.reason, "the user is logging off")
|
|
windows.join(0.3)
|
|
self.assertTrue(windows.is_alive()) # Windows waits for the window's answer
|
|
console.done()
|
|
windows.join(10)
|
|
self.assertEqual(answers, [0])
|
|
|
|
unsaved = self.tc.ConsoleClose(timeout=0.05) # nothing saves: let Windows go anyway
|
|
session_window(unsaved)
|
|
self.assertEqual(send(unsaved, end, 1, 0), 0)
|
|
self.assertEqual(unsaved.reason, "the system is shutting down")
|
|
|
|
|
|
class TrainingProcessStopTests(unittest.TestCase):
|
|
"""piper_train._run asks training to save and stop; it kills it only if it must."""
|
|
|
|
# Stand-ins for the training process, run as ``python -c SCRIPT <launcher> <pid file>``:
|
|
# once started, they write the pid of the interpreter itself (with a virtual
|
|
# environment, the process _run starts is a launcher that runs it). TRAINING
|
|
# checks what _run tells it, then stops like StopOnRequest once the stop file exists.
|
|
TRAINING = (
|
|
"import os, sys, time\n"
|
|
"stop = os.environ['JAVIS_PIPER_STOP_FILE']\n"
|
|
"if os.path.exists(stop) or os.environ['JAVIS_PIPER_LAUNCHER'] != sys.argv[1]:\n"
|
|
" sys.exit(8)\n"
|
|
"open(sys.argv[2], 'w').write(str(os.getpid()))\n"
|
|
"deadline = time.monotonic() + 60\n"
|
|
"while not os.path.exists(stop) and time.monotonic() < deadline:\n"
|
|
" time.sleep(0.05)\n"
|
|
"sys.exit(4 if os.path.exists(stop) else 9)\n"
|
|
)
|
|
IGNORES_STOP = "import os, sys, time; open(sys.argv[2], 'w').write(str(os.getpid())); time.sleep(60)"
|
|
|
|
def setUp(self) -> None:
|
|
from javis.voices import piper_train
|
|
self.pt = piper_train
|
|
piper_train._STOP.clear()
|
|
self.addCleanup(piper_train._STOP.clear)
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.stop_file = Path(d.name) / ".stop_training"
|
|
self.pid_file = Path(d.name) / "pid"
|
|
|
|
def _child_pid(self) -> int:
|
|
"""The pid of the stand-in's interpreter, once it has started."""
|
|
deadline = time.monotonic() + 30
|
|
while time.monotonic() < deadline:
|
|
with contextlib.suppress(OSError, ValueError):
|
|
return int(self.pid_file.read_text(encoding="utf-8"))
|
|
time.sleep(0.05)
|
|
raise AssertionError("the training process did not start")
|
|
|
|
def _assert_ended(self, pid: int) -> None:
|
|
deadline = time.monotonic() + 10
|
|
while process_alive(pid) and time.monotonic() < deadline:
|
|
time.sleep(0.05)
|
|
self.assertFalse(process_alive(pid))
|
|
|
|
def _run(self, child: str, *, ctrl_c: int = 0, **kw):
|
|
"""``(rc, child's exit code, pids passed to on_start, seconds)`` of running *child*.
|
|
|
|
Ctrl+C is pressed *ctrl_c* times, at _run's polls once the child has started.
|
|
*kw* goes to _run; what it printed is in ``self.out`` and ``self.err``.
|
|
"""
|
|
from javis.voices.training_control import this_process
|
|
procs, pids, test = [], [], self
|
|
|
|
class Popen(subprocess.Popen):
|
|
waits = 0
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
procs.append(self)
|
|
|
|
def wait(self, timeout=None):
|
|
if timeout == 1.0: # _run's polling
|
|
Popen.waits += 1
|
|
if 2 <= Popen.waits < 2 + ctrl_c:
|
|
test._child_pid()
|
|
raise KeyboardInterrupt
|
|
return super().wait(timeout)
|
|
|
|
me = this_process()
|
|
cmd = [sys.executable, "-c", child, f"{me['pid']}:{me['start']}", str(self.pid_file)]
|
|
started = time.monotonic()
|
|
with patch.object(self.pt.subprocess, "Popen", Popen):
|
|
rc, self.out, self.err = _captured(lambda: self.pt._run(
|
|
cmd, desc="test", stop_file=self.stop_file, on_start=pids.append, **kw))
|
|
return rc, procs[0].returncode if procs else None, pids, time.monotonic() - started
|
|
|
|
def _stop_after_start(self, seconds: float) -> None:
|
|
"""request_stop() *seconds* after the child has started."""
|
|
def stop():
|
|
self._child_pid()
|
|
time.sleep(seconds)
|
|
self.pt.request_stop()
|
|
|
|
thread = threading.Thread(target=stop, daemon=True)
|
|
thread.start()
|
|
self.addCleanup(thread.join, 60)
|
|
|
|
def test_stop_request_lets_training_save_first(self) -> None:
|
|
self.stop_file.touch() # left over from an earlier run: removed before the start
|
|
self._stop_after_start(0.5)
|
|
rc, child_rc, pids, _ = self._run(self.TRAINING)
|
|
self.assertEqual((rc, child_rc), (EXIT_STOPPED, EXIT_STOPPED))
|
|
self.assertEqual(pids, [pids[0], None])
|
|
self.assertFalse(self.stop_file.exists())
|
|
|
|
def test_ctrl_c_stops_training_after_it_saved(self) -> None:
|
|
rc, child_rc, _, _ = self._run(self.TRAINING, ctrl_c=1)
|
|
self.assertEqual((rc, child_rc), (EXIT_STOPPED, EXIT_STOPPED))
|
|
|
|
def test_second_ctrl_c_stops_training_at_once(self) -> None:
|
|
rc, _, pids, seconds = self._run(self.IGNORES_STOP, ctrl_c=2)
|
|
self.assertEqual(rc, EXIT_STOPPED)
|
|
self.assertLess(seconds, 30)
|
|
self.assertFalse(process_alive(pids[0]))
|
|
self._assert_ended(self._child_pid()) # the interpreter too, not just its launcher
|
|
|
|
def test_training_that_does_not_stop_is_killed_after_the_grace_period(self) -> None:
|
|
self._stop_after_start(0.5)
|
|
with patch.object(self.pt, "STOP_GRACE_SECONDS", 1):
|
|
rc, _, pids, seconds = self._run(self.IGNORES_STOP)
|
|
self.assertEqual(rc, EXIT_STOPPED)
|
|
self.assertLess(seconds, 30)
|
|
self.assertFalse(process_alive(pids[0]))
|
|
self._assert_ended(self._child_pid())
|
|
|
|
def test_nothing_starts_after_a_stop_request(self) -> None:
|
|
self.pt.request_stop()
|
|
self.assertEqual(self._run(self.TRAINING)[:3], (EXIT_STOPPED, None, []))
|
|
|
|
def test_on_poll_runs_while_training_runs_until_it_fails(self) -> None:
|
|
polls = []
|
|
|
|
def on_poll(): # e.g. SnapshotExporter.poll
|
|
polls.append(time.monotonic())
|
|
if len(polls) == 2:
|
|
stop = threading.Timer(1.5, self.pt.request_stop)
|
|
stop.start()
|
|
self.addCleanup(stop.cancel)
|
|
raise OSError("disk full")
|
|
|
|
rc, child_rc, _, _ = self._run(self.TRAINING, on_poll=on_poll)
|
|
self.assertEqual((rc, child_rc), (EXIT_STOPPED, EXIT_STOPPED)) # training went on
|
|
self.assertEqual(len(polls), 2)
|
|
self.assertIn("WARNING: OSError('disk full')", self.err)
|
|
|
|
|
|
class OnnxExportTests(unittest.TestCase):
|
|
"""The voice model is replaced only by a complete export."""
|
|
|
|
def setUp(self) -> None:
|
|
from javis.voices import piper_train
|
|
self.pt = piper_train
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.voice = _voice_dir(Path(d.name))
|
|
self.model = self.voice / "piper_model.onnx"
|
|
self.model_config = self.voice / "piper_model.onnx.json"
|
|
self.model.write_bytes(b"old model")
|
|
self.model_config.write_text("old config", encoding="utf-8")
|
|
self.config = self.voice / "training" / "piper_config.json"
|
|
self.config.write_text(json.dumps({"phoneme_id_map": {"a": [1]}}), encoding="utf-8")
|
|
|
|
def _export(self, writes: bytes | None = b"new model", rc: int = 0) -> tuple[int, str]:
|
|
def fake_run(cmd, **kw):
|
|
if writes is not None:
|
|
Path(cmd[cmd.index("--output-file") + 1]).write_bytes(writes)
|
|
return rc
|
|
|
|
with patch.object(self.pt, "_run", fake_run):
|
|
result, _, err = _captured(lambda: self.pt._export_onnx(
|
|
Path("last.ckpt"), self.model, self.voice, desc="test"))
|
|
return result, err
|
|
|
|
def _leftovers(self) -> list[str]:
|
|
return sorted(p.name for p in self.voice.iterdir() if "partial" in p.name)
|
|
|
|
def test_complete_export_replaces_the_model(self) -> None:
|
|
self.assertEqual(self._export()[0], 0)
|
|
self.assertEqual(self.model.read_bytes(), b"new model")
|
|
self.assertEqual(json.loads(self.model_config.read_text(encoding="utf-8")),
|
|
{"phoneme_id_map": {"a": [1]}})
|
|
self.assertEqual(self._leftovers(), [])
|
|
|
|
def test_failed_export_leaves_the_previous_model(self) -> None:
|
|
for writes, rc in ((b"half a mod", 1), (b"", 0), (None, 0), (b"stopped", EXIT_STOPPED)):
|
|
with self.subTest(writes=writes, rc=rc):
|
|
self.assertNotEqual(self._export(writes, rc)[0], 0)
|
|
self.assertEqual(self.model.read_bytes(), b"old model")
|
|
self.assertEqual(self.model_config.read_text(encoding="utf-8"), "old config")
|
|
self.assertEqual(self._leftovers(), [])
|
|
|
|
def test_damaged_voice_config_is_rebuilt(self) -> None:
|
|
self.config.write_text('{"phoneme_id_map": {"a": [', encoding="utf-8") # cut short
|
|
|
|
def rebuild(lightning_dir, dest):
|
|
dest.write_text(json.dumps({"phoneme_id_map": {"b": [2]}}), encoding="utf-8")
|
|
return True
|
|
|
|
with patch.object(self.pt, "_rebuild_voice_config", side_effect=rebuild):
|
|
self.assertEqual(self._export()[0], 0)
|
|
self.assertIn('"b"', self.model_config.read_text(encoding="utf-8"))
|
|
with patch.object(self.pt, "_rebuild_voice_config", return_value=False):
|
|
rc, err = self._export(b"newer model")
|
|
self.assertEqual(rc, 1)
|
|
self.assertIn("could not be rebuilt", err)
|
|
self.assertEqual(self.model.read_bytes(), b"new model")
|
|
self.assertEqual(self._leftovers(), [])
|
|
|
|
|
|
class VoiceFollowsTrainingTests(unittest.TestCase):
|
|
"""While training runs, the voice is re-exported from its progress in the background."""
|
|
|
|
def setUp(self) -> None:
|
|
from javis.voices import piper_train
|
|
self.pt = piper_train
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.voice = _voice_dir(Path(d.name))
|
|
(self.voice / "training" / "piper_config.json").write_text(
|
|
json.dumps({"phoneme_id_map": {"a": [1]}}), encoding="utf-8")
|
|
self.onnx = self.voice / "piper_model.onnx"
|
|
self.last = _ckpt_dir(self.voice) / "last.ckpt"
|
|
self.last.parent.mkdir(parents=True)
|
|
self.last.write_bytes(b"progress")
|
|
self.exported: list[Path] = []
|
|
self.rc, self.seconds = 0, 0 # how the next exports end, and how long they take
|
|
self.now = 1000.0
|
|
for p in (patch.object(piper_train, "_export_command", self._export_command),
|
|
patch.object(piper_train, "time", SimpleNamespace(monotonic=lambda: self.now))):
|
|
p.start()
|
|
self.addCleanup(p.stop)
|
|
|
|
def _export_command(self, ckpt: Path, files) -> tuple[list[str], dict]:
|
|
"""Stands in for piper's export: writes what it would, takes self.seconds, exits self.rc."""
|
|
self.exported.append(ckpt)
|
|
info = {"checkpoint": os.path.abspath(ckpt), "checkpoint_mtime_ns": ckpt.stat().st_mtime_ns,
|
|
"epoch": 2, "global_step": 21, "batches_done": 5, "batches_per_epoch": 8}
|
|
script = "\n".join([
|
|
"import pathlib, sys, time",
|
|
f"pathlib.Path({str(files.partial_onnx)!r}).write_bytes(b'onnx')",
|
|
f"pathlib.Path({str(files.partial_info)!r}).write_text({json.dumps(info)!r})",
|
|
"print('exported', flush=True)",
|
|
f"time.sleep({self.seconds})",
|
|
f"sys.exit({self.rc})"])
|
|
python, env = self.pt.python_command()
|
|
return [python, "-c", script], {**os.environ, **env}
|
|
|
|
def _poll(self, exporter) -> tuple[str, str]:
|
|
"""Poll *exporter* as training does, until the export it starts (if any) has ended."""
|
|
def polls():
|
|
exporter.poll()
|
|
deadline = time.monotonic() + 60
|
|
while exporter._proc is not None and time.monotonic() < deadline:
|
|
time.sleep(0.05)
|
|
exporter.poll()
|
|
_, out, err = _captured(polls)
|
|
return out, err
|
|
|
|
def _save_progress(self) -> None:
|
|
"""Training saves its progress anew."""
|
|
later = self.last.stat().st_mtime_ns + 10**9
|
|
self.last.write_bytes(b"more progress")
|
|
os.utime(self.last, ns=(later, later))
|
|
|
|
def _leftovers(self) -> list[str]:
|
|
return sorted(p.name for p in self.voice.iterdir() if ".partial-" in p.name)
|
|
|
|
def test_the_voice_is_updated_in_the_background(self) -> None:
|
|
exporter = self.pt.SnapshotExporter(self.voice, 15)
|
|
out, _ = self._poll(exporter) # the first export starts right away
|
|
self.assertEqual(self.exported, [self.last])
|
|
self.assertIn("Voice updated: it speaks as trained up to epoch 2, batch 5/8.", out)
|
|
self.assertEqual(self.onnx.read_bytes(), b"onnx")
|
|
self.assertEqual(json.loads(self.onnx.with_suffix(".onnx.json").read_text(encoding="utf-8")),
|
|
{"phoneme_id_map": {"a": [1]}})
|
|
self.assertEqual(self.pt.export_info(self.onnx)["batches_done"], 5)
|
|
self.assertEqual(self._leftovers(), [])
|
|
self.assertIn("exported", exporter.log_path.read_text(encoding="utf-8"))
|
|
|
|
self._save_progress()
|
|
self.now += 15 * 60 - 1 # the next one, once the interval has passed
|
|
self._poll(exporter)
|
|
self.assertEqual(len(self.exported), 1)
|
|
self.now += 1
|
|
self._poll(exporter)
|
|
self.assertEqual(len(self.exported), 2)
|
|
|
|
self.now += 15 * 60 # due, but training saved nothing new meanwhile
|
|
out, _ = self._poll(exporter)
|
|
self.assertEqual((len(self.exported), out), (2, ""))
|
|
self._save_progress()
|
|
self.now += 59 # it looks for new progress every minute until then
|
|
self._poll(exporter)
|
|
self.assertEqual(len(self.exported), 2)
|
|
self.now += 1
|
|
self._poll(exporter)
|
|
self.assertEqual(self.exported, [self.last] * 3)
|
|
|
|
def test_a_failed_export_leaves_the_voice_as_it_was(self) -> None:
|
|
self.onnx.write_bytes(b"old")
|
|
self.rc = 3
|
|
exporter = self.pt.SnapshotExporter(self.voice, 15)
|
|
out, err = self._poll(exporter)
|
|
self.assertEqual(self.onnx.read_bytes(), b"old")
|
|
self.assertNotIn("Voice updated", out)
|
|
self.assertIn("WARNING: could not update the voice with the training progress "
|
|
f"(export exit 3); see {exporter.log_path}", err)
|
|
self.assertEqual(self._leftovers(), [])
|
|
self.now += 15 * 60 # that progress isn't tried again...
|
|
self._poll(exporter)
|
|
self.assertEqual(len(self.exported), 1)
|
|
self.rc = 0
|
|
self._save_progress() # ...the next is
|
|
self.now += 60
|
|
self._poll(exporter)
|
|
self.assertEqual(len(self.exported), 2)
|
|
self.assertEqual(self.onnx.read_bytes(), b"onnx")
|
|
|
|
def test_stopping_training_abandons_a_running_export(self) -> None:
|
|
self.seconds = 60
|
|
exporter = self.pt.SnapshotExporter(self.voice, 15)
|
|
_captured(exporter.poll)
|
|
proc, deadline = exporter._proc, time.monotonic() + 60
|
|
while "exported" not in exporter.log_path.read_text(encoding="utf-8") \
|
|
and time.monotonic() < deadline:
|
|
time.sleep(0.05)
|
|
self.assertEqual(len(self._leftovers()), 2) # its model and record so far
|
|
started = time.monotonic()
|
|
exporter.stop()
|
|
self.assertLess(time.monotonic() - started, 30)
|
|
self.assertIsNotNone(proc.poll())
|
|
self.assertEqual(self._leftovers(), [])
|
|
self.assertFalse(self.onnx.exists())
|
|
exporter.stop() # nothing left to stop
|
|
|
|
def test_no_snapshots_when_turned_off(self) -> None:
|
|
for minutes in (0, -5, float("nan"), float("inf")):
|
|
with self.subTest(minutes=minutes):
|
|
exporter = self.pt.SnapshotExporter(self.voice, minutes)
|
|
self.now += 10**6
|
|
self._poll(exporter)
|
|
self.assertEqual(self.exported, [])
|
|
|
|
def test_what_the_voice_was_exported_from(self) -> None:
|
|
for info, text in (
|
|
(None, "an unknown point of training"),
|
|
({"epoch": "3"}, "an unknown point of training"),
|
|
({"epoch": 3}, "epoch 3"),
|
|
({"epoch": 3, "batches_done": 700, "batches_per_epoch": 1456},
|
|
"epoch 3, batch 700/1456"),
|
|
({"epoch": 3, "batches_done": 1456, "batches_per_epoch": 1456},
|
|
"the end of epoch 3"),
|
|
({"epoch": 3, "batches_done": None, "batches_per_epoch": None}, "epoch 3")):
|
|
with self.subTest(info=info):
|
|
self.assertEqual(self.pt.describe_export(info), text)
|
|
|
|
record = self.pt.export_info_path(self.onnx)
|
|
self.assertEqual(record.name, "piper_model.info.json")
|
|
self.assertIsNone(self.pt.export_info(self.onnx))
|
|
for content in ("{", "[1]"): # cut short; not a record
|
|
record.write_text(content, encoding="utf-8")
|
|
self.assertIsNone(self.pt.export_info(self.onnx))
|
|
|
|
mtime = self.last.stat().st_mtime_ns
|
|
exported = {"checkpoint": str(self.last), "checkpoint_mtime_ns": mtime}
|
|
for info, current in ((exported, False), # the model is gone
|
|
(exported, True),
|
|
({**exported, "checkpoint_mtime_ns": mtime - 1}, False),
|
|
({**exported, "checkpoint": str(self.last.with_name("x.ckpt"))},
|
|
False),
|
|
({**exported, "checkpoint": None}, False),
|
|
({"checkpoint": str(self.last)}, False)):
|
|
with self.subTest(info=info, model=self.onnx.exists()):
|
|
record.write_text(json.dumps(info), encoding="utf-8")
|
|
self.assertEqual(self.pt._exported_from(self.last, self.onnx), current)
|
|
self.onnx.write_bytes(b"onnx")
|
|
|
|
def test_the_newest_saved_progress_is_exported(self) -> None:
|
|
lightning = self.voice / "training" / "piper_lightning"
|
|
newer = _ckpt_dir(self.voice, 1) / "last-v1.ckpt"
|
|
newer.parent.mkdir(parents=True)
|
|
others = [newer.with_name("epoch=3-step=9.ckpt"), newer.with_name("last.ckpt.partial")]
|
|
base = self.last.stat().st_mtime_ns
|
|
for i, path in enumerate([newer, *others], 1): # a top-k save; one being written
|
|
path.write_bytes(b"x")
|
|
os.utime(path, ns=(base + i * 10**9,) * 2)
|
|
self.assertEqual(self.pt._newest_progress(lightning), newer)
|
|
self.assertIsNone(self.pt._newest_progress(self.voice / "missing"))
|
|
|
|
def test_leftovers_of_exports_whose_launcher_died_are_removed(self) -> None:
|
|
dead, live = 4001, 4002
|
|
ours = self.pt._ExportFiles(self.onnx)
|
|
names = {ours.partial_onnx.name: True,
|
|
f"piper_model.partial-{dead}-1.onnx": False,
|
|
f"piper_model.onnx.partial-{dead}-1.json": False,
|
|
f"piper_model.info.partial-{dead}-1.json": False,
|
|
f"piper_model.partial-{live}-1.onnx": True,
|
|
"piper_model.onnx": True, "notes.partial.txt": True}
|
|
for name in names:
|
|
(self.voice / name).write_bytes(b"x")
|
|
with patch.object(self.pt, "process_alive", lambda pid, *_: pid == live):
|
|
self.pt._remove_stale_exports(self.voice)
|
|
self.assertEqual(sorted(p.name for p in self.voice.iterdir() if p.is_file()),
|
|
sorted(name for name, kept in names.items() if kept))
|
|
|
|
|
|
class ScheduledTrainingTests(unittest.IsolatedAsyncioTestCase):
|
|
"""The overnight scheduler reports how a run ended, and lets it save when Javis quits."""
|
|
|
|
async def asyncSetUp(self) -> None:
|
|
from javis.core.config import TrainingConfig
|
|
from javis.voices import piper_train, training_scheduler
|
|
self.ts, self.pt = training_scheduler, piper_train
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.voice = _voice_dir(Path(d.name))
|
|
self.config = SimpleNamespace(training=TrainingConfig())
|
|
self.events: list[tuple[str, dict]] = []
|
|
self.bus = SimpleNamespace(publish=self._publish)
|
|
for p in (patch.object(training_scheduler, "window_end",
|
|
lambda now, *_: now + timedelta(hours=8)),
|
|
patch.object(training_scheduler, "training_in_progress", return_value=False)):
|
|
p.start()
|
|
self.addCleanup(p.stop)
|
|
self.addCleanup(piper_train._STOP.clear)
|
|
|
|
async def _publish(self, event: str, payload: dict) -> None:
|
|
self.events.append((event, payload))
|
|
|
|
async def test_reports_how_the_run_ended(self) -> None:
|
|
for rc, ok, status in ((0, True, "completed"),
|
|
(EXIT_STOPPED, True, "stopped (progress saved)"),
|
|
(EXIT_BUSY, True, "being trained already"),
|
|
(1, False, "FAILED (exit 1)")):
|
|
with self.subTest(rc=rc), patch.object(self.pt, "run_pipeline", return_value=rc) as run:
|
|
self.events.clear()
|
|
await self.ts._run_once(self.config, self.voice, self.bus)
|
|
run.assert_called_once()
|
|
self.assertEqual([event for event, _ in self.events], ["task_start", "task_done"])
|
|
self.assertEqual(self.events[-1][1]["ok"], ok)
|
|
self.assertIn(status, self.events[-1][1]["msg"])
|
|
|
|
async def test_quitting_javis_lets_training_save_and_stop(self) -> None:
|
|
started, finished = threading.Event(), threading.Event()
|
|
|
|
def run_pipeline(*args, **kwargs): # stops when asked to, like the real one
|
|
started.set()
|
|
self.pt._STOP.wait(30)
|
|
finished.set()
|
|
return EXIT_STOPPED
|
|
|
|
with patch.object(self.pt, "run_pipeline", run_pipeline):
|
|
run = asyncio.create_task(self.ts._run_once(self.config, self.voice, self.bus))
|
|
self.assertTrue(await asyncio.to_thread(started.wait, 30))
|
|
run.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await run
|
|
self.assertTrue(self.pt._STOP.is_set())
|
|
self.assertTrue(finished.is_set()) # waited for it to finish
|
|
|
|
|
|
class PiperVoiceReloadTests(unittest.TestCase):
|
|
"""Javis speaks with the active voice's newest model, which training keeps updating."""
|
|
|
|
def setUp(self) -> None:
|
|
try:
|
|
import piper.config # noqa: F401
|
|
except ImportError:
|
|
self.skipTest("piper-tts not installed")
|
|
from javis.core.config import Config, PathsConfig
|
|
from javis.tts import piper_backend
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.config = Config(paths=PathsConfig(base_dir=Path(d.name)))
|
|
self.config.tts.default_voice = "Rufus"
|
|
self.saves = 0
|
|
self.attempts: list[bytes] = [] # the models loading was tried with
|
|
self.loaded: list[bytes] = [] # ...and those that loaded
|
|
self.broken: set[bytes] = set()
|
|
self.loading = threading.Event() # cleared: loading takes until it is set
|
|
self.loading.set()
|
|
patcher = patch.object(piper_backend.PiperBackend, "_load_voice",
|
|
staticmethod(self._load_voice))
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
self.backend = piper_backend.PiperBackend(self.config)
|
|
self.backend.CHECK_SECONDS = 0
|
|
self.addCleanup(self._reloaded)
|
|
self.addCleanup(self.loading.set) # (cleanups run last first)
|
|
|
|
def _export(self, voice: str, model: bytes) -> None:
|
|
"""Training (or an export) writes *voice*'s model."""
|
|
path = self.config.paths.voices_dir / voice / "piper_model.onnx"
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.saves += 1
|
|
mtime = (1_700_000_000 + self.saves) * 10**9
|
|
for file, content in ((Path(str(path) + ".json"), b"{}"), (path, model)):
|
|
file.write_bytes(content)
|
|
os.utime(file, ns=(mtime, mtime))
|
|
|
|
def _load_voice(self, model_path: Path):
|
|
self.assertTrue(self.loading.wait(10))
|
|
model = Path(model_path).read_bytes()
|
|
self.attempts.append(model)
|
|
if model in self.broken:
|
|
raise RuntimeError(f"cannot load {model!r}")
|
|
self.loaded.append(model)
|
|
number = len(self.loaded)
|
|
|
|
def synthesize(text, syn_config=None):
|
|
yield SimpleNamespace(audio_float_array=np.full(2, number, dtype=np.float32))
|
|
|
|
return SimpleNamespace(config=SimpleNamespace(sample_rate=16000 + number),
|
|
synthesize=synthesize)
|
|
|
|
def _speaking(self) -> bytes:
|
|
"""Synthesize something, as Javis does; returns the model that spoke."""
|
|
audio, sample_rate = self.backend._synthesize_sync("Hallo")
|
|
number = int(audio[0])
|
|
self.assertEqual(sample_rate, 16000 + number)
|
|
return self.loaded[number - 1]
|
|
|
|
def _reloaded(self) -> None:
|
|
"""Wait for the model being loaded in the background, if one is."""
|
|
self.assertTrue(self.backend._reloading.acquire(timeout=10))
|
|
self.backend._reloading.release()
|
|
|
|
def test_the_model_training_updates_takes_over(self) -> None:
|
|
self._export("Rufus", b"epoch 1")
|
|
self.backend.load()
|
|
self.assertEqual(self._speaking(), b"epoch 1")
|
|
self.loading.clear() # loading the new model takes a while...
|
|
self._export("Rufus", b"epoch 2")
|
|
self.assertEqual(self._speaking(), b"epoch 1") # ...the old one speaks meanwhile
|
|
self.assertEqual(self._speaking(), b"epoch 1")
|
|
self.loading.set()
|
|
self._reloaded()
|
|
self.assertEqual(self._speaking(), b"epoch 2")
|
|
self.assertEqual(self.attempts, [b"epoch 1", b"epoch 2"]) # loaded once each
|
|
|
|
def test_a_model_that_fails_to_load_leaves_the_voice_as_it_was(self) -> None:
|
|
self._export("Rufus", b"epoch 1")
|
|
self.backend.load()
|
|
self.broken.add(b"damaged")
|
|
self._export("Rufus", b"damaged")
|
|
self._speaking()
|
|
self._reloaded()
|
|
self.assertEqual(self._speaking(), b"epoch 1")
|
|
self._reloaded()
|
|
self.assertEqual(self.attempts, [b"epoch 1", b"damaged"]) # not tried again...
|
|
self._export("Rufus", b"epoch 2")
|
|
self._speaking()
|
|
self._reloaded()
|
|
self.assertEqual(self._speaking(), b"epoch 2") # ...but the next model is
|
|
|
|
def test_activating_another_voice_switches_to_its_model(self) -> None:
|
|
self._export("Rufus", b"rufus")
|
|
self._export("Anna", b"anna")
|
|
self.backend.load()
|
|
self.config.tts.default_voice = "Anna" # e.g. the GUI activates the voice it trains
|
|
self._speaking()
|
|
self._reloaded()
|
|
self.assertEqual(self._speaking(), b"anna")
|
|
|
|
def test_a_voice_whose_training_has_not_exported_a_model_yet(self) -> None:
|
|
self.assertFalse(self.backend.model_available())
|
|
with self.assertRaisesRegex(FileNotFoundError, "Piper ONNX model not found"):
|
|
self.backend.load()
|
|
self.broken.add(b"damaged")
|
|
self._export("Rufus", b"damaged")
|
|
self.assertTrue(self.backend.model_available())
|
|
with self.assertRaisesRegex(RuntimeError, "cannot load"):
|
|
self.backend.load()
|
|
self.assertFalse(self.backend.model_available()) # not tried again...
|
|
self._export("Rufus", b"epoch 1")
|
|
self.assertTrue(self.backend.model_available()) # ...until training exports anew
|
|
self.backend.load()
|
|
self.assertEqual(self._speaking(), b"epoch 1")
|
|
|
|
|
|
class TtsReviveTests(unittest.IsolatedAsyncioTestCase):
|
|
"""Speech output comes on once the voice that had no usable model at startup has one."""
|
|
|
|
async def asyncSetUp(self) -> None:
|
|
from javis.agent import orchestrator
|
|
self.orch = object.__new__(orchestrator.Orchestrator)
|
|
self.orch._stop = asyncio.Event()
|
|
self.orch.tts = None
|
|
self.spawned: list[str] = []
|
|
|
|
def spawn(coro, name=None):
|
|
coro.close()
|
|
self.spawned.append(name)
|
|
|
|
patcher = patch.object(orchestrator, "spawn", spawn)
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
self.backend = MagicMock()
|
|
|
|
async def _revive(self) -> None:
|
|
await asyncio.wait_for(self.orch._revive_tts(self.backend, interval=0.01), 10)
|
|
|
|
async def test_speech_comes_on_once_the_model_loads(self) -> None:
|
|
self.backend.model_available.side_effect = [False, True, True]
|
|
self.backend.load.side_effect = [RuntimeError("damaged"), None]
|
|
await self._revive()
|
|
self.assertIs(self.orch.tts, self.backend)
|
|
self.assertEqual(self.backend.load.call_count, 2)
|
|
self.assertEqual(self.spawned, ["tts-prewarm"])
|
|
|
|
async def test_quitting_javis_ends_the_wait(self) -> None:
|
|
self.backend.model_available.return_value = False
|
|
revive = asyncio.create_task(self._revive())
|
|
await asyncio.sleep(0.05)
|
|
self.orch._stop.set()
|
|
await revive
|
|
loop = asyncio.get_running_loop()
|
|
self.backend.model_available.return_value = True # quitting while the model loads
|
|
self.backend.load.side_effect = lambda: loop.call_soon_threadsafe(self.orch._stop.set)
|
|
self.orch._stop.clear()
|
|
await self._revive()
|
|
self.assertIsNone(self.orch.tts)
|
|
self.assertEqual(self.spawned, [])
|
|
|
|
|
|
class PiperCompatTests(unittest.TestCase):
|
|
def test_int_resblock_builds_high_quality_blocks(self) -> None:
|
|
try:
|
|
import javis.voices.piper_compat # noqa: F401 — applies the fixes
|
|
from piper.train.vits import models, modules
|
|
from pysilero_vad import SileroVoiceActivityDetector
|
|
except ImportError:
|
|
self.skipTest("piper training extras not installed")
|
|
g = models.Generator(192, 1, (3, 7, 11), ((1, 3, 5),) * 3, (8, 8, 2, 2), 64, (16, 16, 4, 4))
|
|
self.assertIsInstance(g.resblocks[0], modules.ResBlock1)
|
|
self.assertTrue(hasattr(SileroVoiceActivityDetector, "process_array"))
|
|
|
|
|
|
def _piper_compat():
|
|
try:
|
|
from javis.voices import piper_compat
|
|
except ImportError:
|
|
raise unittest.SkipTest("piper training extras not installed")
|
|
return piper_compat
|
|
|
|
|
|
def _pristine(name: str):
|
|
"""A fresh copy of a ``piper.train.vits`` module, untouched by piper_compat's patches."""
|
|
import importlib.util
|
|
|
|
from piper.train import vits
|
|
spec = importlib.util.spec_from_file_location(
|
|
f"_pristine_vits_{name}", Path(vits.__file__).with_name(f"{name}.py"))
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class ExportInfoTests(unittest.TestCase):
|
|
"""piper_compat's export records what it exported, and refuses NaN/inf weights."""
|
|
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
from lightning.pytorch.core import saving
|
|
self.saving = saving
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.tmp = Path(d.name)
|
|
|
|
def test_our_options_are_taken_out_of_piper_s(self) -> None:
|
|
for argv, value, rest in (
|
|
(["x", "--info-file", "i.json", "--checkpoint", "c"], "i.json",
|
|
["x", "--checkpoint", "c"]),
|
|
(["x", "--info-file=i.json"], "i.json", ["x"]),
|
|
(["x", "--info-file"], None, ["x", "--info-file"]),
|
|
(["x", "--info-files=i"], None, ["x", "--info-files=i"]),
|
|
(["x"], None, ["x"])):
|
|
with self.subTest(argv=argv):
|
|
argv = list(argv)
|
|
self.assertEqual(self.pc._pop_option(argv, "--info-file"), value)
|
|
self.assertEqual(argv, rest)
|
|
|
|
def test_epoch_progress_of_damaged_records_is_unknown(self) -> None:
|
|
def saved(loaders):
|
|
return {"loops": {"fit_loop": {"state_dict": {"combined_loader": loaders}}}}
|
|
|
|
self.assertEqual(self.pc._epoch_progress(saved([{"batches_done": 7, "order": [0] * 9}])),
|
|
{"batches_done": 7, "batches_per_epoch": 9})
|
|
for checkpoint in ({}, {"loops": None}, saved(None), saved([]), saved([{}]),
|
|
saved([{"batches_done": "x", "order": [0]}]),
|
|
saved([{"batches_done": 1, "order": None}])):
|
|
with self.subTest(checkpoint=checkpoint):
|
|
self.assertEqual(self.pc._epoch_progress(checkpoint), {})
|
|
|
|
def test_the_checkpoint_read_is_recorded(self) -> None:
|
|
import torch
|
|
path, newer = self.tmp / "last.ckpt", self.tmp / "newer.ckpt"
|
|
path.write_bytes(b"saved progress")
|
|
newer.write_bytes(b"newer progress")
|
|
mtime = 1_700_000_000 * 10**9
|
|
os.utime(path, ns=(mtime, mtime))
|
|
read = []
|
|
|
|
def load(source, *args, **kwargs):
|
|
read.append((source.read(), args, kwargs))
|
|
os.replace(newer, path) # training saves anew meanwhile: the file isn't held open
|
|
return {"state_dict": {"model_g.w": torch.zeros(2), "steps": torch.tensor([3])},
|
|
"epoch": 3, "global_step": 700,
|
|
"loops": {"fit_loop": {"state_dict": {"combined_loader": [
|
|
{"batches_done": 700, "order": list(range(1456))}]}}}}
|
|
|
|
info: dict = {}
|
|
with patch.object(self.saving, "pl_load", load):
|
|
self.pc._inspect_exported_checkpoint(info)
|
|
self.assertEqual(self.saving.pl_load(str(path), map_location="cpu")["epoch"], 3)
|
|
self.assertEqual(read, [(b"saved progress", (), {"map_location": "cpu"})])
|
|
self.assertEqual(info, {"checkpoint": os.path.abspath(path), "checkpoint_mtime_ns": mtime,
|
|
"epoch": 3, "global_step": 700, "batches_done": 700,
|
|
"batches_per_epoch": 1456})
|
|
self.assertEqual(path.read_bytes(), b"newer progress")
|
|
|
|
def test_nan_weights_are_not_exported(self) -> None:
|
|
import torch
|
|
path = self.tmp / "diverged.ckpt"
|
|
path.write_bytes(b"x")
|
|
diverged = {"state_dict": {"model_g.a": torch.zeros(2),
|
|
"model_g.b": torch.tensor([1.0, float("nan")])}}
|
|
with patch.object(self.saving, "pl_load", lambda source, *a, **kw: diverged):
|
|
self.pc._inspect_exported_checkpoint({})
|
|
with self.assertRaisesRegex(SystemExit,
|
|
r"diverged\.ckpt holds NaN/inf weights \(model_g\.b\)"):
|
|
self.saving.pl_load(str(path))
|
|
|
|
def test_export_writes_its_record_when_asked_to(self) -> None:
|
|
from piper.train import export_onnx
|
|
ckpt = _checkpoint(self.tmp / "last.ckpt", step=9)
|
|
record = self.tmp / "piper_model.info.partial-1-1.json"
|
|
argvs = []
|
|
|
|
def main(): # piper's export loads the checkpoint
|
|
argvs.append(list(sys.argv))
|
|
self.saving.pl_load(str(ckpt), map_location="cpu")
|
|
|
|
for argv in (["piper_compat", "--checkpoint", str(ckpt)],
|
|
["piper_compat", "--checkpoint", str(ckpt), "--info-file", str(record)]):
|
|
with self.subTest(argv=argv), \
|
|
patch.object(self.saving, "pl_load", self.saving.pl_load), \
|
|
patch.object(self.pc, "_legacy_onnx_export"), \
|
|
patch.object(export_onnx, "main", main), patch.object(sys, "argv", list(argv)):
|
|
started = time.time()
|
|
self.pc._export(sys.argv)
|
|
self.assertEqual(argvs.pop(), ["piper_compat", "--checkpoint", str(ckpt)])
|
|
self.assertEqual(record.exists(), "--info-file" in argv)
|
|
self.assertEqual(sorted(p.name for p in self.tmp.iterdir()), [ckpt.name, record.name])
|
|
info = json.loads(record.read_text(encoding="utf-8"))
|
|
self.assertGreaterEqual(info.pop("exported_at"), started)
|
|
self.assertEqual(info, {"checkpoint": os.path.abspath(ckpt),
|
|
"checkpoint_mtime_ns": ckpt.stat().st_mtime_ns,
|
|
"epoch": None, "global_step": 9})
|
|
|
|
|
|
class MonotonicAlignmentTests(unittest.TestCase):
|
|
"""piper's alignment search runs on piper_compat's port, not a compiled extension."""
|
|
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
self.kernels = {"numba": self.pc._batch_mas(), "numpy": self.pc._mas_numpy}
|
|
|
|
@staticmethod
|
|
def _align(kernel, values, t_ys, t_xs) -> np.ndarray:
|
|
phoneme = np.full(values.shape[:2], -1, np.int32)
|
|
kernel(phoneme, values.astype(np.float32), np.asarray(t_ys, np.int32),
|
|
np.asarray(t_xs, np.int32))
|
|
return phoneme
|
|
|
|
def test_alignment_is_the_best_monotonic_path(self) -> None:
|
|
import itertools
|
|
rng = np.random.default_rng(0)
|
|
for t_y, t_x in [(1, 1), (5, 1), (4, 4), (6, 3), (7, 4), (8, 2)]:
|
|
for _ in range(5):
|
|
scores = rng.integers(-20, 20, (1, t_y + 2, t_x + 1)).astype(np.float32) # + padding
|
|
# every path from phoneme 0 to t_x - 1 that moves on by 0 or 1 phoneme per frame
|
|
best = max(scores[0, np.arange(t_y), np.cumsum((0,) + steps)].sum()
|
|
for steps in itertools.product((0, 1), repeat=t_y - 1)
|
|
if sum(steps) == t_x - 1)
|
|
for name, kernel in self.kernels.items():
|
|
path = self._align(kernel, scores, [t_y], [t_x])[0]
|
|
with self.subTest(kernel=name, t_y=t_y, t_x=t_x):
|
|
self.assertEqual((path[0], path[t_y - 1]), (0, t_x - 1))
|
|
self.assertTrue(set(np.diff(path[:t_y]).tolist()) <= {0, 1})
|
|
self.assertTrue((path[t_y:] == -1).all())
|
|
self.assertEqual(scores[0, np.arange(t_y), path[:t_y]].sum(), best)
|
|
|
|
def test_kernels_agree_on_padded_batches(self) -> None:
|
|
rng = np.random.default_rng(1)
|
|
values = rng.normal(-200, 30, (6, 90, 25)).astype(np.float32)
|
|
t_ys, t_xs = [90, 61, 40, 25, 7, 1], [25, 20, 3, 25, 7, 1]
|
|
paths = {name: self._align(k, values, t_ys, t_xs) for name, k in self.kernels.items()}
|
|
np.testing.assert_array_equal(paths["numba"], paths["numpy"])
|
|
for path, t_y, t_x in zip(paths["numba"], t_ys, t_xs):
|
|
self.assertEqual((path[0], path[t_y - 1]), (0, t_x - 1))
|
|
self.assertTrue((path[t_y:] == -1).all())
|
|
|
|
def test_maximum_path_keeps_piper_contract(self) -> None:
|
|
import torch
|
|
from piper.train.vits import monotonic_align
|
|
neg_cent = torch.randn(3, 40, 12, generator=torch.Generator().manual_seed(2)) * 30 - 200
|
|
t_ys, t_xs = torch.tensor([40, 31, 12]), torch.tensor([12, 9, 12])
|
|
mask = ((torch.arange(40)[None, :, None] < t_ys[:, None, None])
|
|
& (torch.arange(12)[None, None, :] < t_xs[:, None, None])).float()
|
|
before = neg_cent.clone()
|
|
attn = monotonic_align.maximum_path(neg_cent, mask)
|
|
self.assertIs(monotonic_align.maximum_path, self.pc.maximum_path)
|
|
self.assertEqual((attn.shape, attn.dtype), (neg_cent.shape, neg_cent.dtype))
|
|
self.assertTrue(torch.equal(neg_cent, before)) # works on a copy
|
|
self.assertTrue(torch.equal(attn.sum(2), mask[:, :, 0])) # one phoneme per valid frame
|
|
self.assertTrue(torch.equal(attn * mask, attn))
|
|
|
|
def test_stands_in_for_a_missing_extension(self) -> None:
|
|
from piper.train import vits
|
|
name = "piper.train.vits.monotonic_align"
|
|
saved = vits.__dict__.pop("monotonic_align")
|
|
try:
|
|
with patch.dict(sys.modules, {name: None}): # its import now fails
|
|
self.pc._fix_monotonic_alignment()
|
|
from piper.train.vits import monotonic_align # as SynthesizerTrn.forward does
|
|
self.assertIs(monotonic_align.maximum_path, self.pc.maximum_path)
|
|
finally:
|
|
vits.monotonic_align = saved
|
|
|
|
def test_piper_forward_maps_every_frame_to_one_phoneme(self) -> None:
|
|
import torch
|
|
from piper.train.vits.losses import kl_loss
|
|
from piper.train.vits.models import SynthesizerTrn
|
|
torch.manual_seed(0)
|
|
model = SynthesizerTrn(
|
|
n_vocab=12, spec_channels=17, segment_size=8, inter_channels=8, hidden_channels=8,
|
|
filter_channels=16, n_heads=2, n_layers=1, kernel_size=3, p_dropout=0.0,
|
|
resblock="2", resblock_kernel_sizes=(3,), resblock_dilation_sizes=((1, 3),),
|
|
upsample_rates=(2, 2), upsample_initial_channel=8, upsample_kernel_sizes=(4, 4),
|
|
use_sdp=False)
|
|
x_lengths, y_lengths = torch.tensor([9, 5]), torch.tensor([40, 23])
|
|
with torch.no_grad():
|
|
_, _, attn, _, _, y_mask, (_, z_p, m_p, logs_p, _, logs_q) = model(
|
|
torch.randint(1, 12, (2, 9)), x_lengths, torch.randn(2, 17, 40), y_lengths)
|
|
attn = attn.squeeze(1)
|
|
self.assertTrue(torch.equal(attn.sum(2), y_mask.squeeze(1)))
|
|
phoneme = attn.argmax(2)
|
|
for i, (t_x, t_y) in enumerate(zip(x_lengths.tolist(), y_lengths.tolist())):
|
|
self.assertEqual((phoneme[i, 0].item(), phoneme[i, t_y - 1].item()), (0, t_x - 1))
|
|
self.assertTrue(set(torch.diff(phoneme[i, :t_y]).tolist()) <= {0, 1})
|
|
self.assertTrue(torch.isfinite(kl_loss(z_p, logs_q, m_p, logs_p, y_mask)))
|
|
|
|
|
|
class LengthBucketTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
rng = np.random.default_rng(1)
|
|
self.lengths = [int(x) for x in rng.integers(1, 1000, 1000)]
|
|
|
|
def _sampler(self, **kw):
|
|
kw = {"drop_last": False, "seed": 7, **kw}
|
|
return self.pc.LengthBucketBatchSampler(self.lengths, kw.pop("batch_size", 8), **kw)
|
|
|
|
def test_every_clip_once_per_epoch(self) -> None:
|
|
s = self._sampler(batch_size=7)
|
|
batches = list(s)
|
|
self.assertEqual(len(batches), len(s))
|
|
self.assertEqual(sorted(i for b in batches for i in b), list(range(1000)))
|
|
self.assertTrue(all(1 <= len(b) <= 7 for b in batches))
|
|
|
|
def test_drop_last_keeps_full_batches(self) -> None:
|
|
s = self._sampler(batch_size=7, drop_last=True)
|
|
batches = list(s)
|
|
self.assertEqual(len(batches), len(s))
|
|
self.assertEqual(len(s), 1000 // 7)
|
|
self.assertTrue(all(len(b) == 7 for b in batches))
|
|
flat = [i for b in batches for i in b]
|
|
self.assertEqual(len(flat), len(set(flat)))
|
|
|
|
def test_batches_pad_far_less_than_random_ones(self) -> None:
|
|
s = self._sampler()
|
|
rng = np.random.default_rng(2)
|
|
order = rng.permutation(1000).tolist()
|
|
unbucketed = [order[i:i + 8] for i in range(0, 1000, 8)]
|
|
self.assertLess(self.pc.padding_factor(s.batches(0), self.lengths), 1.1)
|
|
self.assertGreater(self.pc.padding_factor(unbucketed, self.lengths), 1.4)
|
|
|
|
def test_order_is_reproducible_per_epoch_and_reshuffled_across_epochs(self) -> None:
|
|
self.assertEqual(self._sampler().batches(3), self._sampler().batches(3))
|
|
self.assertNotEqual(self._sampler().batches(0), self._sampler().batches(1))
|
|
self.assertNotEqual(self._sampler().batches(0), self._sampler(seed=8).batches(0))
|
|
|
|
def test_first_epoch_of_a_run_starts_with_the_longest_batch(self) -> None:
|
|
s = self._sampler()
|
|
s.set_epoch(4)
|
|
first, again = list(s), list(s) # a DataLoader with workers iterates twice
|
|
self.assertEqual(max(self.lengths[i] for i in first[0]), max(self.lengths))
|
|
self.assertCountEqual(first, s.batches(4))
|
|
self.assertEqual(again, first)
|
|
s.set_epoch(5)
|
|
self.assertEqual(list(s), s.batches(5)) # only the run's first epoch
|
|
|
|
def test_longest_batch_first_with_dataloader_workers(self) -> None:
|
|
from torch.utils.data import DataLoader
|
|
loader = DataLoader(range(1000), batch_sampler=self._sampler(), num_workers=1)
|
|
first = next(iter(loader)).tolist()
|
|
self.assertEqual(max(self.lengths[i] for i in first), max(self.lengths))
|
|
|
|
def test_lightning_sets_the_epoch_through_the_dataloader(self) -> None:
|
|
from lightning.fabric.utilities.data import _set_sampler_epoch
|
|
from torch.utils.data import DataLoader
|
|
s = self._sampler(longest_first=False)
|
|
loader = DataLoader(range(1000), batch_sampler=s)
|
|
_set_sampler_epoch(loader, 5)
|
|
self.assertEqual(s.epoch, 5)
|
|
self.assertEqual([b.tolist() for b in loader], s.batches(5))
|
|
|
|
def _resumed(self, state, epoch: int, **kw):
|
|
"""``(batches, what was printed)`` of a new run resuming from *state* in *epoch*."""
|
|
sampler = self._sampler(**kw)
|
|
with contextlib.redirect_stdout(io.StringIO()) as out:
|
|
sampler.load_state_dict(state)
|
|
sampler.set_epoch(epoch)
|
|
return list(sampler), out.getvalue()
|
|
|
|
def test_an_interrupted_epoch_resumes_with_the_batches_it_had_not_trained(self) -> None:
|
|
s = self._sampler()
|
|
self.assertEqual(s.state_dict(0), {}) # no epoch started
|
|
s.set_epoch(3)
|
|
first = list(s) # the run's first epoch: longest batch first
|
|
state = s.state_dict(40)
|
|
self.assertEqual((state["epoch"], state["batches_done"]), (3, 40))
|
|
resumed = self._sampler()
|
|
with contextlib.redirect_stdout(io.StringIO()) as out:
|
|
resumed.load_state_dict(json.loads(json.dumps(state))) # (as saved and loaded)
|
|
rest = list(resumed) # as Lightning does: before it sets the epoch
|
|
resumed.set_epoch(3)
|
|
self.assertEqual(out.getvalue(), "")
|
|
self.assertEqual(list(resumed), rest) # a DataLoader with workers iterates twice
|
|
self.assertCountEqual(first[:40] + rest, s.batches(3)) # each batch once
|
|
self.assertEqual(max(self.lengths[i] for i in rest[0]), # longest first again
|
|
max(self.lengths[i] for b in first[40:] for i in b))
|
|
resumed.set_epoch(4)
|
|
self.assertEqual(list(resumed), s.batches(4)) # later epochs are whole
|
|
later = self._sampler(longest_first=False)
|
|
later.set_epoch(4)
|
|
list(later)
|
|
rest, _ = self._resumed(later.state_dict(10), 4, longest_first=False)
|
|
self.assertEqual(rest, later.batches(4)[10:]) # in the order it would have had
|
|
|
|
def test_the_loader_saves_and_restores_the_position_in_its_epoch(self) -> None:
|
|
from torch.utils.data import DataLoader
|
|
trained = {"batches": 0}
|
|
loader = self.pc.ResumableDataLoader(range(1000), batch_sampler=self._sampler(),
|
|
batches_done=lambda: trained["batches"])
|
|
self.assertIsInstance(loader, DataLoader)
|
|
self.assertEqual(loader.state_dict(), {}) # (before its first epoch)
|
|
first = [b.tolist() for b in loader]
|
|
trained["batches"] = 100
|
|
state = loader.state_dict()
|
|
self.assertEqual((state["epoch"], state["batches_done"]), (0, 100))
|
|
resumed = self.pc.ResumableDataLoader(range(1000), batch_sampler=self._sampler())
|
|
resumed.load_state_dict(state)
|
|
once, twice = [b.tolist() for b in resumed], [b.tolist() for b in resumed]
|
|
self.assertEqual(once, twice)
|
|
self.assertCountEqual(once, first[100:])
|
|
|
|
def test_saved_positions_that_cannot_resume_start_the_epoch_afresh(self) -> None:
|
|
s = self._sampler(longest_first=False)
|
|
s.set_epoch(2)
|
|
list(s)
|
|
cases = [
|
|
("piper's random batches", {}, 2),
|
|
("nothing", None, 2),
|
|
("damaged", {**s.state_dict(5), "order": "x"}, 2),
|
|
("start of the epoch", s.state_dict(0), 2),
|
|
("end of the epoch", s.state_dict(len(s)), 2), # validated and saved: the next one
|
|
("another epoch", s.state_dict(5), 3),
|
|
]
|
|
for name, state, epoch in cases:
|
|
with self.subTest(name):
|
|
self.assertEqual(self._resumed(state, epoch, longest_first=False),
|
|
(s.batches(epoch), ""))
|
|
|
|
def test_changed_training_data_starts_the_interrupted_epoch_afresh(self) -> None:
|
|
s = self._sampler(longest_first=False)
|
|
list(s)
|
|
state = s.state_dict(5)
|
|
for change in ("a clip more", "a clip longer"):
|
|
with self.subTest(change):
|
|
self.lengths = list(s.lengths)
|
|
if change == "a clip more":
|
|
self.lengths.append(500)
|
|
else: # as many batches, but different ones
|
|
self.lengths[self.lengths.index(min(self.lengths))] = 5000
|
|
fresh = self._sampler(longest_first=False).batches(0)
|
|
self.assertNotEqual(fresh, s.batches(0))
|
|
batches, out = self._resumed(state, 0, longest_first=False)
|
|
self.assertIn("training data changed", out)
|
|
self.assertEqual(batches, fresh)
|
|
|
|
def test_train_dataloader_buckets_piper_utterances(self) -> None:
|
|
import torch
|
|
from lightning.fabric.utilities.data import _set_sampler_epoch
|
|
from piper.train.vits.dataset import CachedUtterance, VitsDataset
|
|
from torch.utils.data import Subset
|
|
|
|
samples = [1000, 20000, 1100, 19000, 5000, 5200, 300]
|
|
with tempfile.TemporaryDirectory() as d:
|
|
utts = []
|
|
for i, n in enumerate(samples):
|
|
stem = Path(d) / f"u{i}"
|
|
torch.save(torch.tensor([1, 2, 3]), f"{stem}.phonemes.pt")
|
|
torch.save(torch.zeros(n), f"{stem}.audio.pt")
|
|
torch.save(torch.zeros(4, n // 256 + 1), f"{stem}.spec.pt")
|
|
utts.append(CachedUtterance(Path(f"{stem}.phonemes.pt"), Path(f"{stem}.audio.pt"),
|
|
Path(f"{stem}.spec.pt"), text="x"))
|
|
# random_split() hands piper a Subset (here nested); clip 6 is held out.
|
|
train = Subset(Subset(VitsDataset(utts), [6, 5, 4, 3, 2, 1, 0]), [1, 2, 3, 4, 5, 6])
|
|
dm = SimpleNamespace(train_dataset=train, batch_size=2, num_workers=0,
|
|
num_speakers=1, segment_size=256)
|
|
with contextlib.redirect_stdout(io.StringIO()) as out:
|
|
loader = self.pc._bucketed_train_dataloader(dm)
|
|
self.assertIn("Length-bucketed batches: 3 per epoch", out.getvalue())
|
|
self.assertEqual(len(loader), 3)
|
|
_set_sampler_epoch(loader, 2)
|
|
loaded = list(loader)
|
|
batches = [sorted(b.audio_lengths.tolist()) for b in loaded]
|
|
self.assertEqual(batches[0], [19000, 20000]) # longest first
|
|
self.assertCountEqual(batches, [[1000, 1100], [5000, 5200], [19000, 20000]])
|
|
self.assertEqual({b.spectrograms.shape[-1] % 32 for b in loaded}, {0}) # rounded up
|
|
self.assertEqual({b.phoneme_ids.shape[-1] for b in loaded}, {16})
|
|
|
|
def test_collate_rounds_padded_lengths_up(self) -> None:
|
|
import pickle
|
|
|
|
import torch
|
|
from piper.train.vits.dataset import UtteranceCollate, UtteranceTensors
|
|
utts = [UtteranceTensors(phoneme_ids=torch.arange(1, p + 1), spectrogram=torch.rand(4, t),
|
|
audio_norm=torch.rand(t * 256))
|
|
for p, t in ((5, 40), (17, 33))]
|
|
plain = UtteranceCollate(is_multispeaker=False, segment_size=256)
|
|
# DataLoader workers get the collate function pickled.
|
|
rounded = pickle.loads(pickle.dumps(self.pc.RoundedCollate(plain)))
|
|
expected, batch = plain(utts), rounded(utts)
|
|
self.assertEqual(tuple(batch.spectrograms.shape), (2, 4, 64))
|
|
self.assertEqual(tuple(batch.phoneme_ids.shape), (2, 32))
|
|
torch.testing.assert_close(batch.spectrograms[..., :40], expected.spectrograms, rtol=0, atol=0)
|
|
torch.testing.assert_close(batch.phoneme_ids[:, :17], expected.phoneme_ids, rtol=0, atol=0)
|
|
self.assertFalse(batch.spectrograms[..., 40:].any())
|
|
self.assertFalse(batch.phoneme_ids[:, 17:].any())
|
|
for name in ("phoneme_lengths", "spectrogram_lengths", "audios", "audio_lengths"):
|
|
torch.testing.assert_close(getattr(batch, name), getattr(expected, name), rtol=0, atol=0)
|
|
exact = [UtteranceTensors(phoneme_ids=torch.arange(1, 17), spectrogram=torch.rand(4, 64),
|
|
audio_norm=torch.rand(64 * 256))]
|
|
self.assertEqual(rounded(exact).spectrograms.shape[-1], 64)
|
|
|
|
def test_extra_padding_leaves_real_frames_alone(self) -> None:
|
|
import torch
|
|
from piper.train.vits.models import SynthesizerTrn
|
|
torch.manual_seed(0)
|
|
model = SynthesizerTrn(
|
|
n_vocab=12, spec_channels=17, segment_size=8, inter_channels=8, hidden_channels=8,
|
|
filter_channels=16, n_heads=2, n_layers=1, kernel_size=3, p_dropout=0.0,
|
|
resblock="2", resblock_kernel_sizes=(3,), resblock_dilation_sizes=((1, 3),),
|
|
upsample_rates=(2, 2), upsample_initial_channel=8, upsample_kernel_sizes=(4, 4),
|
|
use_sdp=False).eval()
|
|
x_lengths, y_lengths = torch.tensor([9, 5]), torch.tensor([40, 23])
|
|
x = torch.randint(1, 12, (2, 9)) * (torch.arange(9) < x_lengths[:, None])
|
|
y = torch.randn(2, 17, 40) * (torch.arange(40) < y_lengths[:, None, None])
|
|
|
|
def encode(x, y): # every part of piper's VITS that sees whole utterances
|
|
with torch.no_grad():
|
|
h, m_p, logs_p, x_mask = model.enc_p(x, x_lengths)
|
|
_, m_q, logs_q, y_mask = model.enc_q(y, y_lengths)
|
|
return {"h": h, "m_p": m_p, "logs_p": logs_p, "logw": model.dp(h, x_mask),
|
|
"m_q": m_q, "logs_q": logs_q, "z_p": model.flow(m_q, y_mask)}
|
|
|
|
padded = encode(self.pc._pad_to_multiple(x, 16), self.pc._pad_to_multiple(y, 32))
|
|
for name, value in encode(x, y).items():
|
|
with self.subTest(name):
|
|
self.assertGreater(padded[name].shape[-1], value.shape[-1])
|
|
torch.testing.assert_close(padded[name][..., :value.shape[-1]], value)
|
|
|
|
def test_training_aids_hook_into_piper(self) -> None:
|
|
from lightning.fabric.utilities.types import _Stateful
|
|
from lightning.pytorch.utilities.combined_loader import CombinedLoader
|
|
from piper.train.vits.dataset import VitsDataModule
|
|
from piper.train.vits.lightning import VitsModel
|
|
from torch.utils.data import DataLoader
|
|
tc = self.pc.training_control
|
|
console = SimpleNamespace(closing=threading.Event())
|
|
for tuning, stop_file in (("1", None), ("0", "stop")):
|
|
with self.subTest(JAVIS_PIPER_CUDNN_TUNING=tuning, stop_file=stop_file), \
|
|
patch.object(VitsModel, "configure_callbacks", VitsModel.configure_callbacks), \
|
|
patch.object(VitsDataModule, "train_dataloader",
|
|
lambda dm: DataLoader(range(3))), \
|
|
patch.object(self.pc, "_avoid_gpu_syncs") as avoid_gpu_syncs, \
|
|
patch.object(self.pc, "_safe_checkpoint_files") as safe_checkpoint_files, \
|
|
patch.object(self.pc, "_checkpoint_every_epoch") as checkpoint_every_epoch, \
|
|
patch.object(self.pc, "_robust_dataset") as robust_dataset, \
|
|
patch.object(self.pc, "_on_cache_write", None), \
|
|
patch.dict(os.environ, {"JAVIS_PIPER_BUCKETING": "1",
|
|
"JAVIS_PIPER_CUDNN_TUNING": tuning,
|
|
"JAVIS_PIPER_PROGRESS_MINUTES": "4"}):
|
|
os.environ.pop(tc.LAUNCHER_ENV, None)
|
|
if stop_file:
|
|
os.environ[tc.STOP_FILE_ENV] = stop_file
|
|
else:
|
|
os.environ.pop(tc.STOP_FILE_ENV, None)
|
|
stopper = self.pc._install_training_aids(console)
|
|
for patched in (avoid_gpu_syncs, safe_checkpoint_files, checkpoint_every_epoch):
|
|
patched.assert_called_once_with()
|
|
robust_dataset.assert_called_once_with(VitsDataModule)
|
|
callbacks = VitsModel.configure_callbacks(SimpleNamespace())
|
|
self.assertTrue(any(isinstance(c, self.pc.NonFiniteGuard) for c in callbacks))
|
|
self.assertEqual(any(isinstance(c, self.pc.CudnnFirstStepTuning) for c in callbacks),
|
|
tuning == "1")
|
|
savers = [c for c in callbacks if isinstance(c, self.pc.ProgressSaver)]
|
|
self.assertEqual([(s.stopper, s.console, s.interval) for s in savers],
|
|
[(stopper, console, 240.0)])
|
|
if stop_file: # started by the launcher: stops on request, even while preprocessing
|
|
self.assertIsInstance(stopper, self.pc.StopOnRequest)
|
|
self.assertIn(stopper, callbacks)
|
|
self.assertEqual(self.pc._on_cache_write, stopper.raise_if_requested)
|
|
else:
|
|
self.assertIsNone(stopper)
|
|
self.assertIsNone(self.pc._on_cache_write)
|
|
with contextlib.redirect_stderr(io.StringIO()) as err: # unusable for bucketing
|
|
loader = VitsDataModule.train_dataloader(SimpleNamespace())
|
|
self.assertIn("length bucketing unavailable", err.getvalue())
|
|
# piper's batches still load the checkpoints of bucketed ones
|
|
self.assertIsInstance(loader, _Stateful)
|
|
self.assertEqual(loader.state_dict(), {})
|
|
CombinedLoader(loader)._load_state_dicts([{"epoch": 1, "batches_done": 5}])
|
|
self.assertEqual(list(loader), [0, 1, 2])
|
|
|
|
|
|
class PrecisionTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
|
|
def _resolve(self, argv, **kw):
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
self.pc._resolve_precision(argv, **kw)
|
|
return argv
|
|
|
|
def test_bf16_falls_back_to_32_without_native_support(self) -> None:
|
|
argv = ["fit", "--trainer.precision", "bf16-mixed", "--trainer.devices", "1"]
|
|
self.assertEqual(self._resolve(list(argv), bf16_supported=True), argv)
|
|
self.assertEqual(self._resolve(list(argv), bf16_supported=False)[2], "32")
|
|
self.assertEqual(self._resolve(["fit", "--trainer.precision=bf16-mixed"],
|
|
bf16_supported=False)[1], "--trainer.precision=32")
|
|
|
|
def test_other_precisions_are_left_alone(self) -> None:
|
|
for value in ("32", "16-mixed"):
|
|
argv = ["fit", "--trainer.precision", value]
|
|
self.assertEqual(self._resolve(list(argv), bf16_supported=False), argv)
|
|
|
|
def test_cpu_training_uses_32(self) -> None:
|
|
with patch.object(self.pc, "_cuda_bf16_supported", return_value=True):
|
|
argv = self._resolve(["fit", "--trainer.accelerator", "cpu",
|
|
"--trainer.precision", "bf16-mixed"])
|
|
self.assertEqual(argv[-1], "32")
|
|
argv = self._resolve(["fit", "--trainer.accelerator", "gpu",
|
|
"--trainer.precision", "bf16-mixed"])
|
|
self.assertEqual(argv[-1], "bf16-mixed")
|
|
|
|
def test_config_defaults_to_bf16(self) -> None:
|
|
from javis.core.config import Config, validate
|
|
cfg = Config()
|
|
self.assertEqual(cfg.training.precision, "bf16-mixed")
|
|
cfg.training.precision = "fp8"
|
|
with self.assertLogs("javis.config", level="WARNING"):
|
|
validate(cfg)
|
|
self.assertEqual(cfg.training.precision, "bf16-mixed")
|
|
|
|
|
|
class NonFiniteGuardTests(unittest.TestCase):
|
|
"""The guard inside a real Lightning loop with piper's GAN-style manual optimization."""
|
|
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
|
|
def _fit(self, nan_batches, guard, optimizer=None):
|
|
import lightning as L
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
|
|
make_optimizer = optimizer or torch.optim.AdamW
|
|
|
|
class GanLike(L.LightningModule):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.automatic_optimization = False
|
|
self.g = torch.nn.Linear(2, 2)
|
|
self.d = torch.nn.Linear(2, 1)
|
|
self.batches = []
|
|
|
|
def training_step(self, batch, batch_idx):
|
|
self.batches.append(batch_idx)
|
|
opt_g, opt_d = self.optimizers()
|
|
loss_g = self.g(batch).sum()
|
|
if batch_idx in nan_batches:
|
|
loss_g = loss_g * float("inf")
|
|
opt_g.zero_grad()
|
|
self.manual_backward(loss_g)
|
|
opt_g.step()
|
|
opt_d.zero_grad()
|
|
self.manual_backward(self.d(batch).sum())
|
|
opt_d.step()
|
|
|
|
def configure_optimizers(self):
|
|
self.opts = [make_optimizer(self.g.parameters(), lr=0.1),
|
|
make_optimizer(self.d.parameters(), lr=0.1)]
|
|
return self.opts
|
|
|
|
def train_dataloader(self):
|
|
return DataLoader(torch.ones(6, 2), batch_size=1)
|
|
|
|
model = self.model = GanLike()
|
|
trainer = L.Trainer(accelerator="cpu", max_epochs=1, callbacks=[guard], logger=False,
|
|
enable_checkpointing=False, enable_progress_bar=False,
|
|
enable_model_summary=False)
|
|
with warnings.catch_warnings(), contextlib.redirect_stderr(io.StringIO()):
|
|
warnings.simplefilter("ignore")
|
|
trainer.fit(model)
|
|
return model
|
|
|
|
def test_bad_steps_are_skipped_and_counted(self) -> None:
|
|
import torch
|
|
guard = self.pc.NonFiniteGuard(max_consecutive=3)
|
|
model = self._fit({1, 2, 4}, guard)
|
|
self.assertEqual(guard.skipped, {0: 3}) # generator only; streak was reset
|
|
self.assertTrue(all(bool(torch.isfinite(p).all()) for p in model.parameters()))
|
|
# piper's AdamW skips on the device: fused, its state untouched by bad steps
|
|
for opt, good_steps in zip(model.opts, (3, 6)):
|
|
self.assertTrue(all(group["fused"] for group in opt.param_groups))
|
|
self.assertEqual({float(s["step"]) for s in opt.state.values()}, {good_steps})
|
|
|
|
def test_other_optimizers_are_checked_on_the_cpu(self) -> None:
|
|
import torch
|
|
guard = self.pc.NonFiniteGuard(max_consecutive=3)
|
|
model = self._fit({1, 2, 4}, guard, optimizer=torch.optim.SGD)
|
|
self.assertEqual(guard.skipped, {0: 3})
|
|
self.assertTrue(all(bool(torch.isfinite(p).all()) for p in model.parameters()))
|
|
|
|
def test_a_diverged_run_stops(self) -> None:
|
|
import torch
|
|
for optimizer in (torch.optim.AdamW, torch.optim.SGD):
|
|
with self.subTest(optimizer=optimizer.__name__):
|
|
guard = self.pc.NonFiniteGuard(max_consecutive=4, check_every=2)
|
|
with self.assertRaises(self.pc.TrainingDivergedError):
|
|
self._fit({1, 2, 3, 4, 5}, guard, optimizer)
|
|
self.assertEqual(self.model.batches, [0, 1, 2, 3, 4]) # at the 4th in a row
|
|
|
|
def test_restored_adam_state_moves_to_the_fused_kernel(self) -> None:
|
|
import copy
|
|
|
|
import torch
|
|
torch.manual_seed(0)
|
|
layer = torch.nn.Linear(3, 3)
|
|
twin = copy.deepcopy(layer)
|
|
|
|
def train(model, optimizer, steps):
|
|
for _ in range(steps):
|
|
optimizer.zero_grad()
|
|
model(torch.ones(2, 3)).square().sum().backward()
|
|
optimizer.step()
|
|
|
|
optimizer = torch.optim.AdamW(layer.parameters(), lr=0.01, foreach=True)
|
|
reference = torch.optim.AdamW(twin.parameters(), lr=0.01, foreach=True)
|
|
train(layer, optimizer, 2)
|
|
train(twin, reference, 2)
|
|
restored = torch.optim.AdamW(layer.parameters(), lr=0.01) # as Lightning resumes it
|
|
restored.load_state_dict(copy.deepcopy(optimizer.state_dict()))
|
|
self.assertTrue(self.pc._use_fused_adam(restored))
|
|
self.assertTrue(all(g["fused"] and not g["foreach"] for g in restored.param_groups))
|
|
train(layer, restored, 2)
|
|
train(twin, reference, 2)
|
|
for p, q in zip(layer.parameters(), twin.parameters()):
|
|
torch.testing.assert_close(p, q)
|
|
self.assertEqual({float(s["step"]) for s in restored.state.values()}, {4.0})
|
|
resumed = torch.optim.AdamW(layer.parameters(), lr=0.01) # its checkpoint, resumed
|
|
resumed.load_state_dict(restored.state_dict())
|
|
self.assertTrue(all(g["fused"] for g in resumed.param_groups))
|
|
self.assertTrue(self.pc._use_fused_adam(resumed))
|
|
self.assertFalse(self.pc._use_fused_adam(torch.optim.SGD(layer.parameters(), lr=0.1)))
|
|
|
|
def test_skipping_on_the_gpu_never_waits_for_it(self) -> None:
|
|
import torch
|
|
if not torch.cuda.is_available():
|
|
self.skipTest("needs CUDA")
|
|
layer = torch.nn.Linear(4, 4, device="cuda")
|
|
optimizer = torch.optim.AdamW(layer.parameters(), lr=0.1)
|
|
trainer = SimpleNamespace(optimizers=[optimizer], global_step=0)
|
|
guard = self.pc.NonFiniteGuard()
|
|
guard.on_train_start(trainer, None)
|
|
start = layer.weight.detach().clone()
|
|
x = torch.ones(2, 4, device="cuda")
|
|
weights = []
|
|
for scale in (float("inf"), 1.0):
|
|
optimizer.zero_grad()
|
|
(layer(x).sum() * scale).backward()
|
|
torch.cuda.synchronize()
|
|
torch.cuda.set_sync_debug_mode("error")
|
|
try:
|
|
guard.on_before_optimizer_step(trainer, None, optimizer)
|
|
optimizer.step()
|
|
finally:
|
|
torch.cuda.set_sync_debug_mode("default")
|
|
weights.append(layer.weight.detach().clone())
|
|
torch.testing.assert_close(weights[0], start, rtol=0, atol=0) # skipped
|
|
self.assertFalse(torch.equal(weights[1], start))
|
|
self.assertEqual({float(s["step"]) for s in optimizer.state.values()}, {1.0})
|
|
with contextlib.redirect_stderr(io.StringIO()):
|
|
guard.on_train_end(trainer, None)
|
|
self.assertEqual(guard.skipped, {0: 1})
|
|
|
|
def test_collapsed_fp16_scale_stops_the_run(self) -> None:
|
|
guard = self.pc.NonFiniteGuard()
|
|
module = SimpleNamespace(log_dict=lambda *a, **k: None)
|
|
for scale, fails in ((65536.0, False), (0.0, True), (float("inf"), True)):
|
|
scaler = SimpleNamespace(is_enabled=lambda: True, get_scale=lambda s=scale: s)
|
|
trainer = SimpleNamespace(precision_plugin=SimpleNamespace(scaler=scaler))
|
|
if fails:
|
|
with self.assertRaises(self.pc.TrainingDivergedError):
|
|
guard.on_train_batch_end(trainer, module, None, None, 0)
|
|
else:
|
|
guard.on_train_batch_end(trainer, module, None, None, 0)
|
|
|
|
def test_nan_weights_are_never_saved(self) -> None:
|
|
import torch
|
|
guard = self.pc.NonFiniteGuard()
|
|
layer = torch.nn.Linear(2, 2)
|
|
guard.on_save_checkpoint(None, layer, {})
|
|
with torch.no_grad():
|
|
layer.weight[0, 0] = float("nan")
|
|
with self.assertRaises(self.pc.TrainingDivergedError):
|
|
guard.on_save_checkpoint(None, layer, {})
|
|
|
|
|
|
class GpuSyncTests(unittest.TestCase):
|
|
"""piper_compat's sync-free replacements compute exactly what piper's code computes."""
|
|
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
import torch
|
|
self.torch = torch
|
|
torch.manual_seed(0)
|
|
|
|
def _run(self, fn, *inputs):
|
|
"""``fn``'s outputs and the gradients of a fixed random projection of them."""
|
|
torch = self.torch
|
|
leaves = [t.detach().clone().requires_grad_(t.is_floating_point()) for t in inputs]
|
|
outs = fn(*leaves)
|
|
outs = outs if isinstance(outs, tuple) else (outs,)
|
|
gen = torch.Generator().manual_seed(1)
|
|
sum((o * torch.randn(o.shape, generator=gen, dtype=o.dtype)).sum() for o in outs).backward()
|
|
return [o.detach() for o in outs] + [t.grad for t in leaves if t.requires_grad]
|
|
|
|
def _assert_same(self, fn, reference, *inputs) -> None:
|
|
for mine, theirs in zip(self._run(fn, *inputs), self._run(reference, *inputs), strict=True):
|
|
self.torch.testing.assert_close(mine, theirs, rtol=0, atol=0)
|
|
|
|
def test_slice_segments_matches_piper(self) -> None:
|
|
torch = self.torch
|
|
piper = _pristine("commons").slice_segments
|
|
cases = [(torch.randn(3, 4, 50, dtype=torch.float64), torch.tensor([0, 7, -2]), 16),
|
|
(torch.randn(2, 4, 10, dtype=torch.float64), torch.tensor([0, -1]), 16),
|
|
(torch.randn(2, 1, 300, dtype=torch.float64), torch.tensor([3, 5]) * 16, 64)]
|
|
for x, ids, size in cases:
|
|
self._assert_same(lambda t, ids=ids, size=size: self.pc.slice_segments(t, ids, size),
|
|
lambda t, ids=ids, size=size: piper(t, ids, size), x)
|
|
|
|
def test_rand_slice_segments_matches_piper(self) -> None:
|
|
torch = self.torch
|
|
piper = _pristine("commons").rand_slice_segments
|
|
x, lengths = torch.randn(3, 4, 50), torch.tensor([50, 40, 10]) # 10: shorter than a window
|
|
starts = set()
|
|
for seed in range(10):
|
|
torch.manual_seed(seed) # on the CPU both draw the same starts
|
|
out, ids = self.pc.rand_slice_segments(x, lengths, 16)
|
|
torch.manual_seed(seed)
|
|
ref_out, ref_ids = piper(x, lengths, 16)
|
|
torch.testing.assert_close(ids, ref_ids, rtol=0, atol=0)
|
|
torch.testing.assert_close(out, ref_out, rtol=0, atol=0)
|
|
self.assertTrue(bool((ids.clamp_min(0) + 16 <= lengths.clamp_min(16)).all()))
|
|
starts.add(int(ids[0]))
|
|
self.assertGreater(len(starts), 1)
|
|
|
|
def test_discriminator_loss_matches_piper(self) -> None:
|
|
torch = self.torch
|
|
piper = _pristine("losses").discriminator_loss
|
|
outputs = [torch.randn(2, 1, 30, dtype=torch.bfloat16), torch.randn(2, 5),
|
|
torch.randn(2, 1, 30, dtype=torch.bfloat16), torch.randn(2, 5)]
|
|
|
|
def run(fn):
|
|
leaves = [t.clone().requires_grad_(True) for t in outputs]
|
|
loss, r_losses, g_losses = fn(leaves[:2], leaves[2:])
|
|
loss.backward()
|
|
return loss.detach(), [float(v) for v in r_losses + g_losses], [t.grad for t in leaves]
|
|
|
|
(loss, terms, grads), (ref_loss, ref_terms, ref_grads) = run(self.pc.discriminator_loss), run(piper)
|
|
torch.testing.assert_close(loss, ref_loss, rtol=0, atol=0)
|
|
self.assertEqual(terms, ref_terms)
|
|
for grad, ref_grad in zip(grads, ref_grads, strict=True):
|
|
torch.testing.assert_close(grad, ref_grad, rtol=0, atol=0)
|
|
|
|
def test_spline_matches_piper(self) -> None:
|
|
torch = self.torch
|
|
reference, patched = _pristine("transforms"), _pristine("transforms")
|
|
patched.unconstrained_rational_quadratic_spline = self.pc._mask_free_spline(
|
|
patched.unconstrained_rational_quadratic_spline)
|
|
inputs = torch.linspace(-7, 7, 60, dtype=torch.float64).view(2, 1, 30) # partly outside ±5
|
|
params = [torch.randn(2, 1, 30, n, dtype=torch.float64) for n in (10, 10, 9)]
|
|
for inverse, x in ((False, inputs), (True, inputs * 0.6)): # inverse: inference, piper's code
|
|
self._assert_same(
|
|
lambda *a, inverse=inverse: patched.piecewise_rational_quadratic_transform(
|
|
*a, inverse=inverse, tails="linear", tail_bound=5.0),
|
|
lambda *a, inverse=inverse: reference.piecewise_rational_quadratic_transform(
|
|
*a, inverse=inverse, tails="linear", tail_bound=5.0),
|
|
x, *params)
|
|
|
|
def test_mel_spectrogram_matches_piper(self) -> None:
|
|
torch = self.torch
|
|
piper = _pristine("mel_processing").mel_spectrogram_torch
|
|
args = (1024, 80, 22050, 256, 1024, 0.0, None)
|
|
y = torch.rand(2, 8192) * 2 - 1
|
|
self._assert_same(lambda t: self.pc.mel_spectrogram_torch(t, *args),
|
|
lambda t: piper(t, *args), y)
|
|
|
|
def test_flip_matches_piper(self) -> None:
|
|
torch = self.torch
|
|
from piper.train.vits import modules
|
|
flip, x = modules.Flip(), torch.randn(2, 4, 7, dtype=torch.float64)
|
|
with patch.object(modules.Flip, "forward", self.pc._flip_forward):
|
|
flipped, logdet = flip(x)
|
|
restored = flip(flipped, reverse=True)
|
|
torch.testing.assert_close(flipped, torch.flip(x, [1]), rtol=0, atol=0)
|
|
torch.testing.assert_close(restored, x, rtol=0, atol=0)
|
|
torch.testing.assert_close(logdet, torch.zeros(2, dtype=torch.float64), rtol=0, atol=0)
|
|
|
|
def test_batches_are_pinned(self) -> None:
|
|
import dataclasses
|
|
|
|
from piper.train.vits.dataset import Batch
|
|
from torch.utils.data._utils.pin_memory import pin_memory
|
|
torch = self.torch
|
|
if not torch.cuda.is_available():
|
|
self.skipTest("pinned memory needs CUDA")
|
|
batch = Batch(phoneme_ids=torch.ones(2, 3, dtype=torch.long),
|
|
phoneme_lengths=torch.tensor([3, 2]), spectrograms=torch.randn(2, 5, 4),
|
|
spectrogram_lengths=torch.tensor([4, 3]), audios=torch.randn(2, 1, 1024),
|
|
audio_lengths=torch.tensor([1024, 768]), speaker_ids=None)
|
|
with patch.object(Batch, "pin_memory", self.pc._pin_batch, create=True):
|
|
pinned = pin_memory(batch) # what the DataLoader's pinning thread calls
|
|
self.assertIsInstance(pinned, Batch)
|
|
self.assertIsNone(pinned.speaker_ids)
|
|
for field in dataclasses.fields(Batch):
|
|
value = getattr(batch, field.name)
|
|
if value is not None:
|
|
self.assertTrue(getattr(pinned, field.name).is_pinned(), field.name)
|
|
torch.testing.assert_close(getattr(pinned, field.name), value, rtol=0, atol=0)
|
|
|
|
def test_training_math_never_waits_for_the_gpu(self) -> None:
|
|
torch = self.torch
|
|
if not torch.cuda.is_available():
|
|
self.skipTest("needs CUDA")
|
|
cuda = torch.device("cuda")
|
|
spline = self.pc._mask_free_spline(
|
|
_pristine("transforms").unconstrained_rational_quadratic_spline)
|
|
z = torch.randn(4, 8, 100, device=cuda, requires_grad=True)
|
|
lengths = torch.tensor([100, 80, 60, 20], device=cuda)
|
|
audio = (torch.rand(4, 1, 100 * 256, device=cuda) * 2 - 1).requires_grad_()
|
|
flow_in = (torch.randn(4, 1, 50, device=cuda) * 4).requires_grad_()
|
|
flow_params = [torch.randn(4, 1, 50, n, device=cuda, requires_grad=True) for n in (10, 10, 9)]
|
|
mel_args = (1024, 80, 22050, 256, 1024, 0.0, None)
|
|
self.pc.mel_spectrogram_torch(audio[:, 0], *mel_args) # caches window and mel basis
|
|
torch.cuda.synchronize()
|
|
torch.cuda.set_sync_debug_mode("error")
|
|
try:
|
|
z_slice, ids = self.pc.rand_slice_segments(z, lengths, 32)
|
|
mel = self.pc.mel_spectrogram_torch(
|
|
self.pc.slice_segments(audio, ids * 256, 32 * 256).squeeze(1), *mel_args)
|
|
flow_out, flow_logdet = spline(flow_in, *flow_params, tails="linear", tail_bound=5.0)
|
|
flipped, flip_logdet = self.pc._flip_forward(None, z_slice)
|
|
d_loss, _, _ = self.pc.discriminator_loss([z_slice], [flipped])
|
|
(mel.mean() + flow_out.sum() + flow_logdet.sum() + flip_logdet.sum() + d_loss).backward()
|
|
finally:
|
|
torch.cuda.set_sync_debug_mode("default")
|
|
torch.cuda.synchronize()
|
|
self.assertTrue(all(t.grad is not None for t in (z, audio, flow_in, *flow_params)))
|
|
|
|
def test_installed_into_piper(self) -> None:
|
|
from piper.train.vits import commons, losses, modules, transforms
|
|
from piper.train.vits import lightning as vits_lightning
|
|
from piper.train.vits.dataset import Batch
|
|
targets = [(commons, "slice_segments"), (commons, "rand_slice_segments"),
|
|
(losses, "discriminator_loss"), (transforms, "unconstrained_rational_quadratic_spline"),
|
|
(vits_lightning, "slice_segments"), (vits_lightning, "discriminator_loss"),
|
|
(vits_lightning, "mel_spectrogram_torch"), (modules.Flip, "forward")]
|
|
had_pin_memory = "pin_memory" in vars(Batch)
|
|
|
|
def drop_pin_memory():
|
|
if not had_pin_memory and "pin_memory" in vars(Batch):
|
|
del Batch.pin_memory
|
|
|
|
with contextlib.ExitStack() as restore:
|
|
for owner, name in targets:
|
|
restore.enter_context(patch.object(owner, name, getattr(owner, name)))
|
|
restore.callback(drop_pin_memory)
|
|
self.pc._avoid_gpu_syncs()
|
|
for module in (commons, vits_lightning):
|
|
self.assertIs(module.slice_segments, self.pc.slice_segments)
|
|
for module in (losses, vits_lightning):
|
|
self.assertIs(module.discriminator_loss, self.pc.discriminator_loss)
|
|
self.assertIs(commons.rand_slice_segments, self.pc.rand_slice_segments)
|
|
self.assertIs(vits_lightning.mel_spectrogram_torch, self.pc.mel_spectrogram_torch)
|
|
self.assertIs(modules.Flip.forward, self.pc._flip_forward)
|
|
self.assertTrue(hasattr(Batch, "pin_memory"))
|
|
spline = transforms.unconstrained_rational_quadratic_spline
|
|
self.assertTrue(spline._javis_mask_free)
|
|
self.pc._avoid_gpu_syncs()
|
|
self.assertIs(transforms.unconstrained_rational_quadratic_spline, spline) # not wrapped twice
|
|
|
|
|
|
class CudnnTuningTests(unittest.TestCase):
|
|
def test_cudnn_times_algorithms_during_the_first_step_only(self) -> None:
|
|
import lightning as L
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
pc = _piper_compat()
|
|
self.addCleanup(setattr, torch.backends.cudnn, "benchmark", torch.backends.cudnn.benchmark)
|
|
seen = []
|
|
|
|
class Probe(L.LightningModule):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.w = torch.nn.Linear(1, 1)
|
|
|
|
def training_step(self, batch, batch_idx):
|
|
seen.append(torch.backends.cudnn.benchmark)
|
|
return self.w(batch).sum()
|
|
|
|
def configure_optimizers(self):
|
|
return torch.optim.SGD(self.parameters(), lr=0.1)
|
|
|
|
def train_dataloader(self):
|
|
return DataLoader(torch.ones(3, 1), batch_size=1)
|
|
|
|
for asked in (False, True): # --trainer.benchmark
|
|
seen.clear()
|
|
trainer = L.Trainer(accelerator="cpu", max_epochs=1, benchmark=asked, logger=False,
|
|
callbacks=[pc.CudnnFirstStepTuning()], enable_checkpointing=False,
|
|
enable_progress_bar=False, enable_model_summary=False)
|
|
with warnings.catch_warnings(), contextlib.redirect_stderr(io.StringIO()):
|
|
warnings.simplefilter("ignore")
|
|
trainer.fit(Probe())
|
|
with self.subTest(asked=asked):
|
|
self.assertEqual(seen, [True, asked, asked])
|
|
self.assertEqual(torch.backends.cudnn.benchmark, asked)
|
|
|
|
|
|
class CheckpointSavingTests(unittest.TestCase):
|
|
"""Checkpoints are written whole, last.ckpt every epoch; a missing metric isn't fatal."""
|
|
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.root = Path(d.name)
|
|
|
|
def _files(self) -> list[str]:
|
|
return sorted(p.name for p in self.root.iterdir())
|
|
|
|
def test_last_checkpoint_is_saved_every_epoch(self) -> None:
|
|
import torch
|
|
with _checkpoint_patches(self.pc):
|
|
model, trainer = _toy_training(self.root, epochs=3)
|
|
_, _, err = _captured(lambda: trainer.fit(model))
|
|
last = torch.load(self.root / "last.ckpt", map_location="cpu", weights_only=False)
|
|
self.assertEqual((last["epoch"], last["global_step"]), (2, 12))
|
|
self.assertEqual(self._files(), ["epoch=0-val_mel=1.0000.ckpt", "last.ckpt"])
|
|
self.assertEqual(err.count("val_mos was not logged"), 1)
|
|
|
|
def test_a_failed_save_keeps_the_previous_checkpoint(self) -> None:
|
|
import torch
|
|
from lightning.fabric.plugins.io import torch_io
|
|
|
|
class Unsaveable:
|
|
def __reduce__(self):
|
|
raise RuntimeError("disk full")
|
|
|
|
path = self.root / "last.ckpt"
|
|
path.write_bytes(b"previous")
|
|
with _checkpoint_patches(self.pc):
|
|
with self.assertRaises(RuntimeError):
|
|
torch_io._atomic_save({"weights": Unsaveable()}, path)
|
|
self.assertEqual((path.read_bytes(), self._files()), (b"previous", ["last.ckpt"]))
|
|
with patch.object(self.pc.training_control, "replace_file",
|
|
side_effect=PermissionError(errno.EACCES, "in use")), \
|
|
self.assertRaises(PermissionError): # Lightning's own save ignored this
|
|
torch_io._atomic_save({"step": 2}, path)
|
|
self.assertEqual((path.read_bytes(), self._files()), (b"previous", ["last.ckpt"]))
|
|
torch_io._atomic_save({"step": 3}, path)
|
|
self.assertEqual(torch.load(path, weights_only=False), {"step": 3})
|
|
self.assertEqual(self._files(), ["last.ckpt"])
|
|
|
|
def test_an_old_checkpoint_open_elsewhere_is_deleted_once_closed(self) -> None:
|
|
from lightning.fabric.plugins.io import torch_io
|
|
attempts = []
|
|
|
|
def busy_once(checkpoint_io, path):
|
|
attempts.append(path)
|
|
if len(attempts) == 1:
|
|
raise PermissionError(errno.EACCES, "in use")
|
|
|
|
with patch.object(torch_io.TorchCheckpointIO, "remove_checkpoint", busy_once), \
|
|
_checkpoint_patches(self.pc):
|
|
torch_io.TorchCheckpointIO().remove_checkpoint("old.ckpt")
|
|
self.assertEqual(attempts, ["old.ckpt", "old.ckpt"])
|
|
|
|
|
|
class StopOnRequestTests(unittest.TestCase):
|
|
"""Training saves its progress and stops when asked to, or when its launcher or console goes."""
|
|
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.stop_file = Path(d.name) / ".stop_training"
|
|
self.checkpoints = Path(d.name) / "checkpoints"
|
|
|
|
def _fit(self, model, trainer, stopper, console=None) -> tuple[int, str]:
|
|
rc, out, _ = _captured(lambda: self.pc._fit(lambda: trainer.fit(model), stopper, console))
|
|
return rc, out
|
|
|
|
def _last(self) -> dict:
|
|
import torch
|
|
return torch.load(self.checkpoints / "last.ckpt", map_location="cpu", weights_only=False)
|
|
|
|
def _resume(self, *, epochs: int = 2):
|
|
"""The model of a run resumed from last.ckpt, once that has finished."""
|
|
model, trainer = _toy_training(self.checkpoints, epochs=epochs, resumable=True)
|
|
_captured(lambda: trainer.fit(model, ckpt_path=self.checkpoints / "last.ckpt"))
|
|
return model
|
|
|
|
def _stop_at(self, stopper, epoch: int, batch: int):
|
|
def on_step(model, batch_idx):
|
|
if (model.current_epoch, batch_idx) == (epoch, batch):
|
|
self.stop_file.touch()
|
|
stopper._next_poll = 0.0 # look now, not within the second
|
|
return on_step
|
|
|
|
def test_stop_mid_epoch_saves_the_epoch_so_far_and_resumes_with_its_next_batch(self) -> None:
|
|
stopper = self.pc.StopOnRequest(str(self.stop_file))
|
|
with _checkpoint_patches(self.pc):
|
|
model, trainer = _toy_training(
|
|
self.checkpoints, epochs=2, resumable=True, on_step=self._stop_at(stopper, 0, 1),
|
|
callbacks=[stopper, self.pc.ProgressSaver(stopper)])
|
|
rc, out = self._fit(model, trainer, stopper)
|
|
self.assertEqual((rc, model.steps), (EXIT_STOPPED, [(0, 0), (0, 1)]))
|
|
self.assertIn("Stop requested", out)
|
|
self.assertIn("Saved the training progress: epoch 0, batch 2/4.", out)
|
|
last = self._last()
|
|
self.assertEqual((last["epoch"], last["global_step"]), (0, 2))
|
|
self.assertEqual(self.pc._epoch_progress(last),
|
|
{"batches_done": 2, "batches_per_epoch": 4})
|
|
self.assertEqual([p.name for p in self.checkpoints.iterdir()], ["last.ckpt"])
|
|
resumed = self._resume()
|
|
self.assertEqual(resumed.steps, [(0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3)])
|
|
self.assertCountEqual(model.items + resumed.items[:2], range(4)) # each clip once
|
|
epoch_1 = self.pc.LengthBucketBatchSampler([3, 1, 4, 2], 1, drop_last=False).batches(1)
|
|
self.assertEqual(resumed.items[2:], [clip for batch in epoch_1 for clip in batch])
|
|
self.assertTrue((self.checkpoints / "epoch=0-val_mel=1.0000.ckpt").exists()) # validated
|
|
|
|
def test_stop_at_the_end_of_an_epoch_validates_and_saves_it(self) -> None:
|
|
from lightning.pytorch.callbacks import Callback
|
|
stop_file = self.stop_file
|
|
|
|
class StopDuringValidation(Callback):
|
|
def on_validation_start(self, trainer, pl_module):
|
|
stop_file.touch()
|
|
|
|
for when in ("at its last batch", "during its validation"):
|
|
with self.subTest(when), _checkpoint_patches(self.pc):
|
|
stop_file.unlink(missing_ok=True)
|
|
shutil.rmtree(self.checkpoints, ignore_errors=True)
|
|
stopper = self.pc.StopOnRequest(str(stop_file))
|
|
extra = [StopDuringValidation()] if "validation" in when else []
|
|
model, trainer = _toy_training(
|
|
self.checkpoints, epochs=3, resumable=True,
|
|
on_step=None if extra else self._stop_at(stopper, 0, 3),
|
|
callbacks=[stopper, self.pc.ProgressSaver(stopper), *extra])
|
|
rc, out = self._fit(model, trainer, stopper)
|
|
self.assertEqual((rc, model.steps), (EXIT_STOPPED, [(0, i) for i in range(4)]))
|
|
self.assertNotIn("Saved the training progress", out)
|
|
last = self._last()
|
|
self.assertEqual((last["epoch"], last["global_step"]), (0, 4))
|
|
self.assertEqual(self.pc._epoch_progress(last),
|
|
{"batches_done": 4, "batches_per_epoch": 4})
|
|
self.assertTrue((self.checkpoints / "epoch=0-val_mel=1.0000.ckpt").exists())
|
|
self.assertEqual(self._resume().steps, [(1, i) for i in range(4)])
|
|
|
|
def test_a_crash_loses_only_the_training_since_the_last_save(self) -> None:
|
|
class Clock: # every look at piper_compat's clock is an hour later
|
|
now = 0.0
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(time, name)
|
|
|
|
def monotonic(self):
|
|
Clock.now += 3600
|
|
return Clock.now
|
|
|
|
def on_step(model, batch_idx):
|
|
if (model.current_epoch, batch_idx) == (1, 2):
|
|
raise RuntimeError("crash")
|
|
|
|
out = io.StringIO()
|
|
with _checkpoint_patches(self.pc), patch.object(self.pc, "time", Clock()):
|
|
model, trainer = _toy_training(self.checkpoints, epochs=2, resumable=True,
|
|
callbacks=[self.pc.ProgressSaver()], on_step=on_step)
|
|
with self.assertRaisesRegex(RuntimeError, "crash"), contextlib.redirect_stdout(out), \
|
|
contextlib.redirect_stderr(io.StringIO()), warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
trainer.fit(model)
|
|
# Saved after every step but an epoch's last: its checkpoint comes after validation.
|
|
self.assertEqual(re.findall(r"progress: epoch (\d), batch (\d)/4", out.getvalue()),
|
|
[("0", "1"), ("0", "2"), ("0", "3"), ("1", "1"), ("1", "2")])
|
|
self.assertEqual(self._last()["global_step"], 6)
|
|
with _checkpoint_patches(self.pc):
|
|
resumed = self._resume()
|
|
self.assertEqual(resumed.steps, [(1, 2), (1, 3)])
|
|
self.assertCountEqual(model.items[4:6] + resumed.items, range(4))
|
|
|
|
def test_closing_the_console_window_saves_before_windows_ends_training(self) -> None:
|
|
console = self.pc.training_control.ConsoleClose(timeout=30)
|
|
handled, handlers = [], []
|
|
last = self.checkpoints / "last.ckpt"
|
|
|
|
def close_window(): # what Windows does, in a thread of its own
|
|
handled.append(console.handle(2))
|
|
handled.append(last.exists()) # Windows ends the process once the handler returns
|
|
|
|
def on_step(model, batch_idx):
|
|
if batch_idx == 1:
|
|
handler = threading.Thread(target=close_window)
|
|
handler.start()
|
|
handlers.append(handler)
|
|
self.assertTrue(console.closing.wait(10))
|
|
|
|
with _checkpoint_patches(self.pc):
|
|
model, trainer = _toy_training(self.checkpoints, epochs=2, resumable=True,
|
|
callbacks=[self.pc.ProgressSaver(console=console)],
|
|
on_step=on_step)
|
|
rc, out = self._fit(model, trainer, None, console)
|
|
handlers[0].join(10)
|
|
self.assertEqual(handled, [True, True])
|
|
self.assertEqual((rc, model.steps), (EXIT_STOPPED, [(0, 0), (0, 1)]))
|
|
self.assertNotIn("Saved the training progress", out) # the console is going
|
|
self.assertEqual(self._last()["global_step"], 2)
|
|
self.assertEqual(self._resume().steps[:2], [(0, 2), (0, 3)])
|
|
|
|
def test_closing_the_console_during_a_due_save_ends_training_after_that_save(self) -> None:
|
|
from lightning.pytorch.callbacks import Callback
|
|
console = self.pc.training_control.ConsoleClose(timeout=30)
|
|
handled, handlers, saves = [], [], []
|
|
last = self.checkpoints / "last.ckpt"
|
|
|
|
def close_window():
|
|
handled.append(console.handle(2))
|
|
handled.append(last.exists())
|
|
|
|
class CloseWhileSaving(Callback):
|
|
def on_save_checkpoint(self, trainer, pl_module, checkpoint):
|
|
saves.append(trainer.global_step)
|
|
if not handlers:
|
|
handlers.append(threading.Thread(target=close_window))
|
|
handlers[0].start()
|
|
assert console.closing.wait(10)
|
|
|
|
class Clock: # every look is an hour later: a save is due after every step
|
|
now = 0.0
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(time, name)
|
|
|
|
def monotonic(self):
|
|
Clock.now += 3600
|
|
return Clock.now
|
|
|
|
with _checkpoint_patches(self.pc), patch.object(self.pc, "time", Clock()):
|
|
model, trainer = _toy_training(
|
|
self.checkpoints, epochs=2, resumable=True,
|
|
callbacks=[self.pc.ProgressSaver(console=console), CloseWhileSaving()])
|
|
rc, out = self._fit(model, trainer, None, console)
|
|
handlers[0].join(10)
|
|
self.assertEqual(handled, [True, True])
|
|
# Not another step and save: Windows allows only seconds.
|
|
self.assertEqual((rc, model.steps, saves), (EXIT_STOPPED, [(0, 0)], [1]))
|
|
self.assertNotIn("Saved the training progress", out)
|
|
self.assertEqual(self._last()["global_step"], 1)
|
|
|
|
def test_stop_requested_before_training_starts(self) -> None:
|
|
self.stop_file.touch()
|
|
stopper = self.pc.StopOnRequest(str(self.stop_file))
|
|
with _checkpoint_patches(self.pc):
|
|
model, trainer = _toy_training(self.checkpoints, epochs=3, callbacks=[stopper])
|
|
rc, _ = self._fit(model, trainer, stopper)
|
|
self.assertEqual((rc, model.steps), (EXIT_STOPPED, []))
|
|
self.assertEqual(list(self.checkpoints.glob("*.ckpt")), [])
|
|
|
|
def test_looks_for_the_stop_file_at_most_once_a_second(self) -> None:
|
|
stopper = self.pc.StopOnRequest(str(self.stop_file))
|
|
self.assertFalse(stopper.requested())
|
|
self.stop_file.touch()
|
|
self.assertFalse(stopper.requested()) # looked less than a second ago
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
self.assertTrue(stopper.requested(force=True))
|
|
self.assertEqual(stopper.reason, "stop requested")
|
|
|
|
def test_training_stops_when_its_launcher_has_gone(self) -> None:
|
|
me = self.pc.training_control.this_process()
|
|
self.assertFalse(self.pc.StopOnRequest(launcher=(me["pid"], me["start"])).requested())
|
|
stopper = self.pc.StopOnRequest(launcher=_exited_process())
|
|
with contextlib.redirect_stdout(io.StringIO()) as out:
|
|
self.assertTrue(stopper.requested())
|
|
with self.assertRaises(self.pc.StopTraining):
|
|
stopper.raise_if_requested()
|
|
self.assertIn("has exited", out.getvalue())
|
|
|
|
def test_from_environment(self) -> None:
|
|
tc = self.pc.training_control
|
|
with patch.dict(os.environ, {tc.STOP_FILE_ENV: "", tc.LAUNCHER_ENV: ""}):
|
|
self.assertIsNone(self.pc.StopOnRequest.from_environment())
|
|
with patch.dict(os.environ, {tc.STOP_FILE_ENV: "stop", tc.LAUNCHER_ENV: "12:345"}):
|
|
stopper = self.pc.StopOnRequest.from_environment()
|
|
self.assertEqual((stopper.stop_file, stopper.launcher), ("stop", (12, 345)))
|
|
|
|
def test_exit_codes(self) -> None:
|
|
def raises(exc):
|
|
def main():
|
|
raise exc
|
|
return main
|
|
|
|
stopped = self.pc.StopOnRequest()
|
|
stopped.reason = "stop requested"
|
|
for main, stopper, expected in (
|
|
(lambda: None, None, 0), (lambda: None, self.pc.StopOnRequest(), 0),
|
|
(lambda: None, stopped, EXIT_STOPPED),
|
|
(raises(self.pc.StopTraining("stop requested")), None, EXIT_STOPPED),
|
|
(raises(self.pc.TrainingDivergedError("NaN")), None, EXIT_DIVERGED)):
|
|
with self.subTest(expected=expected, stopper=stopper):
|
|
self.assertEqual(_captured(functools.partial(self.pc._fit, main, stopper))[0],
|
|
expected)
|
|
with self.assertRaises(ValueError): # a crash: the launcher restarts training
|
|
self.pc._fit(raises(ValueError("bug")))
|
|
err = io.StringIO()
|
|
with contextlib.redirect_stderr(err), self.assertRaises(self.pc._OUT_OF_MEMORY):
|
|
self.pc._fit(raises(self.pc._OUT_OF_MEMORY("CUDA out of memory")))
|
|
self.assertIn("smaller --batch-size", err.getvalue())
|
|
closed = SimpleNamespace(closing=threading.Event()) # ended at an epoch's end
|
|
closed.closing.set()
|
|
self.assertEqual(self.pc._fit(lambda: None, None, closed), EXIT_STOPPED)
|
|
|
|
|
|
class ProgressSaverTests(unittest.TestCase):
|
|
"""When the progress within an epoch is saved."""
|
|
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
|
|
def _trainer(self, *, batches: float = 10, keeps_last: bool = True):
|
|
"""``(trainer, its save of last.ckpt)``."""
|
|
save = MagicMock()
|
|
callback = SimpleNamespace(save_last=keeps_last, _save_last_checkpoint=save,
|
|
_monitor_candidates=lambda trainer: {})
|
|
return SimpleNamespace(num_training_batches=batches, current_epoch=3, should_stop=False,
|
|
checkpoint_callbacks=[callback]), save
|
|
|
|
def test_every_few_minutes_and_before_an_epoch_ends(self) -> None:
|
|
saver = self.pc.ProgressSaver(minutes=10)
|
|
trainer, save = self._trainer()
|
|
batch_end = functools.partial(saver.on_train_batch_end, trainer, None, None, None)
|
|
saver._last_save = time.monotonic() - 400
|
|
with contextlib.redirect_stdout(io.StringIO()) as out:
|
|
batch_end(5)
|
|
save.assert_not_called()
|
|
batch_end(8) # half due: validation and the epoch's checkpoint come next
|
|
self.assertEqual(save.call_count, 1)
|
|
saver._last_save = time.monotonic() - 601
|
|
batch_end(9) # the epoch's last batch: never
|
|
self.assertEqual(save.call_count, 1)
|
|
batch_end(0)
|
|
self.assertEqual(save.call_count, 2)
|
|
self.assertIn("epoch 3, batch 9/10", out.getvalue())
|
|
self.assertFalse(trainer.should_stop)
|
|
for minutes in (0, -5):
|
|
with self.subTest(minutes=minutes):
|
|
saver = self.pc.ProgressSaver(minutes=minutes)
|
|
saver._last_save = -math.inf
|
|
trainer, save = self._trainer()
|
|
saver.on_train_batch_end(trainer, None, None, None, 5)
|
|
save.assert_not_called()
|
|
|
|
def test_stop_without_a_way_to_save_mid_epoch(self) -> None:
|
|
stopper = self.pc.StopOnRequest()
|
|
stopper.reason = "stop requested"
|
|
for batches, keeps_last in ((10, False), (math.inf, True)): # inf: an iterable dataset
|
|
with self.subTest(batches=batches, keeps_last=keeps_last):
|
|
saver = self.pc.ProgressSaver(stopper)
|
|
trainer, save = self._trainer(batches=batches, keeps_last=keeps_last)
|
|
with contextlib.redirect_stdout(io.StringIO()) as out:
|
|
saver.on_train_batch_end(trainer, None, None, None, 2)
|
|
saver.on_train_batch_end(trainer, None, None, None, 3)
|
|
save.assert_not_called()
|
|
self.assertTrue(trainer.should_stop) # validates and saves at the epoch's end
|
|
self.assertEqual(out.getvalue().count("progress within an epoch is not saved"),
|
|
0 if keeps_last else 1)
|
|
|
|
def test_a_closing_console_is_let_go_even_if_saving_fails(self) -> None:
|
|
console = SimpleNamespace(closing=threading.Event(), reason="closed", done=MagicMock())
|
|
console.closing.set()
|
|
saver = self.pc.ProgressSaver(console=console)
|
|
trainer, save = self._trainer()
|
|
save.side_effect = OSError("disk full")
|
|
with self.assertRaises(OSError):
|
|
saver.on_train_batch_end(trainer, None, None, None, 2)
|
|
console.done.assert_called_once_with()
|
|
|
|
def test_minutes_from_the_environment(self) -> None:
|
|
for value, expected in (("5", 5.0), ("0.5", 0.5), ("0", 0.0), ("-1", 0.0),
|
|
("nan", 0.0), ("inf", 0.0), ("ten", 10.0), (None, 10.0)):
|
|
with self.subTest(value), patch.dict(os.environ):
|
|
os.environ.pop("JAVIS_PIPER_PROGRESS_MINUTES", None)
|
|
if value is not None:
|
|
os.environ["JAVIS_PIPER_PROGRESS_MINUTES"] = value
|
|
self.assertEqual(self.pc._progress_minutes(), expected)
|
|
|
|
|
|
class CacheRobustnessTests(unittest.TestCase):
|
|
"""A damaged preprocessing cache file is rebuilt or passed over, never fatal."""
|
|
|
|
def setUp(self) -> None:
|
|
self.pc = _piper_compat()
|
|
d = tempfile.TemporaryDirectory()
|
|
self.addCleanup(d.cleanup)
|
|
self.root = Path(d.name)
|
|
|
|
def test_piper_uses_the_safe_cache_functions(self) -> None:
|
|
import torch
|
|
from piper.train.vits import dataset
|
|
self.assertIsInstance(dataset.torch, self.pc._CacheTorch)
|
|
self.assertIs(dataset.torch.LongTensor, torch.LongTensor)
|
|
|
|
def test_cache_files_are_written_whole(self) -> None:
|
|
import torch
|
|
cache, path = self.pc._CacheTorch(), self.root / "a.pt"
|
|
written = MagicMock() # stands in for StopOnRequest.raise_if_requested
|
|
with patch.object(self.pc, "_on_cache_write", written):
|
|
cache.save(torch.ones(2), path)
|
|
written.assert_called_once_with()
|
|
self.assertTrue(torch.equal(cache.load(path), torch.ones(2)))
|
|
|
|
class Unsaveable:
|
|
def __reduce__(self):
|
|
raise RuntimeError("disk full")
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
cache.save([Unsaveable()], self.root / "b.pt")
|
|
self.assertEqual([p.name for p in self.root.iterdir()], ["a.pt"])
|
|
|
|
def test_a_damaged_cache_file_is_deleted_to_be_rebuilt(self) -> None:
|
|
import torch
|
|
cache, path = self.pc._CacheTorch(), self.root / "a.pt"
|
|
torch.save(torch.ones(1000), path)
|
|
_truncate(path)
|
|
with self.assertRaises(self.pc.CorruptCacheFile):
|
|
cache.load(path)
|
|
self.assertFalse(path.exists())
|
|
with self.assertRaises(FileNotFoundError): # missing isn't damaged
|
|
cache.load(path)
|
|
|
|
def test_what_counts_as_damaged_content(self) -> None:
|
|
import pickle
|
|
damaged = (RuntimeError("PytorchStreamReader failed reading zip archive"), EOFError(),
|
|
pickle.UnpicklingError("invalid load key"),
|
|
UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid start byte"),
|
|
OSError(errno.EINVAL, "Invalid argument"))
|
|
not_damaged = (FileNotFoundError(errno.ENOENT, "missing"),
|
|
PermissionError(errno.EACCES, "in use"), MemoryError(),
|
|
RuntimeError("CUDA out of memory"),
|
|
RuntimeError("DefaultCPUAllocator: not enough memory"), KeyboardInterrupt())
|
|
for exc in damaged + not_damaged:
|
|
with self.subTest(exc=exc):
|
|
self.assertEqual(self.pc._unreadable_content(exc), exc in damaged)
|
|
|
|
def test_unreadable_utterances_are_replaced_by_readable_ones(self) -> None:
|
|
pc, unreadable = self.pc, {2, 3}
|
|
|
|
class Utterances:
|
|
def __getitem__(self, i):
|
|
if i in unreadable:
|
|
raise pc.CorruptCacheFile(i) if i == 2 else FileNotFoundError(i)
|
|
return i
|
|
|
|
part = pc.RobustSubset(Utterances(), [0, 1, 2, 3, 4])
|
|
with patch.object(pc, "_unreadable_utterances", 0), \
|
|
contextlib.redirect_stderr(io.StringIO()) as err:
|
|
self.assertEqual([part[i] for i in range(5)], [0, 1, 4, 4, 4])
|
|
self.assertEqual(part.__getitems__([2, 0]), [4, 0])
|
|
self.assertEqual(part[-1], 4)
|
|
with self.assertRaises(IndexError): # ends piper's loop over the test split
|
|
part[5]
|
|
unreadable.update(range(5))
|
|
with self.assertRaises(pc.CorruptCacheFile): # nothing readable: the first error
|
|
part[2]
|
|
self.assertIn("used another utterance", err.getvalue())
|
|
|
|
def test_data_module_rebuilds_damaged_files(self) -> None:
|
|
from torch.utils.data import Subset
|
|
pc, root, attempts = self.pc, self.root, []
|
|
(root / "x.pt.partial").write_bytes(b"cut short") # from a run killed while writing
|
|
|
|
class DataModule:
|
|
cache_dir = root
|
|
damaged = 2
|
|
|
|
def prepare_data(self):
|
|
attempts.append(sorted(p.name for p in root.iterdir()))
|
|
if len(attempts) <= self.damaged:
|
|
raise pc.CorruptCacheFile(f"damaged file {len(attempts)}")
|
|
return "prepared"
|
|
|
|
def setup(self, stage):
|
|
self.train_dataset = Subset(range(4), [0, 1])
|
|
self.val_dataset = Subset(range(4), [2])
|
|
self.test_dataset = None
|
|
|
|
setup = DataModule.setup
|
|
pc._robust_dataset(DataModule)
|
|
pc._robust_dataset(DataModule)
|
|
self.assertIs(DataModule.setup.__wrapped__, setup) # patched once
|
|
dm = DataModule()
|
|
with contextlib.redirect_stderr(io.StringIO()) as err:
|
|
self.assertEqual(dm.prepare_data(), "prepared")
|
|
self.assertEqual(attempts, [[]] * 3)
|
|
self.assertEqual(err.getvalue().count("WARNING: damaged file"), 2)
|
|
dm.setup("fit")
|
|
self.assertIsInstance(dm.train_dataset, pc.RobustSubset)
|
|
self.assertIsInstance(dm.val_dataset, pc.RobustSubset)
|
|
self.assertEqual((dm.train_dataset.indices, dm.val_dataset.indices), ([0, 1], [2]))
|
|
self.assertIsNone(dm.test_dataset)
|
|
|
|
attempts.clear()
|
|
DataModule.damaged = 100
|
|
with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(pc.CorruptCacheFile):
|
|
dm.prepare_data()
|
|
self.assertEqual(len(attempts), 50)
|
|
|
|
|
|
class TranscriptGateTests(unittest.TestCase):
|
|
def test_unreliable_transcripts_are_rejected(self) -> None:
|
|
from javis.voices.audiobook_ingest import _transcript_ok
|
|
|
|
def tr(**kw):
|
|
return SimpleNamespace(**{"avg_logprob": -0.2, "no_speech_prob": 0.01,
|
|
"compression_ratio": 1.4, **kw})
|
|
self.assertTrue(_transcript_ok(tr()))
|
|
self.assertTrue(_transcript_ok(tr(avg_logprob=None, no_speech_prob=None,
|
|
compression_ratio=None)))
|
|
self.assertFalse(_transcript_ok(tr(avg_logprob=-1.3)))
|
|
self.assertFalse(_transcript_ok(tr(no_speech_prob=0.9)))
|
|
self.assertFalse(_transcript_ok(tr(compression_ratio=3.1)))
|
|
|
|
|
|
class NarratorFilterTests(unittest.TestCase):
|
|
def test_dominant_speaker_scores_high(self) -> None:
|
|
from javis.voices.dataset_clean import narrator_scores
|
|
|
|
rng = np.random.default_rng(0)
|
|
a, b = np.eye(8)[0], np.eye(8)[1]
|
|
embs = np.array([v + rng.normal(0, 0.1, 8) for v in [a] * 70 + [b] * 30])
|
|
embs /= np.linalg.norm(embs, axis=1, keepdims=True)
|
|
s = narrator_scores(embs)
|
|
self.assertGreater(s[:70].min(), 0.8)
|
|
self.assertLess(s[70:].max(), 0.4)
|
|
|
|
def test_clean_dataset_filters_and_keeps_new_imports(self) -> None:
|
|
import soundfile as sf
|
|
from javis.voices.dataset_clean import clean_dataset
|
|
|
|
def embed(audio, sr):
|
|
# The first sample marks the speaker: 0.5 = narrator, else character voice.
|
|
return np.eye(4)[0] if audio[0] > 0.4 else np.eye(4)[1]
|
|
|
|
with tempfile.TemporaryDirectory() as d:
|
|
training = Path(d)
|
|
(training / "wavs").mkdir()
|
|
|
|
def add(stem: str, narrator: bool) -> str:
|
|
sf.write(str(training / "wavs" / f"{stem}.wav"),
|
|
np.full(100, 0.5 if narrator else 0.1, dtype=np.float32), 22050)
|
|
return f"{stem}|text {stem}|text {stem}\n"
|
|
|
|
lines = [add(f"clip_{i:05d}", narrator=i % 4 != 0) for i in range(8)]
|
|
(training / "metadata.csv").write_text("".join(lines), encoding="utf-8")
|
|
|
|
report = clean_dataset(training, embed, dry_run=True)
|
|
self.assertEqual((report["kept"], report["dropped"]), (6, 2))
|
|
self.assertEqual(len((training / "metadata.csv").read_text().splitlines()), 8)
|
|
|
|
clean_dataset(training, embed)
|
|
self.assertEqual(len((training / "metadata.csv").read_text().splitlines()), 6)
|
|
self.assertEqual(len((training / "metadata.unfiltered.csv").read_text().splitlines()), 8)
|
|
|
|
# A later import appends to metadata.csv; re-filtering must keep it.
|
|
with open(training / "metadata.csv", "a", encoding="utf-8") as f:
|
|
f.write(add("clip_00100", narrator=True))
|
|
report = clean_dataset(training, embed)
|
|
self.assertEqual((report["clips"], report["kept"]), (9, 7))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(unittest.main())
|