Initial commit
This commit is contained in:
Yatharth Sharma
570
zipvoice/utils/checkpoint.py
Normal file
570
zipvoice/utils/checkpoint.py
Normal file
@@ -0,0 +1,570 @@
|
||||
# Copyright 2021-2025 Xiaomi Corporation (authors: Fangjun Kuang,
|
||||
# Zengwei Yao)
|
||||
#
|
||||
# See ../../LICENSE for clarification regarding multiple authors
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import glob
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from lhotse.dataset.sampling.base import CutSampler
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from torch.optim import Optimizer
|
||||
|
||||
from zipvoice.utils.common import AttributeDict, GradScaler
|
||||
|
||||
# use duck typing for LRScheduler since we have different possibilities, see
|
||||
# our class LRScheduler.
|
||||
LRSchedulerType = object
|
||||
|
||||
|
||||
def save_checkpoint(
|
||||
filename: Path,
|
||||
model: Union[nn.Module, DDP],
|
||||
model_avg: Optional[nn.Module] = None,
|
||||
model_ema: Optional[nn.Module] = None,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
optimizer: Optional[Optimizer] = None,
|
||||
scheduler: Optional[LRSchedulerType] = None,
|
||||
scaler: Optional[GradScaler] = None,
|
||||
sampler: Optional[CutSampler] = None,
|
||||
rank: int = 0,
|
||||
) -> None:
|
||||
"""Save training information to a file.
|
||||
|
||||
Args:
|
||||
filename:
|
||||
The checkpoint filename.
|
||||
model:
|
||||
The model to be saved. We only save its `state_dict()`.
|
||||
model_avg:
|
||||
The stored model averaged from the start of training.
|
||||
model_ema:
|
||||
The EMA version of model.
|
||||
params:
|
||||
User defined parameters, e.g., epoch, loss.
|
||||
optimizer:
|
||||
The optimizer to be saved. We only save its `state_dict()`.
|
||||
scheduler:
|
||||
The scheduler to be saved. We only save its `state_dict()`.
|
||||
scalar:
|
||||
The GradScaler to be saved. We only save its `state_dict()`.
|
||||
sampler:
|
||||
The sampler used in the labeled training dataset. We only
|
||||
save its `state_dict()`.
|
||||
rank:
|
||||
Used in DDP. We save checkpoint only for the node whose
|
||||
rank is 0.
|
||||
Returns:
|
||||
Return None.
|
||||
"""
|
||||
if rank != 0:
|
||||
return
|
||||
|
||||
logging.info(f"Saving checkpoint to {filename}")
|
||||
|
||||
if isinstance(model, DDP):
|
||||
model = model.module
|
||||
|
||||
checkpoint = {
|
||||
"model": model.state_dict(),
|
||||
"optimizer": optimizer.state_dict() if optimizer is not None else None,
|
||||
"scheduler": scheduler.state_dict() if scheduler is not None else None,
|
||||
"grad_scaler": scaler.state_dict() if scaler is not None else None,
|
||||
"sampler": sampler.state_dict() if sampler is not None else None,
|
||||
}
|
||||
|
||||
if model_avg is not None:
|
||||
checkpoint["model_avg"] = model_avg.to(torch.float32).state_dict()
|
||||
if model_ema is not None:
|
||||
checkpoint["model_ema"] = model_ema.to(torch.float32).state_dict()
|
||||
|
||||
if params:
|
||||
for k, v in params.items():
|
||||
assert k not in checkpoint
|
||||
checkpoint[k] = v
|
||||
|
||||
torch.save(checkpoint, filename)
|
||||
|
||||
|
||||
def load_checkpoint(
|
||||
filename: Path,
|
||||
model: Optional[nn.Module] = None,
|
||||
model_avg: Optional[nn.Module] = None,
|
||||
model_ema: Optional[nn.Module] = None,
|
||||
strict: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
logging.info(f"Loading checkpoint from {filename}")
|
||||
checkpoint = torch.load(filename, map_location="cpu", weights_only=False)
|
||||
|
||||
if model is not None:
|
||||
|
||||
if next(iter(checkpoint["model"])).startswith("module."):
|
||||
logging.debug("Loading checkpoint saved by DDP")
|
||||
dst_state_dict = model.state_dict()
|
||||
src_state_dict = checkpoint["model"]
|
||||
for key in dst_state_dict.keys():
|
||||
src_key = "{}.{}".format("module", key)
|
||||
dst_state_dict[key] = src_state_dict.pop(src_key)
|
||||
assert len(src_state_dict) == 0
|
||||
model.load_state_dict(dst_state_dict, strict=strict)
|
||||
else:
|
||||
logging.debug("Loading checkpoint")
|
||||
model.load_state_dict(checkpoint["model"], strict=strict)
|
||||
|
||||
checkpoint.pop("model")
|
||||
|
||||
if model_avg is not None and "model_avg" in checkpoint:
|
||||
logging.info("Loading averaged model")
|
||||
model_avg.load_state_dict(checkpoint["model_avg"], strict=strict)
|
||||
checkpoint.pop("model_avg")
|
||||
|
||||
if model_ema is not None and "model_ema" in checkpoint:
|
||||
logging.info("Loading ema model")
|
||||
model_ema.load_state_dict(checkpoint["model_ema"], strict=strict)
|
||||
checkpoint.pop("model_ema")
|
||||
|
||||
return checkpoint
|
||||
|
||||
|
||||
def load_checkpoint_extend_vocab_size(
|
||||
filename: Path, extend_size: int, model: nn.Module, strict: bool = True
|
||||
) -> Dict[str, Any]:
|
||||
logging.info(f"Loading checkpoint from {filename}")
|
||||
checkpoint = torch.load(filename, map_location="cpu", weights_only=False)
|
||||
|
||||
if model is not None:
|
||||
if next(iter(checkpoint["model"])).startswith("module."):
|
||||
logging.info("Loading checkpoint saved by DDP")
|
||||
dst_state_dict = model.state_dict()
|
||||
src_state_dict = checkpoint["model"]
|
||||
for key in dst_state_dict.keys():
|
||||
src_key = "{}.{}".format("module", key)
|
||||
dst_state_dict[key] = src_state_dict.pop(src_key)
|
||||
assert len(src_state_dict) == 0
|
||||
else:
|
||||
logging.info("Loading checkpoint")
|
||||
dst_state_dict = checkpoint["model"]
|
||||
dst_state_dict["spk_embed.weight"] = model.state_dict()["spk_embed.weight"]
|
||||
embed_weight = model.state_dict()["embed.weight"]
|
||||
embed_weight[:-extend_size, :] = dst_state_dict["embed.weight"]
|
||||
dst_state_dict["embed.weight"] = embed_weight
|
||||
|
||||
model.load_state_dict(dst_state_dict, strict=strict)
|
||||
|
||||
|
||||
def load_checkpoint_copy_proj_three_channel_alter(
|
||||
filename: Path,
|
||||
in_proj_key: str,
|
||||
out_proj_key: str,
|
||||
dim: int,
|
||||
model: nn.Module,
|
||||
) -> Dict[str, Any]:
|
||||
logging.info(f"Loading checkpoint from {filename}")
|
||||
checkpoint = torch.load(filename, map_location="cpu", weights_only=False)
|
||||
|
||||
if model is not None:
|
||||
if next(iter(checkpoint["model"])).startswith("module."):
|
||||
logging.info("Loading checkpoint saved by DDP")
|
||||
|
||||
dst_state_dict = dict()
|
||||
src_state_dict = checkpoint["model"]
|
||||
for key in src_state_dict.keys():
|
||||
dst_state_dict[key.lstrip("module.")] = src_state_dict.pop(key)
|
||||
assert len(src_state_dict) == 0
|
||||
else:
|
||||
logging.info("Loading checkpoint")
|
||||
dst_state_dict = checkpoint["model"]
|
||||
keys = list(dst_state_dict.keys())
|
||||
for key in keys:
|
||||
if in_proj_key in key:
|
||||
if "weight" in key:
|
||||
weight = dst_state_dict.pop(key)
|
||||
dst_state_dict[key.replace("weight", "0.weight")] = torch.cat(
|
||||
[
|
||||
weight[:, :dim] / 2,
|
||||
weight[:, :dim] / 2,
|
||||
weight[:, dim : dim * 2],
|
||||
weight[:, dim * 2 :] / 2,
|
||||
weight[:, dim * 2 :] / 2,
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
dst_state_dict[key.replace("weight", "1.weight")] = weight
|
||||
if "bias" in key:
|
||||
bias = dst_state_dict.pop(key)
|
||||
dst_state_dict[key.replace("bias", "0.bias")] = bias
|
||||
dst_state_dict[key.replace("bias", "1.bias")] = bias
|
||||
if out_proj_key in key:
|
||||
if "weight" in key:
|
||||
weight = dst_state_dict.pop(key)
|
||||
dst_state_dict[key.replace("weight", "0.weight")] = torch.cat(
|
||||
[weight, weight], dim=0
|
||||
)
|
||||
dst_state_dict[key.replace("weight", "1.weight")] = weight
|
||||
elif "bias" in key:
|
||||
bias = dst_state_dict.pop(key)
|
||||
dst_state_dict[key.replace("bias", "0.bias")] = torch.cat(
|
||||
[bias, bias], dim=0
|
||||
)
|
||||
dst_state_dict[key.replace("bias", "1.bias")] = bias
|
||||
|
||||
model.load_state_dict(dst_state_dict, strict=True)
|
||||
|
||||
|
||||
def find_checkpoints(out_dir: Path, iteration: int = 0) -> List[str]:
|
||||
"""Find all available checkpoints in a directory.
|
||||
|
||||
The checkpoint filenames have the form: `checkpoint-xxx.pt`
|
||||
where xxx is a numerical value.
|
||||
|
||||
Assume you have the following checkpoints in the folder `foo`:
|
||||
|
||||
- checkpoint-1.pt
|
||||
- checkpoint-20.pt
|
||||
- checkpoint-300.pt
|
||||
- checkpoint-4000.pt
|
||||
|
||||
Case 1 (Return all checkpoints)::
|
||||
|
||||
find_checkpoints(out_dir='foo')
|
||||
|
||||
Case 2 (Return checkpoints newer than checkpoint-20.pt, i.e.,
|
||||
checkpoint-4000.pt, checkpoint-300.pt, and checkpoint-20.pt)
|
||||
|
||||
find_checkpoints(out_dir='foo', iteration=20)
|
||||
|
||||
Case 3 (Return checkpoints older than checkpoint-20.pt, i.e.,
|
||||
checkpoint-20.pt, checkpoint-1.pt)::
|
||||
|
||||
find_checkpoints(out_dir='foo', iteration=-20)
|
||||
|
||||
Args:
|
||||
out_dir:
|
||||
The directory where to search for checkpoints.
|
||||
iteration:
|
||||
If it is 0, return all available checkpoints.
|
||||
If it is positive, return the checkpoints whose iteration number is
|
||||
greater than or equal to `iteration`.
|
||||
If it is negative, return the checkpoints whose iteration number is
|
||||
less than or equal to `-iteration`.
|
||||
Returns:
|
||||
Return a list of checkpoint filenames, sorted in descending
|
||||
order by the numerical value in the filename.
|
||||
"""
|
||||
checkpoints = list(glob.glob(f"{out_dir}/checkpoint-[0-9]*.pt"))
|
||||
pattern = re.compile(r"checkpoint-([0-9]+).pt")
|
||||
iter_checkpoints = []
|
||||
for c in checkpoints:
|
||||
result = pattern.search(c)
|
||||
if not result:
|
||||
logging.warn(f"Invalid checkpoint filename {c}")
|
||||
continue
|
||||
|
||||
iter_checkpoints.append((int(result.group(1)), c))
|
||||
|
||||
# iter_checkpoints is a list of tuples. Each tuple contains
|
||||
# two elements: (iteration_number, checkpoint-iteration_number.pt)
|
||||
|
||||
iter_checkpoints = sorted(iter_checkpoints, reverse=True, key=lambda x: x[0])
|
||||
if iteration >= 0:
|
||||
ans = [ic[1] for ic in iter_checkpoints if ic[0] >= iteration]
|
||||
else:
|
||||
ans = [ic[1] for ic in iter_checkpoints if ic[0] <= -iteration]
|
||||
|
||||
return ans
|
||||
|
||||
|
||||
def average_checkpoints_with_averaged_model(
|
||||
filename_start: str,
|
||||
filename_end: str,
|
||||
device: torch.device = torch.device("cpu"),
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Average model parameters over the range with given
|
||||
start model (excluded) and end model.
|
||||
|
||||
Let start = batch_idx_train of model-start;
|
||||
end = batch_idx_train of model-end;
|
||||
interval = end - start.
|
||||
Then the average model over range from start (excluded) to end is
|
||||
(1) avg = (model_end * end - model_start * start) / interval.
|
||||
It can be written as
|
||||
(2) avg = model_end * weight_end + model_start * weight_start,
|
||||
where weight_end = end / interval,
|
||||
weight_start = -start / interval = 1 - weight_end.
|
||||
Since the terms `weight_end` and `weight_start` would be large
|
||||
if the model has been trained for lots of batches, which would cause
|
||||
overflow when multiplying the model parameters.
|
||||
To avoid this, we rewrite (2) as:
|
||||
(3) avg = (model_end + model_start * (weight_start / weight_end))
|
||||
* weight_end
|
||||
|
||||
The model index could be epoch number or iteration number.
|
||||
|
||||
Args:
|
||||
filename_start:
|
||||
Checkpoint filename of the start model. We assume it
|
||||
is saved by :func:`save_checkpoint`.
|
||||
filename_end:
|
||||
Checkpoint filename of the end model. We assume it
|
||||
is saved by :func:`save_checkpoint`.
|
||||
device:
|
||||
Move checkpoints to this device before averaging.
|
||||
"""
|
||||
state_dict_start = torch.load(
|
||||
filename_start, map_location=device, weights_only=False
|
||||
)
|
||||
state_dict_end = torch.load(filename_end, map_location=device, weights_only=False)
|
||||
|
||||
average_period = state_dict_start["average_period"]
|
||||
|
||||
batch_idx_train_start = state_dict_start["batch_idx_train"]
|
||||
batch_idx_train_start = (batch_idx_train_start // average_period) * average_period
|
||||
batch_idx_train_end = state_dict_end["batch_idx_train"]
|
||||
batch_idx_train_end = (batch_idx_train_end // average_period) * average_period
|
||||
interval = batch_idx_train_end - batch_idx_train_start
|
||||
assert interval > 0, interval
|
||||
weight_end = batch_idx_train_end / interval
|
||||
weight_start = 1 - weight_end
|
||||
|
||||
model_end = state_dict_end["model_avg"]
|
||||
model_start = state_dict_start["model_avg"]
|
||||
avg = model_end
|
||||
|
||||
# scale the weight to avoid overflow
|
||||
average_state_dict(
|
||||
state_dict_1=avg,
|
||||
state_dict_2=model_start,
|
||||
weight_1=1.0,
|
||||
weight_2=weight_start / weight_end,
|
||||
scaling_factor=weight_end,
|
||||
)
|
||||
|
||||
return avg
|
||||
|
||||
|
||||
def remove_checkpoints(
|
||||
out_dir: Path,
|
||||
topk: int,
|
||||
rank: int = 0,
|
||||
):
|
||||
"""Remove checkpoints from the given directory.
|
||||
|
||||
We assume that checkpoint filename has the form `checkpoint-xxx.pt`
|
||||
where xxx is a number, representing the number of processed batches
|
||||
when saving that checkpoint. We sort checkpoints by filename and keep
|
||||
only the `topk` checkpoints with the highest `xxx`.
|
||||
|
||||
Args:
|
||||
out_dir:
|
||||
The directory containing checkpoints to be removed.
|
||||
topk:
|
||||
Number of checkpoints to keep.
|
||||
rank:
|
||||
If using DDP for training, it is the rank of the current node.
|
||||
Use 0 if no DDP is used for training.
|
||||
"""
|
||||
assert topk >= 1, topk
|
||||
if rank != 0:
|
||||
return
|
||||
checkpoints = find_checkpoints(out_dir)
|
||||
|
||||
if len(checkpoints) == 0:
|
||||
logging.warn(f"No checkpoints found in {out_dir}")
|
||||
return
|
||||
|
||||
if len(checkpoints) <= topk:
|
||||
return
|
||||
|
||||
to_remove = checkpoints[topk:]
|
||||
for c in to_remove:
|
||||
os.remove(c)
|
||||
|
||||
|
||||
def resume_checkpoint(
|
||||
params: AttributeDict,
|
||||
model: nn.Module,
|
||||
model_avg: nn.Module,
|
||||
model_ema: Optional[nn.Module] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Load checkpoint from file.
|
||||
|
||||
If params.start_epoch is larger than 1, it will load the checkpoint from
|
||||
`params.start_epoch - 1`.
|
||||
|
||||
Apart from loading state dict for `model` and `optimizer` it also updates
|
||||
`best_train_epoch`, `best_train_loss`, `best_valid_epoch`,
|
||||
and `best_valid_loss` in `params`.
|
||||
|
||||
Args:
|
||||
params:
|
||||
The return value of :func:`get_params`.
|
||||
model:
|
||||
The training model.
|
||||
Returns:
|
||||
Return a dict containing previously saved training info.
|
||||
"""
|
||||
filename = params.exp_dir / f"epoch-{params.start_epoch - 1}.pt"
|
||||
|
||||
assert filename.is_file(), f"{filename} does not exist!"
|
||||
|
||||
saved_params = load_checkpoint(
|
||||
filename,
|
||||
model=model,
|
||||
model_avg=model_avg,
|
||||
model_ema=model_ema,
|
||||
strict=True,
|
||||
)
|
||||
|
||||
if params.start_epoch > 1:
|
||||
keys = [
|
||||
"best_train_epoch",
|
||||
"best_valid_epoch",
|
||||
"batch_idx_train",
|
||||
"best_train_loss",
|
||||
"best_valid_loss",
|
||||
]
|
||||
for k in keys:
|
||||
params[k] = saved_params[k]
|
||||
|
||||
return saved_params
|
||||
|
||||
|
||||
def average_state_dict(
|
||||
state_dict_1: Dict[str, torch.Tensor],
|
||||
state_dict_2: Dict[str, torch.Tensor],
|
||||
weight_1: float,
|
||||
weight_2: float,
|
||||
scaling_factor: float = 1.0,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Average two state_dict with given weights:
|
||||
state_dict_1 = (state_dict_1 * weight_1 + state_dict_2 * weight_2)
|
||||
* scaling_factor
|
||||
It is an in-place operation on state_dict_1 itself.
|
||||
"""
|
||||
# Identify shared parameters. Two parameters are said to be shared
|
||||
# if they have the same data_ptr
|
||||
uniqued: Dict[int, str] = dict()
|
||||
for k, v in state_dict_1.items():
|
||||
v_data_ptr = v.data_ptr()
|
||||
if v_data_ptr in uniqued:
|
||||
continue
|
||||
uniqued[v_data_ptr] = k
|
||||
|
||||
uniqued_names = list(uniqued.values())
|
||||
for k in uniqued_names:
|
||||
v = state_dict_1[k]
|
||||
if torch.is_floating_point(v):
|
||||
v *= weight_1
|
||||
v += state_dict_2[k].to(device=state_dict_1[k].device) * weight_2
|
||||
v *= scaling_factor
|
||||
|
||||
|
||||
def update_averaged_model(
|
||||
params: Dict[str, torch.Tensor],
|
||||
model_cur: Union[nn.Module, DDP],
|
||||
model_avg: nn.Module,
|
||||
) -> None:
|
||||
"""Update the averaged model:
|
||||
model_avg = model_cur * (average_period / batch_idx_train)
|
||||
+ model_avg * ((batch_idx_train - average_period) / batch_idx_train)
|
||||
|
||||
Args:
|
||||
params:
|
||||
User defined parameters, e.g., epoch, loss.
|
||||
model_cur:
|
||||
The current model.
|
||||
model_avg:
|
||||
The averaged model to be updated.
|
||||
"""
|
||||
weight_cur = params.average_period / params.batch_idx_train
|
||||
weight_avg = 1 - weight_cur
|
||||
|
||||
if isinstance(model_cur, DDP):
|
||||
model_cur = model_cur.module
|
||||
|
||||
cur = model_cur.state_dict()
|
||||
avg = model_avg.state_dict()
|
||||
|
||||
average_state_dict(
|
||||
state_dict_1=avg,
|
||||
state_dict_2=cur,
|
||||
weight_1=weight_avg,
|
||||
weight_2=weight_cur,
|
||||
)
|
||||
|
||||
|
||||
def save_checkpoint_with_global_batch_idx(
|
||||
out_dir: Path,
|
||||
global_batch_idx: int,
|
||||
model: Union[nn.Module, DDP],
|
||||
model_avg: Optional[nn.Module] = None,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
optimizer: Optional[Optimizer] = None,
|
||||
scheduler: Optional[LRSchedulerType] = None,
|
||||
scaler: Optional[GradScaler] = None,
|
||||
sampler: Optional[CutSampler] = None,
|
||||
rank: int = 0,
|
||||
):
|
||||
"""Save training info after processing given number of batches.
|
||||
|
||||
Args:
|
||||
out_dir:
|
||||
The directory to save the checkpoint.
|
||||
global_batch_idx:
|
||||
The number of batches processed so far from the very start of the
|
||||
training. The saved checkpoint will have the following filename:
|
||||
|
||||
f'out_dir / checkpoint-{global_batch_idx}.pt'
|
||||
model:
|
||||
The neural network model whose `state_dict` will be saved in the
|
||||
checkpoint.
|
||||
model_avg:
|
||||
The stored model averaged from the start of training.
|
||||
params:
|
||||
A dict of training configurations to be saved.
|
||||
optimizer:
|
||||
The optimizer used in the training. Its `state_dict` will be saved.
|
||||
scheduler:
|
||||
The learning rate scheduler used in the training. Its `state_dict` will
|
||||
be saved.
|
||||
scaler:
|
||||
The scaler used for mix precision training. Its `state_dict` will
|
||||
be saved.
|
||||
sampler:
|
||||
The sampler used in the training dataset.
|
||||
rank:
|
||||
The rank ID used in DDP training of the current node. Set it to 0
|
||||
if DDP is not used.
|
||||
"""
|
||||
out_dir = Path(out_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
filename = out_dir / f"checkpoint-{global_batch_idx}.pt"
|
||||
save_checkpoint(
|
||||
filename=filename,
|
||||
model=model,
|
||||
model_avg=model_avg,
|
||||
params=params,
|
||||
optimizer=optimizer,
|
||||
scheduler=scheduler,
|
||||
scaler=scaler,
|
||||
sampler=sampler,
|
||||
rank=rank,
|
||||
)
|
||||
Reference in New Issue
Block a user