mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Paul Sajna
f86b44d8fc
* 48KHz softvc * add some notes * almost forgot the inference script * fix readme link * requirements.txt
59 lines
1.3 KiB
Python
59 lines
1.3 KiB
Python
import torch
|
|
|
|
|
|
class Metric:
|
|
def __init__(self):
|
|
self.steps = 0
|
|
self.value = 0
|
|
|
|
def update(self, value):
|
|
self.steps += 1
|
|
self.value += (value - self.value) / self.steps
|
|
return self.value
|
|
|
|
def reset(self):
|
|
self.steps = 0
|
|
self.value = 0
|
|
|
|
|
|
def save_checkpoint(
|
|
checkpoint_dir,
|
|
hubert,
|
|
optimizer,
|
|
scaler,
|
|
step,
|
|
loss,
|
|
best,
|
|
logger,
|
|
):
|
|
state = {
|
|
"hubert": hubert.state_dict(),
|
|
"optimizer": optimizer.state_dict(),
|
|
"scaler": scaler.state_dict(),
|
|
"step": step,
|
|
"loss": loss,
|
|
}
|
|
checkpoint_dir.mkdir(exist_ok=True, parents=True)
|
|
checkpoint_path = checkpoint_dir / f"model-{step}.pt"
|
|
torch.save(state, checkpoint_path)
|
|
if best:
|
|
best_path = checkpoint_dir / "model-best.pt"
|
|
torch.save(state, best_path)
|
|
logger.info(f"Saved checkpoint: {checkpoint_path.stem}")
|
|
|
|
|
|
def load_checkpoint(
|
|
load_path,
|
|
hubert,
|
|
optimizer,
|
|
scaler,
|
|
rank,
|
|
logger,
|
|
):
|
|
logger.info(f"Loading checkpoint from {load_path}")
|
|
checkpoint = torch.load(load_path, map_location={"cuda:0": f"cuda:{rank}"})
|
|
hubert.load_state_dict(checkpoint["hubert"])
|
|
scaler.load_state_dict(checkpoint["scaler"])
|
|
optimizer.load_state_dict(checkpoint["optimizer"])
|
|
return checkpoint["step"], checkpoint["loss"]
|