Initial commit

zipvoice/utils/checkpoint.py

@@ -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,
+    )

zipvoice/utils/common.py

@@ -0,0 +1,670 @@
+import argparse
+import collections
+import json
+import logging
+import os
+import socket
+import subprocess
+import sys
+import warnings
+from collections import defaultdict
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Tuple, Union
+
+import torch
+from packaging import version
+from torch import distributed as dist
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+
+
+if hasattr(torch.amp, "GradScaler"):
+    from torch.amp import GradScaler
+else:
+    from torch.cuda.amp import GradScaler
+
+Pathlike = Union[str, Path]
+
+
+class AttributeDict(dict):
+    def __getattr__(self, key):
+        if key in self:
+            return self[key]
+        raise AttributeError(f"No such attribute '{key}'")
+
+    def __setattr__(self, key, value):
+        self[key] = value
+
+    def __delattr__(self, key):
+        if key in self:
+            del self[key]
+            return
+        raise AttributeError(f"No such attribute '{key}'")
+
+    def __str__(self, indent: int = 2):
+        tmp = {}
+        for k, v in self.items():
+            # PosixPath is ont JSON serializable
+            if isinstance(v, (Path, torch.device, torch.dtype)):
+                v = str(v)
+            tmp[k] = v
+        return json.dumps(tmp, indent=indent, sort_keys=True)
+
+
+class MetricsTracker(collections.defaultdict):
+    def __init__(self):
+        # Passing the type 'int' to the base-class constructor
+        # makes undefined items default to int() which is zero.
+        # This class will play a role as metrics tracker.
+        # It can record many metrics, including but not limited to loss.
+        super(MetricsTracker, self).__init__(int)
+
+    def __add__(self, other: "MetricsTracker") -> "MetricsTracker":
+        ans = MetricsTracker()
+        for k, v in self.items():
+            ans[k] = v
+        for k, v in other.items():
+            if v - v == 0:
+                ans[k] = ans[k] + v
+        return ans
+
+    def __mul__(self, alpha: float) -> "MetricsTracker":
+        ans = MetricsTracker()
+        for k, v in self.items():
+            ans[k] = v * alpha
+        return ans
+
+    def __str__(self) -> str:
+        ans_frames = ""
+        ans_utterances = ""
+        for k, v in self.norm_items():
+            norm_value = "%.4g" % v
+            if "utt_" not in k:
+                ans_frames += str(k) + "=" + str(norm_value) + ", "
+            else:
+                ans_utterances += str(k) + "=" + str(norm_value)
+                if k == "utt_duration":
+                    ans_utterances += " frames, "
+                elif k == "utt_pad_proportion":
+                    ans_utterances += ", "
+                else:
+                    raise ValueError(f"Unexpected key: {k}")
+        frames = "%.2f" % self["frames"]
+        ans_frames += "over " + str(frames) + " frames. "
+        if ans_utterances != "":
+            utterances = "%.2f" % self["utterances"]
+            ans_utterances += "over " + str(utterances) + " utterances."
+
+        return ans_frames + ans_utterances
+
+    def norm_items(self) -> List[Tuple[str, float]]:
+        """
+        Returns a list of pairs, like:
+          [('ctc_loss', 0.1), ('att_loss', 0.07)]
+        """
+        num_frames = self["frames"] if "frames" in self else 1
+        num_utterances = self["utterances"] if "utterances" in self else 1
+        ans = []
+        for k, v in self.items():
+            if k == "frames" or k == "utterances":
+                continue
+            norm_value = (
+                float(v) / num_frames if "utt_" not in k else float(v) / num_utterances
+            )
+            ans.append((k, norm_value))
+        return ans
+
+    def reduce(self, device):
+        """
+        Reduce using torch.distributed, which I believe ensures that
+        all processes get the total.
+        """
+        keys = sorted(self.keys())
+        s = torch.tensor([float(self[k]) for k in keys], device=device)
+        dist.all_reduce(s, op=dist.ReduceOp.SUM)
+        for k, v in zip(keys, s.cpu().tolist()):
+            self[k] = v
+
+    def write_summary(
+        self,
+        tb_writer: SummaryWriter,
+        prefix: str,
+        batch_idx: int,
+    ) -> None:
+        """Add logging information to a TensorBoard writer.
+
+        Args:
+            tb_writer: a TensorBoard writer
+            prefix: a prefix for the name of the loss, e.g. "train/valid_",
+                or "train/current_"
+            batch_idx: The current batch index, used as the x-axis of the plot.
+        """
+        for k, v in self.norm_items():
+            tb_writer.add_scalar(prefix + k, v, batch_idx)
+
+
+@contextmanager
+def torch_autocast(device_type="cuda", **kwargs):
+    """
+    To fix the following warnings:
+    FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated.
+    Please use `torch.amp.autocast('cuda', args...)` instead.
+      with torch.cuda.amp.autocast(enabled=False):
+    """
+    if version.parse(torch.__version__) >= version.parse("2.3.0"):
+        # Use new unified API
+        with torch.amp.autocast(device_type=device_type, **kwargs):
+            yield
+    else:
+        # Suppress deprecation warning and use old CUDA-specific autocast
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=FutureWarning)
+            with torch.cuda.amp.autocast(**kwargs):
+                yield
+
+
+def create_grad_scaler(device="cuda", **kwargs):
+    """
+    Creates a GradScaler compatible with both torch < 2.3.0 and >= 2.3.0.
+    Accepts all kwargs like: enabled, init_scale, growth_factor, etc.
+
+    FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated.
+    Please use `torch.amp.GradScaler('cuda', args...)` instead.
+    """
+    if version.parse(torch.__version__) >= version.parse("2.3.0"):
+        from torch.amp import GradScaler
+
+        return GradScaler(device=device, **kwargs)
+    else:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=FutureWarning)
+            return torch.cuda.amp.GradScaler(**kwargs)
+
+
+def setup_dist(
+    rank=None,
+    world_size=None,
+    master_port=None,
+    use_ddp_launch=False,
+    master_addr=None,
+):
+    """
+    rank and world_size are used only if use_ddp_launch is False.
+    """
+    if "MASTER_ADDR" not in os.environ:
+        os.environ["MASTER_ADDR"] = (
+            "localhost" if master_addr is None else str(master_addr)
+        )
+
+    if "MASTER_PORT" not in os.environ:
+        os.environ["MASTER_PORT"] = "12354" if master_port is None else str(master_port)
+
+    if use_ddp_launch is False:
+        dist.init_process_group("nccl", rank=rank, world_size=world_size)
+        torch.cuda.set_device(rank)
+    else:
+        dist.init_process_group("nccl")
+
+
+def cleanup_dist():
+    dist.destroy_process_group()
+
+
+def prepare_input(
+    params: AttributeDict,
+    batch: dict,
+    device: torch.device,
+    return_tokens: bool = True,
+    return_feature: bool = True,
+    return_audio: bool = False,
+):
+    """
+    Parse the features and targets of the current batch.
+    Args:
+      params:
+        It is returned by :func:`get_params`.
+      batch:
+        It is the return value from iterating
+        `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+        for the format of the `batch`.
+      device:
+        The device of Tensor.
+    """
+    return_list = []
+
+    if return_tokens:
+        return_list += [batch["tokens"]]
+
+    if return_feature:
+        features = batch["features"].to(device)
+        features_lens = batch["features_lens"].to(device)
+        return_list += [features * params.feat_scale, features_lens]
+
+    if return_audio:
+        return_list += [batch["audio"], batch["audio_lens"]]
+
+    return return_list
+
+
+def prepare_avg_tokens_durations(features_lens, tokens_lens):
+    tokens_durations = []
+    for i in range(len(features_lens)):
+        utt_duration = features_lens[i]
+        avg_token_duration = utt_duration // tokens_lens[i]
+        tokens_durations.append([avg_token_duration] * tokens_lens[i])
+    return tokens_durations
+
+
+def pad_labels(y: List[List[int]], pad_id: int, device: torch.device):
+    """
+    Pad the transcripts to the same length with zeros.
+
+    Args:
+      y: the transcripts, which is a list of a list
+
+    Returns:
+      Return a Tensor of padded transcripts.
+    """
+    y = [token_ids + [pad_id] for token_ids in y]
+    length = max([len(token_ids) for token_ids in y])
+    y = [token_ids + [pad_id] * (length - len(token_ids)) for token_ids in y]
+    return torch.tensor(y, dtype=torch.int64, device=device)
+
+
+def get_tokens_index(durations: List[List[int]], num_frames: int) -> torch.Tensor:
+    """
+    Gets position in the transcript for each frame, i.e. the position
+    in the symbol-sequence to look up.
+
+    Args:
+      durations:
+        Duration of each token in transcripts.
+      num_frames:
+        The maximum frame length of the current batch.
+
+    Returns:
+      Return a Tensor of shape (batch_size, num_frames)
+    """
+    durations = [x + [num_frames - sum(x)] for x in durations]
+    batch_size = len(durations)
+    ans = torch.zeros(batch_size, num_frames, dtype=torch.int64)
+    for b in range(batch_size):
+        this_dur = durations[b]
+        cur_frame = 0
+        for i, d in enumerate(this_dur):
+            ans[b, cur_frame : cur_frame + d] = i
+            cur_frame += d
+        assert cur_frame == num_frames, (cur_frame, num_frames)
+    return ans
+
+
+def to_int_tuple(s: Union[str, int]):
+    if isinstance(s, int):
+        return (s,)
+    return tuple(map(int, s.split(",")))
+
+
+def get_adjusted_batch_count(params: AttributeDict) -> float:
+    # returns the number of batches we would have used so far if we had used the
+    # reference duration.  This is for purposes of set_batch_count().
+    return (
+        params.batch_idx_train
+        * (params.max_duration * params.world_size)
+        / params.ref_duration
+    )
+
+
+def set_batch_count(model: Union[nn.Module, DDP], batch_count: float) -> None:
+    if isinstance(model, DDP):
+        # get underlying nn.Module
+        model = model.module
+    for name, module in model.named_modules():
+        if hasattr(module, "batch_count"):
+            module.batch_count = batch_count
+        if hasattr(module, "name"):
+            module.name = name
+
+
+def condition_time_mask(
+    features_lens: torch.Tensor,
+    mask_percent: Tuple[float, float],
+    max_len: int = 0,
+) -> torch.Tensor:
+    """
+    Apply Time masking.
+    Args:
+        features_lens:
+            input tensor of shape ``(B)``
+        mask_size:
+            the width size for masking.
+        max_len:
+            the maximum length of the mask.
+    Returns:
+        Return a 2-D bool tensor (B, T), where masked positions
+        are filled with `True` and non-masked positions are
+        filled with `False`.
+    """
+    mask_size = (
+        torch.zeros_like(features_lens, dtype=torch.float32).uniform_(*mask_percent)
+        * features_lens
+    ).to(torch.int64)
+    mask_starts = (
+        torch.rand_like(mask_size, dtype=torch.float32) * (features_lens - mask_size)
+    ).to(torch.int64)
+    mask_ends = mask_starts + mask_size
+    max_len = max(max_len, features_lens.max())
+    seq_range = torch.arange(0, max_len, device=features_lens.device)
+    mask = (seq_range[None, :] >= mask_starts[:, None]) & (
+        seq_range[None, :] < mask_ends[:, None]
+    )
+    return mask
+
+
+def condition_time_mask_suffix(
+    features_lens: torch.Tensor,
+    mask_percent: Tuple[float, float],
+    max_len: int = 0,
+) -> torch.Tensor:
+    """
+    Apply Time masking, mask from the end time index.
+    Args:
+        features_lens:
+            input tensor of shape ``(B)``
+        mask_size:
+            the width size for masking.
+        max_len:
+            the maximum length of the mask.
+    Returns:
+        Return a 2-D bool tensor (B, T), where masked positions
+        are filled with `True` and non-masked positions are
+        filled with `False`.
+    """
+    mask_size = (
+        torch.zeros_like(features_lens, dtype=torch.float32).uniform_(*mask_percent)
+        * features_lens
+    ).to(torch.int64)
+    mask_starts = (
+        torch.ones_like(mask_size, dtype=torch.float32) * (features_lens - mask_size)
+    ).to(torch.int64)
+    mask_ends = mask_starts + mask_size
+    max_len = max(max_len, features_lens.max())
+    seq_range = torch.arange(0, max_len, device=features_lens.device)
+    mask = (seq_range[None, :] >= mask_starts[:, None]) & (
+        seq_range[None, :] < mask_ends[:, None]
+    )
+    return mask
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """
+    Args:
+      lengths:
+        A 1-D tensor containing sentence lengths.
+      max_len:
+        The length of masks.
+    Returns:
+      Return a 2-D bool tensor, where masked positions
+      are filled with `True` and non-masked positions are
+      filled with `False`.
+
+    >>> lengths = torch.tensor([1, 3, 2, 5])
+    >>> make_pad_mask(lengths)
+    tensor([[False,  True,  True,  True,  True],
+            [False, False, False,  True,  True],
+            [False, False,  True,  True,  True],
+            [False, False, False, False, False]])
+    """
+    assert lengths.ndim == 1, lengths.ndim
+    max_len = max(max_len, lengths.max())
+    n = lengths.size(0)
+    seq_range = torch.arange(0, max_len, device=lengths.device)
+    expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)
+
+    return expaned_lengths >= lengths.unsqueeze(-1)
+
+
+def str2bool(v):
+    """Used in argparse.ArgumentParser.add_argument to indicate
+    that a type is a bool type and user can enter
+
+        - yes, true, t, y, 1, to represent True
+        - no, false, f, n, 0, to represent False
+
+    See https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse  # noqa
+    """
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+def setup_logger(
+    log_filename: Pathlike,
+    log_level: str = "info",
+    use_console: bool = True,
+) -> None:
+    """Setup log level.
+
+    Args:
+      log_filename:
+        The filename to save the log.
+      log_level:
+        The log level to use, e.g., "debug", "info", "warning", "error",
+        "critical"
+      use_console:
+        True to also print logs to console.
+    """
+    now = datetime.now()
+    date_time = now.strftime("%Y-%m-%d-%H-%M-%S")
+    if dist.is_available() and dist.is_initialized():
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+        formatter = f"%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] ({rank}/{world_size}) %(message)s"  # noqa
+        log_filename = f"{log_filename}-{date_time}-{rank}"
+    else:
+        formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
+        log_filename = f"{log_filename}-{date_time}"
+
+    os.makedirs(os.path.dirname(log_filename), exist_ok=True)
+
+    level = logging.ERROR
+    if log_level == "debug":
+        level = logging.DEBUG
+    elif log_level == "info":
+        level = logging.INFO
+    elif log_level == "warning":
+        level = logging.WARNING
+    elif log_level == "critical":
+        level = logging.CRITICAL
+
+    logging.basicConfig(
+        filename=log_filename,
+        format=formatter,
+        level=level,
+        filemode="w",
+        force=True,
+    )
+    if use_console:
+        console = logging.StreamHandler()
+        console.setLevel(level)
+        console.setFormatter(logging.Formatter(formatter))
+        logging.getLogger("").addHandler(console)
+
+
+def get_git_sha1():
+    try:
+        git_commit = (
+            subprocess.run(
+                ["git", "rev-parse", "--short", "HEAD"],
+                check=True,
+                stdout=subprocess.PIPE,
+            )
+            .stdout.decode()
+            .rstrip("\n")
+            .strip()
+        )
+        dirty_commit = (
+            len(
+                subprocess.run(
+                    ["git", "diff", "--shortstat"],
+                    check=True,
+                    stdout=subprocess.PIPE,
+                )
+                .stdout.decode()
+                .rstrip("\n")
+                .strip()
+            )
+            > 0
+        )
+        git_commit = git_commit + "-dirty" if dirty_commit else git_commit + "-clean"
+    except:  # noqa
+        return None
+
+    return git_commit
+
+
+def get_git_date():
+    try:
+        git_date = (
+            subprocess.run(
+                ["git", "log", "-1", "--format=%ad", "--date=local"],
+                check=True,
+                stdout=subprocess.PIPE,
+            )
+            .stdout.decode()
+            .rstrip("\n")
+            .strip()
+        )
+    except:  # noqa
+        return None
+
+    return git_date
+
+
+def get_git_branch_name():
+    try:
+        git_date = (
+            subprocess.run(
+                ["git", "rev-parse", "--abbrev-ref", "HEAD"],
+                check=True,
+                stdout=subprocess.PIPE,
+            )
+            .stdout.decode()
+            .rstrip("\n")
+            .strip()
+        )
+    except:  # noqa
+        return None
+
+    return git_date
+
+
+def get_env_info() -> Dict[str, Any]:
+    """Get the environment information."""
+    return {
+        "torch-version": str(torch.__version__),
+        "torch-cuda-available": torch.cuda.is_available(),
+        "torch-cuda-version": torch.version.cuda,
+        "python-version": sys.version[:4],
+        "zipvoice-git-branch": get_git_branch_name(),
+        "zipvoice-git-sha1": get_git_sha1(),
+        "zipvoice-git-date": get_git_date(),
+        "zipvoice-path": str(Path(__file__).resolve().parent.parent),
+        "hostname": socket.gethostname(),
+        "IP address": socket.gethostbyname(socket.gethostname()),
+    }
+
+
+def get_parameter_groups_with_lrs(
+    model: nn.Module,
+    lr: float,
+    include_names: bool = False,
+    freeze_modules: List[str] = [],
+    unfreeze_modules: List[str] = [],
+) -> List[dict]:
+    """
+    This is for use with the ScaledAdam optimizers (more recent versions that accept
+    lists of named-parameters; we can, if needed, create a version without the names).
+
+    It provides a way to specify learning-rate scales inside the module, so that if
+    any nn.Module in the hierarchy has a floating-point parameter 'lr_scale', it will
+    scale the LR of any parameters inside that module or its submodules.  Note: you
+    can set module parameters outside the __init__ function, e.g.:
+      >>> a = nn.Linear(10, 10)
+      >>> a.lr_scale = 0.5
+
+    Returns: a list of dicts, of the following form:
+      if include_names == False:
+        [  { 'params': [ tensor1, tensor2, ... ], 'lr': 0.01 },
+           { 'params': [ tensor3, tensor4, ... ], 'lr': 0.005 },
+         ...   ]
+      if include_names == true:
+        [  { 'named_params': [ (name1, tensor1, (name2, tensor2), ... ], 'lr': 0.01 },
+           { 'named_params': [ (name3, tensor3), (name4, tensor4), ... ], 'lr': 0.005 },
+         ...   ]
+
+    """
+    # Use freeze_modules or unfreeze_modules to freeze or unfreeze modules
+    assert not (len(freeze_modules) and len(unfreeze_modules))
+
+    # flat_lr_scale just contains the lr_scale explicitly specified
+    # for each prefix of the name, e.g. 'encoder.layers.3', these need
+    # to be multiplied for all prefix of the name of any given parameter.
+    flat_lr_scale = defaultdict(lambda: 1.0)
+    names = []
+    for name, m in model.named_modules():
+        names.append(name)
+        if hasattr(m, "lr_scale"):
+            flat_lr_scale[name] = m.lr_scale
+
+    # lr_to_parames is a dict from learning rate (floating point) to: if
+    # include_names == true, a list of (name, parameter) for that learning rate;
+    # otherwise a list of parameters for that learning rate.
+    lr_to_params = defaultdict(list)
+
+    for name, parameter in model.named_parameters():
+        if not parameter.requires_grad:
+            logging.info(f"Remove {name} from parameter")
+            continue
+        split_name = name.split(".")
+        # caution: as a special case, if the name is '', split_name will be [ '' ].
+        prefix = split_name[0]
+        if len(freeze_modules) > 0:
+            if prefix == "module":  # DDP
+                module_name = split_name[1]
+                if module_name in freeze_modules:
+                    logging.info(f"Remove {name} from parameters")
+                    continue
+            else:
+                if prefix in freeze_modules:
+                    logging.info(f"Remove {name} from parameters")
+                    continue
+        elif len(unfreeze_modules) > 0:
+            if prefix == "module":  # DDP
+                module_name = split_name[1]
+                if module_name not in unfreeze_modules:
+                    logging.info(f"Remove {name} from parameters")
+                    continue
+            else:
+                if prefix not in unfreeze_modules:
+                    logging.info(f"Remove {name} from parameters")
+                    continue
+        cur_lr = lr * flat_lr_scale[prefix]
+        if prefix != "":
+            cur_lr *= flat_lr_scale[""]
+        for part in split_name[1:]:
+            prefix = ".".join([prefix, part])
+            cur_lr *= flat_lr_scale[prefix]
+        lr_to_params[cur_lr].append((name, parameter) if include_names else parameter)
+
+    if include_names:
+        return [{"named_params": pairs, "lr": lr} for lr, pairs in lr_to_params.items()]
+    else:
+        return [{"params": params, "lr": lr} for lr, params in lr_to_params.items()]

zipvoice/utils/diagnostics.py

@@ -0,0 +1,723 @@
+# Copyright  2022-2024  Xiaomi Corp.       (authors: Daniel Povey
+#                                                    Zengwei Yao
+#                                                    Mingshuang Luo,
+#                                                    Zengrui Jin,)
+#
+# 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 logging
+import random
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor, nn
+
+
+class TensorDiagnosticOptions(object):
+    """Options object for tensor diagnostics:
+
+    Args:
+      max_eig_dim:
+        The maximum dimension for which we print out eigenvalues
+        (limited for speed reasons).
+    """
+
+    def __init__(self, max_eig_dim: int = 512):
+        self.max_eig_dim = max_eig_dim
+
+    def dim_is_summarized(self, size: int):
+        return size > 10 and size != 31
+
+
+def get_tensor_stats(
+    x: Tensor,
+    dim: int,
+    stats_type: str,
+) -> Tuple[Tensor, int]:
+    """
+    Returns the specified transformation of the Tensor (either x or x.abs()
+    or (x > 0), summed over all but the index `dim`.
+
+    Args:
+      x:
+        Tensor, tensor to be analyzed
+      dim:
+        Dimension with 0 <= dim < x.ndim
+      stats_type:
+        The stats_type includes several types:
+        "abs" -> take abs() before summing
+        "positive" -> take (x > 0) before summing
+        "rms" -> square before summing, we'll take sqrt later
+        "value"  -> just sum x itself
+        "max", "min" -> take the maximum or minimum [over all other dims but dim]
+            instead of summing
+        "rms-sort" -> this is a bit different than the others, it's based on computing
+            the rms over the specified dim and returning percentiles of the result
+            (11 of them).
+    Returns:
+      stats: a Tensor of shape (x.shape[dim],).
+      count: an integer saying how many items were counted in each element
+      of stats.
+    """
+
+    if stats_type == "rms-sort":
+        rms = (x**2).mean(dim=dim).sqrt()
+        rms = rms.flatten()
+        rms = rms.sort()[0]
+        rms = rms[(torch.arange(11) * rms.numel() // 10).clamp(max=rms.numel() - 1)]
+        count = 1.0
+        return rms, count
+
+    count = x.numel() // x.shape[dim]
+
+    if stats_type == "eigs":
+        x = x.transpose(dim, -1)
+        x = x.reshape(-1, x.shape[-1])
+        # shape of returned tensor: (s, s),
+        # where s is size of dimension `dim` of original x.
+        return torch.matmul(x.transpose(0, 1), x), count
+    elif stats_type == "abs":
+        x = x.abs()
+    elif stats_type == "rms":
+        x = x**2
+    elif stats_type == "positive":
+        x = (x > 0).to(dtype=torch.float)
+    else:
+        assert stats_type in ["value", "max", "min"]
+
+    sum_dims = [d for d in range(x.ndim) if d != dim]
+    if len(sum_dims) > 0:
+        if stats_type == "max":
+            for dim in reversed(sum_dims):
+                x = torch.max(x, dim=dim)[0]
+        elif stats_type == "min":
+            for dim in reversed(sum_dims):
+                x = torch.min(x, dim=dim)[0]
+        else:
+            x = torch.sum(x, dim=sum_dims)
+    x = x.flatten().clone()
+    return x, count
+
+
+@dataclass
+class TensorAndCount:
+    tensor: Tensor
+    count: int
+
+
+class TensorDiagnostic(object):
+    """This class is not directly used by the user, it is responsible for
+    collecting diagnostics for a module or parameter tensor of a torch.nn.Module.
+
+    Args:
+      opts:
+        Options object.
+      name:
+        The name associated with this diagnostics object, will probably be
+            {module_name}.X where X is "output" or "grad", or {parameter_name}.
+            Y where Y is param_value or param_grad.
+    """
+
+    def __init__(self, opts: TensorDiagnosticOptions, name: str):
+        self.opts = opts
+        self.name = name
+        self.class_name = None  # will assign in accumulate()
+
+        self.stats = None  # we'll later assign a list to self.stats.
+        # It's a list of dicts, indexed by dim (i.e. by the
+        # axis of the tensor).  The dicts, in turn, are
+        # indexed by `stats-type` which are strings in
+        # ["abs", "max", "min", "positive", "value", "rms"].
+
+        # scalar_stats contains some analysis of the activations and gradients,
+        self.scalar_stats = None
+
+        # the keys into self.stats[dim] are strings, whose values can be
+        # "abs", "max", "min" ,"value", "positive", "rms", "value".
+        # The values e.g. self.stats[dim]["rms"] are lists of dataclass TensorAndCount,
+        # containing a tensor and its associated count (which is the sum of the other
+        # dims that we aggregated over, e.g. the number of frames and/or batch elements
+        # and/or channels.
+        # ... we actually accumulate the Tensors / counts any time we have the same-dim
+        # tensor, only adding a new element to the list if there was a different dim.
+        # if the string in the key is "eigs", if we detect a length mismatch we put None
+        # as the value.
+
+    def accumulate(self, x, class_name: Optional[str] = None):
+        """
+        Accumulate tensors.
+        """
+        if class_name is not None:
+            self.class_name = class_name
+        if isinstance(x, Tuple):
+            x = x[0]
+        if not isinstance(x, Tensor):
+            return
+        if x.numel() == 0:  # for empty tensor
+            return
+        x = x.detach().clone()
+        if x.ndim == 0:
+            x = x.unsqueeze(0)
+        ndim = x.ndim
+        if self.stats is None:
+            self.stats = [dict() for _ in range(ndim)]
+
+        for dim in range(ndim):
+            this_dim_stats = self.stats[dim]
+            if ndim > 1:
+                # rms-sort is different from the others, it's based on summing over just
+                # this dim, then sorting and returning the percentiles.
+                stats_types = [
+                    "abs",
+                    "max",
+                    "min",
+                    "positive",
+                    "value",
+                    "rms",
+                    "rms-sort",
+                ]
+                if x.shape[dim] <= self.opts.max_eig_dim:
+                    stats_types.append("eigs")
+            else:
+                stats_types = ["value", "abs", "max", "min"]
+
+            for stats_type in stats_types:
+                stats, count = get_tensor_stats(x, dim, stats_type)
+                if stats_type not in this_dim_stats:
+                    this_dim_stats[stats_type] = []  # list of TensorAndCount
+
+                done = False
+                if this_dim_stats[stats_type] is None:
+                    # we can reach here if we detected for stats_type "eigs" that
+                    # where was more than one different size for this dim.  Then we
+                    # disable accumulating this stats type, as it uses too much memory.
+                    continue
+                for s in this_dim_stats[stats_type]:
+                    if s.tensor.shape == stats.shape:
+                        if stats_type == "max":
+                            s.tensor = torch.maximum(s.tensor, stats)
+
+                        elif stats_type == "min":
+                            s.tensor = torch.minimum(s.tensor, stats)
+                        else:
+                            assert stats_type != "max"
+                            s.tensor += stats
+                        s.count += count
+                        done = True
+                        break
+                if not done:
+                    if this_dim_stats[stats_type] != [] and stats_type == "eigs":
+                        # >1 size encountered on this dim, e.g. it's a batch or time
+                        # dimension, don't accumulat "eigs" stats type, it uses too much
+                        # memory
+                        this_dim_stats[stats_type] = None
+                    else:
+                        this_dim_stats[stats_type].append(TensorAndCount(stats, count))
+
+    def print_diagnostics(self):
+        """Print diagnostics for each dimension of the tensor."""
+        if self.stats is None:
+            print(f"Warning: the stats of {self.name} is None.")
+            return
+        for dim, this_dim_stats in enumerate(self.stats):
+            if "rms" in this_dim_stats and "value" in this_dim_stats:
+                # produce "stddev" stats, which is centered RMS.
+                rms_stats_list = this_dim_stats["rms"]
+                value_stats_list = this_dim_stats["value"]
+                if len(rms_stats_list) == len(value_stats_list):
+                    stddev_stats_list = []
+                    for r, v in zip(rms_stats_list, value_stats_list):
+                        stddev_stats_list.append(
+                            # r.count and v.count should be the same, but we don't check
+                            # this.
+                            TensorAndCount(
+                                r.tensor - v.tensor * v.tensor / (v.count + 1.0e-20),
+                                r.count,
+                            )
+                        )
+                    this_dim_stats["stddev"] = stddev_stats_list
+
+            for stats_type, stats_list in this_dim_stats.items():
+                # stats_type could be "rms", "value", "abs", "eigs", "positive", "min"
+                # or "max". "stats_list" could be a list of TensorAndCount (one list per
+                # distinct tensor shape of the stats), or None
+                if stats_list is None:
+                    assert stats_type == "eigs"
+                    continue
+
+                def get_count(count):
+                    return 1 if stats_type in ["max", "min"] else count
+
+                if len(stats_list) == 1:
+                    stats = stats_list[0].tensor / get_count(stats_list[0].count)
+                else:
+                    # a dimension that has variable size in different nnet
+                    # forwards, e.g. a time dimension in an ASR model.
+                    stats = torch.cat(
+                        [x.tensor / get_count(x.count) for x in stats_list], dim=0
+                    )
+
+                if stats_type == "eigs":
+                    try:
+                        if hasattr(torch, "linalg") and hasattr(torch.linalg, "eigh"):
+                            eigs, _ = torch.linalg.eigh(stats)
+                        else:
+                            eigs, _ = torch.symeig(stats)
+                        stats = eigs.abs().sqrt()
+                    except:  # noqa
+                        print("Error getting eigenvalues, trying another method.")
+                        if hasattr(torch, "linalg") and hasattr(torch.linalg, "eig"):
+                            eigs, _ = torch.linalg.eig(stats)
+                            eigs = eigs.abs()
+                        else:
+                            eigs, _ = torch.eig(stats)
+                            eigs = eigs.norm(dim=1)
+                        stats = eigs.sqrt()
+                        # sqrt so it reflects data magnitude, like stddev- not variance
+
+                if stats_type in ["rms", "stddev"]:
+                    # we stored the square; after aggregation we need to take sqrt.
+                    stats = stats.sqrt()
+
+                # if `summarize` we print percentiles of the stats; else,
+                # we print out individual elements.
+                summarize = (len(stats_list) > 1) or self.opts.dim_is_summarized(
+                    stats.numel()
+                )
+                if summarize:  # usually `summarize` will be true
+                    # print out percentiles.
+                    stats = stats.sort()[0]
+                    num_percentiles = 10
+                    size = stats.numel()
+                    percentiles = []
+                    for i in range(num_percentiles + 1):
+                        index = (i * (size - 1)) // num_percentiles
+                        percentiles.append(stats[index].item())
+                    percentiles = ["%.2g" % x for x in percentiles]
+                    percentiles = " ".join(percentiles)
+                    ans = f"percentiles: [{percentiles}]"
+                else:
+                    ans = stats.tolist()
+                    ans = ["%.2g" % x for x in ans]
+                    ans = "[" + " ".join(ans) + "]"
+                if stats_type in ["value", "rms", "stddev", "eigs"]:
+                    # This norm is useful because it is strictly less than the largest
+                    # sqrt(eigenvalue) of the variance, which we print out, and shows,
+                    # speaking in an approximate way, how much of that largest
+                    # eigenvalue can be attributed to the mean of the distribution.
+                    norm = (stats**2).sum().sqrt().item()
+                    ans += f", norm={norm:.2g}"
+                mean = stats.mean().item()
+                rms = (stats**2).mean().sqrt().item()
+                ans += f", mean={mean:.3g}, rms={rms:.3g}"
+
+                # OK, "ans" contains the actual stats, e.g.
+                # ans = "percentiles: \
+                # [0.43 0.46 0.48 0.49 0.49 0.5 0.51 0.52 0.53 0.54 0.59], \
+                # mean=0.5, rms=0.5"
+
+                sizes = [x.tensor.shape[0] for x in stats_list]
+                size_str = (
+                    f"{sizes[0]}" if len(sizes) == 1 else f"{min(sizes)}..{max(sizes)}"
+                )
+                maybe_class_name = (
+                    f" type={self.class_name}," if self.class_name is not None else ""
+                )
+                print(
+                    f"module={self.name},{maybe_class_name} dim={dim}, size={size_str}, "
+                    f"{stats_type} {ans}"
+                )
+
+
+class ScalarDiagnostic(object):
+    """This class is not directly used by the user, it is responsible for
+    collecting diagnostics for a single module (subclass of torch.nn.Module) that
+    represents some kind of nonlinearity, e.g. ReLU, sigmoid, etc.
+    """
+
+    def __init__(self, opts: TensorDiagnosticOptions, name: str):
+        self.opts = opts
+        self.name = name
+        self.class_name = None  # will assign in accumulate()
+        self.is_forward_pass = True
+
+        self.tick_scale = None
+
+        self.saved_inputs = []
+        self.is_ok = True
+
+        self.counts = None
+        self.sum_grad = None
+        self.sum_gradsq = None
+        self.sum_abs_grad = None
+
+    def accumulate_input(self, x: Tensor, class_name: Optional[str] = None):
+        """
+        Called in forward pass.
+        """
+        if not self.is_forward_pass:
+            # in case we did a forward pass without a backward pass, for some reason.
+            self.saved_inputs = []
+            self.is_forward_pass = True
+
+        if class_name is not None:
+            self.class_name = class_name
+        if not self.is_ok:
+            return
+
+        limit = 10
+        if len(self.saved_inputs) > limit:
+            print(
+                f"ERROR: forward pass called for this module over {limit} times "
+                f"with no backward pass. Will not accumulate scalar stats."
+            )
+            self.is_ok = False
+            return
+        self.saved_inputs.append(x)
+
+    def accumulate_output_grad(self, grad: Tensor):
+        if not self.is_ok:
+            return
+        if self.is_forward_pass:
+            self.is_forward_pass = False
+
+        last_shape = (
+            "n/a" if len(self.saved_inputs) == 0 else self.saved_inputs[-1].shape
+        )
+        if len(self.saved_inputs) == 0 or grad.shape != last_shape:
+            print(
+                f"ERROR: shape mismatch or no forward activation present when backward "
+                f"pass called: grad shape ={tuple(grad.shape)}"
+                f", num-saved-inputs={len(self.saved_inputs)}"
+                f", shape-of-last-saved-input={last_shape}"
+            )
+            self.is_ok = False
+            return
+
+        x = self.saved_inputs.pop()
+        self.process_input_and_grad(x, grad)
+
+    def process_input_and_grad(self, x: Tensor, grad: Tensor):
+        assert x.shape == grad.shape
+        x = x.flatten()
+        grad = grad.flatten()
+
+        num_ticks_per_side = 256
+
+        if self.tick_scale is None:
+            x_abs_sorted = x.abs().sort()[0]
+            # take the 98th percentile as the largest value we count separately.
+            index = int(x.numel() * 0.98)
+            self.tick_scale = float(x_abs_sorted[index] / num_ticks_per_side)
+
+            # integerize from tick * (-num ticks_per_side ..  num_ticks_per_side - 1]
+            self.counts = torch.zeros(
+                2 * num_ticks_per_side, dtype=torch.long, device=x.device
+            )
+            self.sum_grad = torch.zeros(
+                2 * num_ticks_per_side, dtype=torch.double, device=x.device
+            )
+            # sum_gradsq is for getting error bars.
+            self.sum_gradsq = torch.zeros(
+                2 * num_ticks_per_side, dtype=torch.double, device=x.device
+            )
+            self.sum_abs_grad = torch.zeros(
+                2 * num_ticks_per_side, dtype=torch.double, device=x.device
+            )
+
+        # this will round down.
+        x = (x / self.tick_scale).to(torch.long)
+        x = x.clamp_(min=-num_ticks_per_side, max=num_ticks_per_side - 1)
+        x = x + num_ticks_per_side
+
+        self.counts.index_add_(dim=0, index=x, source=torch.ones_like(x))
+        self.sum_grad.index_add_(dim=0, index=x, source=grad.to(torch.double))
+        self.sum_gradsq.index_add_(
+            dim=0, index=x, source=(grad * grad).to(torch.double)
+        )
+        self.sum_abs_grad.index_add_(dim=0, index=x, source=grad.abs().to(torch.double))
+
+    def print_diagnostics(self):
+        """Print diagnostics."""
+        if self.is_ok is False or self.counts is None:
+            print(f"Warning: no stats accumulated for {self.name}, is_ok={self.is_ok}")
+            return
+
+        counts = self.counts.to("cpu")
+        sum_grad = self.sum_grad.to(device="cpu", dtype=torch.float32)
+        sum_gradsq = self.sum_gradsq.to(device="cpu", dtype=torch.float32)
+        sum_abs_grad = self.sum_abs_grad.to(device="cpu", dtype=torch.float32)
+
+        counts_cumsum = counts.cumsum(dim=0)
+        counts_tot = counts_cumsum[-1]
+
+        # subdivide the distribution up into `num_bins` intervals for analysis, for
+        # greater statistical significance.  each bin corresponds to multiple of the
+        # original 'tick' intervals.
+        num_bins = 20
+
+        # integer division
+        counts_per_bin = (counts_tot // num_bins) + 1
+        bin_indexes = counts_cumsum // counts_per_bin
+        bin_indexes = bin_indexes.clamp(min=0, max=num_bins).to(torch.long)
+
+        bin_counts = torch.zeros(num_bins, dtype=torch.long)
+        bin_counts.index_add_(dim=0, index=bin_indexes, source=counts)
+        bin_grad = torch.zeros(num_bins)
+        bin_grad.index_add_(dim=0, index=bin_indexes, source=sum_grad)
+        bin_gradsq = torch.zeros(num_bins)
+        bin_gradsq.index_add_(dim=0, index=bin_indexes, source=sum_gradsq)
+        bin_abs_grad = torch.zeros(num_bins)
+        bin_abs_grad.index_add_(dim=0, index=bin_indexes, source=sum_abs_grad)
+
+        bin_boundary_counts = (
+            torch.arange(num_bins + 1, dtype=torch.long) * counts_per_bin
+        )
+        bin_tick_indexes = torch.searchsorted(counts_cumsum, bin_boundary_counts)
+        # boundaries are the "x" values between the bins, e.g. corresponding to the
+        # locations of percentiles of the distribution.
+        num_ticks_per_side = counts.numel() // 2
+        bin_boundaries = (bin_tick_indexes - num_ticks_per_side) * self.tick_scale
+
+        bin_grad = bin_grad / (bin_counts + 1)
+        bin_conf_interval = bin_gradsq.sqrt() / (
+            bin_counts + 1
+        )  # consider this a standard deviation.
+        # bin_grad / bin_abs_grad will give us a sense for how important in a practical
+        # sense, the gradients are.
+        bin_abs_grad = bin_abs_grad / (bin_counts + 1)
+
+        bin_rel_grad = bin_grad / (bin_abs_grad + 1.0e-20)
+        bin_conf = bin_grad / (bin_conf_interval + 1.0e-20)
+
+        def tensor_to_str(x: Tensor):
+            x = ["%.2g" % f for f in x]
+            x = "[" + " ".join(x) + "]"
+            return x
+
+        maybe_class_name = (
+            f" type={self.class_name}," if self.class_name is not None else ""
+        )
+
+        print(
+            f"module={self.name},{maybe_class_name} "
+            f"bin-boundaries={tensor_to_str(bin_boundaries)}, "
+            f"rel_grad={tensor_to_str(bin_rel_grad)}, "
+            f"grad_conf={tensor_to_str(bin_conf)}"
+        )
+
+
+class ModelDiagnostic(object):
+    """This class stores diagnostics for all tensors in the torch.nn.Module.
+
+    Args:
+      opts:
+        Options object.
+    """
+
+    def __init__(self, opts: Optional[TensorDiagnosticOptions] = None):
+        # In this dictionary, the keys are tensors names and the values
+        # are corresponding TensorDiagnostic objects.
+        if opts is None:
+            self.opts = TensorDiagnosticOptions()
+        else:
+            self.opts = opts
+        self.diagnostics = dict()
+
+    def __getitem__(self, name: str):
+        T = ScalarDiagnostic if name[-7:] == ".scalar" else TensorDiagnostic
+        if name not in self.diagnostics:
+            self.diagnostics[name] = T(self.opts, name)
+        return self.diagnostics[name]
+
+    def print_diagnostics(self):
+        """Print diagnostics for each tensor."""
+        for k in sorted(self.diagnostics.keys()):
+            self.diagnostics[k].print_diagnostics()
+
+
+def get_class_name(module: nn.Module):
+    ans = type(module).__name__
+    # we put the below in try blocks in case anyone is using a different version of
+    # these modules that might have different member names.
+    if ans == "Balancer" or ans == "ActivationBalancer":
+        try:
+            ans += f"[{float(module.min_positive)},{float(module.max_positive)},"
+            f"{float(module.min_abs)},{float(module.max_abs)}]"
+        except:
+            pass
+    elif ans == "AbsValuePenalizer":
+        try:
+            ans += f"[{module.limit}]"
+        except:
+            pass
+    return ans
+
+
+def attach_diagnostics(
+    model: nn.Module, opts: Optional[TensorDiagnosticOptions] = None
+) -> ModelDiagnostic:
+    """Attach a ModelDiagnostic object to the model by
+    1) registering forward hook and backward hook on each module, to accumulate
+    its output tensors and gradient tensors, respectively;
+    2) registering backward hook on each module parameter, to accumulate its
+    values and gradients.
+
+    Args:
+      model:
+        the model to be analyzed.
+      opts:
+        Options object.
+
+    Returns:
+      The ModelDiagnostic object attached to the model.
+    """
+
+    ans = ModelDiagnostic(opts)
+    for name, module in model.named_modules():
+        if name == "":
+            name = "<top-level>"
+
+        # Setting model_diagnostic=ans and n=name below, instead of trying to
+        # capture the variables, ensures that we use the current values.
+        # (this matters for `name`, since the variable gets overwritten).
+        # These closures don't really capture by value, only by
+        # "the final value the variable got in the function" :-(
+        def forward_hook(_module, _input, _output, _model_diagnostic=ans, _name=name):
+            if isinstance(_output, tuple) and len(_output) == 1:
+                _output = _output[0]
+
+            if isinstance(_output, Tensor) and _output.dtype in (
+                torch.float32,
+                torch.float16,
+                torch.float64,
+            ):
+                _model_diagnostic[f"{_name}.output"].accumulate(
+                    _output, class_name=get_class_name(_module)
+                )
+            elif isinstance(_output, tuple):
+                for i, o in enumerate(_output):
+                    if isinstance(o, Tensor) and o.dtype in (
+                        torch.float32,
+                        torch.float16,
+                        torch.float64,
+                    ):
+                        _model_diagnostic[f"{_name}.output[{i}]"].accumulate(
+                            o, class_name=get_class_name(_module)
+                        )
+
+        def backward_hook(_module, _input, _output, _model_diagnostic=ans, _name=name):
+            if isinstance(_output, tuple) and len(_output) == 1:
+                _output = _output[0]
+            if isinstance(_output, Tensor) and _output.dtype in (
+                torch.float32,
+                torch.float16,
+                torch.float64,
+            ):
+                _model_diagnostic[f"{_name}.grad"].accumulate(
+                    _output, class_name=get_class_name(_module)
+                )
+            elif isinstance(_output, tuple):
+                for i, o in enumerate(_output):
+                    if isinstance(o, Tensor) and o.dtype in (
+                        torch.float32,
+                        torch.float16,
+                        torch.float64,
+                    ):
+                        _model_diagnostic[f"{_name}.grad[{i}]"].accumulate(
+                            o, class_name=get_class_name(_module)
+                        )
+
+        module.register_forward_hook(forward_hook)
+        module.register_backward_hook(backward_hook)
+
+        if type(module).__name__ in [
+            "Sigmoid",
+            "Tanh",
+            "ReLU",
+            "TanSwish",
+            "Swish",
+            "DoubleSwish",
+            "Swoosh",
+        ]:
+            # For these specific module types, accumulate some additional diagnostics
+            # that can help us improve the activation function.  These require a lot of
+            # memory, to save the forward activations, so limit this to some select
+            # classes. Note: this will not work correctly for all model types.
+            def scalar_forward_hook(
+                _module, _input, _output, _model_diagnostic=ans, _name=name
+            ):
+                if isinstance(_input, tuple):
+                    (_input,) = _input
+                assert isinstance(_input, Tensor)
+                _model_diagnostic[f"{_name}.scalar"].accumulate_input(
+                    _input, class_name=get_class_name(_module)
+                )
+
+            def scalar_backward_hook(
+                _module, _input, _output, _model_diagnostic=ans, _name=name
+            ):
+                if isinstance(_output, tuple):
+                    (_output,) = _output
+                assert isinstance(_output, Tensor)
+                _model_diagnostic[f"{_name}.scalar"].accumulate_output_grad(_output)
+
+            module.register_forward_hook(scalar_forward_hook)
+            module.register_backward_hook(scalar_backward_hook)
+
+    for name, parameter in model.named_parameters():
+
+        def param_backward_hook(
+            grad, _parameter=parameter, _model_diagnostic=ans, _name=name
+        ):
+            _model_diagnostic[f"{_name}.param_value"].accumulate(_parameter)
+            _model_diagnostic[f"{_name}.param_grad"].accumulate(grad)
+
+        try:
+            parameter.register_hook(param_backward_hook)
+        except:
+            logging.warning(
+                f"Warning: could not register backward hook for parameter {name}, "
+                f"it might not be differentiable."
+            )
+
+    return ans
+
+
+def _test_tensor_diagnostic():
+    opts = TensorDiagnosticOptions(512)
+
+    diagnostic = TensorDiagnostic(opts, "foo")
+
+    for _ in range(10):
+        diagnostic.accumulate(torch.randn(50, 100) * 10.0)
+
+    diagnostic.print_diagnostics()
+
+    model = nn.Sequential(nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 80))
+
+    diagnostic = attach_diagnostics(model, opts)
+    for _ in range(10):
+        T = random.randint(200, 300)
+        x = torch.randn(T, 100)
+        y = model(x)
+        y.sum().backward()
+
+    diagnostic.print_diagnostics()
+
+
+if __name__ == "__main__":
+    _test_tensor_diagnostic()

zipvoice/utils/feature.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+# Copyright         2024  Xiaomi Corp.        (authors: Han Zhu)
+#
+# 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.
+
+from dataclasses import dataclass
+from typing import Union
+
+import numpy as np
+import torch
+import torchaudio
+from lhotse.features.base import FeatureExtractor, register_extractor
+from lhotse.utils import Seconds, compute_num_frames
+
+
+@dataclass
+class VocosFbankConfig:
+    sampling_rate: int = 24000
+    n_mels: int = 100
+    n_fft: int = 1024
+    hop_length: int = 256
+
+
+@register_extractor
+class VocosFbank(FeatureExtractor):
+
+    name = "VocosFbank"
+    config_type = VocosFbankConfig
+
+    def __init__(self, num_channels: int = 1):
+        config = VocosFbankConfig
+        super().__init__(config=config)
+        assert num_channels in (1, 2)
+        self.num_channels = num_channels
+        self.fbank = torchaudio.transforms.MelSpectrogram(
+            sample_rate=self.config.sampling_rate,
+            n_fft=self.config.n_fft,
+            hop_length=self.config.hop_length,
+            n_mels=self.config.n_mels,
+            center=True,
+            power=1,
+        )
+
+    def _feature_fn(self, sample):
+        mel = self.fbank(sample)
+        logmel = mel.clamp(min=1e-7).log()
+
+        return logmel
+
+    @property
+    def device(self) -> Union[str, torch.device]:
+        return self.config.device
+
+    def feature_dim(self, sampling_rate: int) -> int:
+        return self.config.n_mels
+
+    def extract(
+        self,
+        samples: Union[np.ndarray, torch.Tensor],
+        sampling_rate: int,
+    ) -> Union[np.ndarray, torch.Tensor]:
+        # Check for sampling rate compatibility.
+        expected_sr = self.config.sampling_rate
+        assert sampling_rate == expected_sr, (
+            f"Mismatched sampling rate: extractor expects {expected_sr}, "
+            f"got {sampling_rate}"
+        )
+        is_numpy = False
+        if not isinstance(samples, torch.Tensor):
+            samples = torch.from_numpy(samples)
+            is_numpy = True
+
+        if len(samples.shape) == 1:
+            samples = samples.unsqueeze(0)
+        else:
+            assert samples.ndim == 2, samples.shape
+
+        if self.num_channels == 1:
+            if samples.shape[0] == 2:
+                samples = samples.mean(dim=0, keepdims=True)
+        else:
+            assert samples.shape[0] == 2, samples.shape
+
+        mel = self._feature_fn(samples)
+        # (1, n_mels, time) or (2, n_mels, time)
+        mel = mel.reshape(-1, mel.shape[-1]).t()
+        # (time, n_mels) or (time, 2 * n_mels)
+
+        num_frames = compute_num_frames(
+            samples.shape[1] / sampling_rate, self.frame_shift, sampling_rate
+        )
+
+        if mel.shape[0] > num_frames:
+            mel = mel[:num_frames]
+        elif mel.shape[0] < num_frames:
+            mel = mel.unsqueeze(0)
+            mel = torch.nn.functional.pad(
+                mel, (0, 0, 0, num_frames - mel.shape[1]), mode="replicate"
+            ).squeeze(0)
+
+        if is_numpy:
+            return mel.cpu().numpy()
+        else:
+            return mel
+
+    @property
+    def frame_shift(self) -> Seconds:
+        return self.config.hop_length / self.config.sampling_rate

zipvoice/utils/hooks.py

@@ -0,0 +1,111 @@
+# Copyright  2021-2024  Xiaomi Corporation  (authors: Zengwei Yao,
+#                                                     Daniel Povey,
+#                                                     Zengrui Jin,)
+#
+# 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 logging
+import random
+
+import torch
+from torch import Tensor, nn
+
+
+def register_inf_check_hooks(model: nn.Module) -> None:
+    """Registering forward hook on each module, to check
+    whether its output tensors is not finite.
+
+    Args:
+      model:
+        the model to be analyzed.
+    """
+
+    for name, module in model.named_modules():
+        if name == "":
+            name = "<top-level>"
+
+        # default param _name is a way to capture the current value of the variable
+        # "name".
+        def forward_hook(_module, _input, _output, _name=name):
+            if isinstance(_output, Tensor):
+                try:
+                    if not torch.isfinite(_output.to(torch.float32).sum()):
+                        logging.warning(f"The sum of {_name}.output is not finite")
+                except RuntimeError:  # e.g. CUDA out of memory
+                    pass
+            elif isinstance(_output, tuple):
+                for i, o in enumerate(_output):
+                    if isinstance(o, tuple):
+                        o = o[0]
+                    if not isinstance(o, Tensor):
+                        continue
+                    try:
+                        if not torch.isfinite(o.to(torch.float32).sum()):
+                            logging.warning(
+                                f"The sum of {_name}.output[{i}] is not finite"
+                            )
+                    except RuntimeError:  # e.g. CUDA out of memory
+                        pass
+
+        # default param _name is a way to capture the current value of the variable
+        # "name".
+        def backward_hook(_module, _input, _output, _name=name):
+            if isinstance(_output, Tensor):
+                try:
+                    if not torch.isfinite(_output.to(torch.float32).sum()):
+                        logging.warning(f"The sum of {_name}.grad is not finite")
+                except RuntimeError:  # e.g. CUDA out of memory
+                    pass
+
+            elif isinstance(_output, tuple):
+                for i, o in enumerate(_output):
+                    if isinstance(o, tuple):
+                        o = o[0]
+                    if not isinstance(o, Tensor):
+                        continue
+                    if not torch.isfinite(o.to(torch.float32).sum()):
+                        logging.warning(f"The sum of {_name}.grad[{i}] is not finite")
+
+        module.register_forward_hook(forward_hook)
+        module.register_backward_hook(backward_hook)
+
+    for name, parameter in model.named_parameters():
+
+        def param_backward_hook(grad, _name=name):
+            if not torch.isfinite(grad.to(torch.float32).sum()):
+                logging.warning(f"The sum of {_name}.param_grad is not finite")
+
+        try:
+            parameter.register_hook(param_backward_hook)
+        except Exception as e:
+            logging.warning(
+                f"Warning: could not register backward hook for parameter {name}"
+                f" with error {e}, it might not be differentiable."
+            )
+
+
+def _test_inf_check_hooks():
+    model = nn.Sequential(nn.Linear(100, 50), nn.Linear(50, 80))
+
+    register_inf_check_hooks(model)
+    for _ in range(10):
+        T = random.randint(200, 300)
+        x = torch.randn(T, 100) + float("inf") * (T % 2)
+        y = model(x)
+        y.sum().backward()
+
+
+if __name__ == "__main__":
+    _test_inf_check_hooks()

zipvoice/utils/infer.py

@@ -0,0 +1,414 @@
+from typing import List
+
+import numpy as np
+import torch
+import torchaudio
+from pydub import AudioSegment
+from pydub.silence import detect_leading_silence, split_on_silence
+
+punctuation = {";", ":", ",", ".", "!", "?", "；", "：", "，", "。", "！", "？"}
+
+
+def chunk_tokens_punctuation(tokens_list: List[str], max_tokens: int = 100):
+    """
+    Splits the input tokens list into chunks according to punctuations,
+        each with a maximum number of tokens.
+
+    Args:
+        token_list (list of str): The list of tokens to be split.
+        max_tokens (int): The maximum number of tokens per chunk.
+
+    Returns:
+        List[str]: A list of text chunks.
+    """
+
+    # 1. Split the tokens according to punctuations.
+    sentences = []
+    current_sentence = []
+    for token in tokens_list:
+        # If the first token of current sentence is punctuation or blank,
+        # append it to the end of the previous sentence.
+        if (
+            len(current_sentence) == 0
+            and len(sentences) != 0
+            and (token in punctuation or token == " ")
+        ):
+            sentences[-1].append(token)
+        # Otherwise, append the current token to the current sentence.
+        else:
+            current_sentence.append(token)
+            # Split the sentence in positions of punctuations.
+            if token in punctuation:
+                sentences.append(current_sentence)
+                current_sentence = []
+    # Assume the last few tokens are also a sentence
+    if len(current_sentence) != 0:
+        sentences.append(current_sentence)
+
+    # 2. Merge short sentences.
+    chunks = []
+    current_chunk = []
+    for sentence in sentences:
+        if len(current_chunk) + len(sentence) <= max_tokens:
+            current_chunk.extend(sentence)
+        else:
+            if len(current_chunk) > 0:
+                chunks.append(current_chunk)
+            current_chunk = sentence
+
+    if len(current_chunk) > 0:
+        chunks.append(current_chunk)
+
+    return chunks
+
+
+def chunk_tokens_dialog(tokens_list: List[str], max_tokens: int = 100):
+    """
+    Splits the input tokens list into chunks according to speaker-turn
+        symbol [S1], each with a maximum number of tokens.
+
+    Args:
+        token_list (list of str): The list of tokens to be split.
+        max_tokens (int): The maximum number of tokens per chunk.
+
+    Returns:
+        List[str]: A list of text chunks.
+    """
+
+    # 1. Split the tokens according to speaker-turn symbol [S1].
+    dialogs = []
+    current_dialog = []
+    for token in tokens_list:
+        if token == "[S1]":
+            if len(current_dialog) != 0:
+                dialogs.append(current_dialog)
+            current_dialog = []
+        current_dialog.append(token)
+    # Assume the last few tokens are also a dialog
+    if len(current_dialog) != 0:
+        dialogs.append(current_dialog)
+
+    # 2. Merge short dialogs.
+    chunks = []
+    current_chunk = []
+    for dialog in dialogs:
+        if len(current_chunk) + len(dialog) <= max_tokens:
+            current_chunk.extend(dialog)
+        else:
+            if len(current_chunk) > 0:
+                chunks.append(current_chunk)
+            current_chunk = dialog
+
+    if len(current_chunk) > 0:
+        chunks.append(current_chunk)
+
+    return chunks
+
+
+def batchify_tokens(
+    tokens_list: List[List[int]],
+    max_duration: float,
+    prompt_duration: float,
+    token_duration: float,
+):
+    """
+    Sort and group the input list of token sequences into batches, where each batch's
+        total duration does not exceed the maximum.
+
+    Args:
+        tokens_list (List[List[int]]): A list of token sequences, where each inner
+            list represents a sequence of tokens.
+        max_duration (float): The maximum allowed total duration for each batch.
+        prompt_duration (float): The duration cost per prompt in the batch.
+        token_duration (float): The duration cost per token.
+
+    Returns:
+        batches: List[List[List[int]]]: A list of batches, where each batch is a list of
+            token sequences that fit within the max duration.
+        index: List[int]: The original index of each sentence, used to recover the
+            sequential order in the future.
+    """
+    # Create index for each sentence
+    indexed_tokens = list(enumerate(tokens_list))
+
+    # Sort according to sentence length (for less padding)
+    indexed_sorted_tokens = sorted(indexed_tokens, key=lambda x: len(x[1]))
+    index = [indexed_sorted_tokens[i][0] for i in range(len(indexed_sorted_tokens))]
+    sorted_tokens = [
+        indexed_sorted_tokens[i][1] for i in range(len(indexed_sorted_tokens))
+    ]
+
+    batches = []
+    batch = []
+    batch_size = 0  # Total number of tokens in current batch
+
+    for tokens in sorted_tokens:
+        # Calculate if adding current token sequence would exceed max duration
+        # Formula considers: existing tokens' duration + existing
+        # prompts' duration + new tokens' duration
+        if (
+            batch_size * token_duration
+            + len(batch) * prompt_duration
+            + len(tokens) * token_duration
+            <= max_duration
+        ):
+            # Add to current batch if within duration limit
+            batch.append(tokens)
+            batch_size += len(tokens)
+        else:
+            # If exceeding limit, finalize current batch (if not empty)
+            if len(batch) > 0:
+                batches.append(batch)
+            # Start new batch with current token sequence
+            batch = [tokens]
+            batch_size = len(tokens)
+
+    # Add the last batch if it's not empty
+    if len(batch) > 0:
+        batches.append(batch)
+
+    return batches, index
+
+
+def cross_fade_concat(
+    chunks: List[torch.Tensor], fade_duration: float = 0.1, sample_rate: int = 24000
+) -> torch.Tensor:
+    """
+    Concatenates audio chunks with cross-fading between consecutive chunks.
+
+    Args:
+        chunks: List of audio tensors, each with shape (C, T) where
+                C = number of channel, T = time dimension (samples)
+        fade_duration: Duration of cross-fade in seconds
+        sample_rate: Audio sample rate in Hz
+
+    Returns:
+        Concatenated audio tensor with shape (N, T_total)
+    """
+    # Handle edge cases: empty input or single chunk
+    if len(chunks) <= 1:
+        return chunks[0] if chunks else torch.tensor([])
+
+    # Calculate total fade samples from duration and sample rate
+    fade_samples = int(fade_duration * sample_rate)
+
+    # Use simple concatenation if fade duration is non-positive
+    if fade_samples <= 0:
+        return torch.cat(chunks, dim=-1)
+
+    # Initialize final tensor with the first chunk
+    final = chunks[0]
+
+    # Iterate through remaining chunks to apply cross-fading
+    for next_chunk in chunks[1:]:
+        # Calculate safe fade length (cannot exceed either chunk's duration)
+        k = min(fade_samples, final.shape[-1], next_chunk.shape[-1])
+
+        # Fall back to simple concatenation if safe fade length is invalid
+        if k <= 0:
+            final = torch.cat([final, next_chunk], dim=-1)
+            continue
+
+        # Create fade curve (1 -> 0) with shape (1, k) for broadcasting
+        fade = torch.linspace(1, 0, k, device=final.device)[None]
+
+        # Concatenate three parts:
+        # 1. Non-overlapping part of previous audio
+        # 2. Cross-faded overlapping region
+        # 3. Non-overlapping part of next audio
+        final = torch.cat(
+            [
+                final[..., :-k],  # All samples except last k from previous
+                final[..., -k:] * fade
+                + next_chunk[..., :k] * (1 - fade),  # Cross-fade region
+                next_chunk[..., k:],  # All samples except first k from next
+            ],
+            dim=-1,
+        )
+
+    return final
+
+
+def add_punctuation(text: str):
+    """Add punctuation if there is not in the end of text"""
+    text = text.strip()
+    if text[-1] not in punctuation:
+        text += "."
+    return text
+
+
+def load_prompt_wav(prompt_wav: str, sampling_rate: int):
+    """
+    Load the waveform with torchaudio and resampling if needed.
+
+    Parameters:
+        prompt_wav: path of the prompt wav.
+        sampling_rate: target sampling rate.
+
+    Returns:
+        Loaded prompt waveform with target sampling rate,
+        PyTorch tensor of shape (C, T)
+    """
+    prompt_wav, prompt_sampling_rate = torchaudio.load(prompt_wav)
+
+    if prompt_sampling_rate != sampling_rate:
+        resampler = torchaudio.transforms.Resample(
+            orig_freq=prompt_sampling_rate, new_freq=sampling_rate
+        )
+        prompt_wav = resampler(prompt_wav)
+    return prompt_wav
+
+
+def rms_norm(prompt_wav: torch.Tensor, target_rms: float):
+    """
+    Normalize the rms of prompt_wav is it is smaller than target rms.
+
+    Parameters:
+        prompt_wav: PyTorch tensor with shape (C, T).
+        target_rms: target rms value
+
+    Returns:
+        prompt_wav: normalized prompt wav with shape (C, T).
+        promt_rms: rms of original prompt wav. Will be used to
+            re-normalize the generated wav.
+    """
+    prompt_rms = torch.sqrt(torch.mean(torch.square(prompt_wav)))
+    if prompt_rms < target_rms:
+        prompt_wav = prompt_wav * target_rms / prompt_rms
+    return prompt_wav, prompt_rms
+
+
+def remove_silence(
+    audio: torch.Tensor,
+    sampling_rate: int,
+    only_edge: bool = False,
+    trail_sil: float = 0,
+):
+    """
+    Remove silences longer than 1 second, and edge silences longer than 0.1 seconds
+
+    Parameters:
+        audio: PyTorch tensor with shape (C, T).
+        sampling_rate: sampling rate of the audio.
+        only_edge: If true, only remove edge silences.
+        trail_sil: the duration of added trailing silence in ms.
+
+    Returns:
+        PyTorch tensor with shape (C, T), where C is number of channels
+            and T is number of audio samples
+    """
+    # Load audio file
+    wave = tensor_to_audiosegment(audio, sampling_rate)
+
+    if not only_edge:
+        # Split audio using silences longer than 1 second
+        non_silent_segs = split_on_silence(
+            wave,
+            min_silence_len=1000,  # Silences longer than 1 second (1000ms)
+            silence_thresh=-50,
+            keep_silence=1000,  # Keep 1.0 second of silence around segments
+            seek_step=10,
+        )
+
+        # Concatenate all non-silent segments
+        wave = AudioSegment.silent(duration=0)
+        for seg in non_silent_segs:
+            wave += seg
+
+    # Remove silence longer than 0.1 seconds in the begining and ending of wave
+    wave = remove_silence_edges(wave, 100, -50)
+
+    # Add trailing silence to avoid leaking prompt to generated speech.
+    wave = wave + AudioSegment.silent(duration=trail_sil)
+
+    # Convert to PyTorch tensor
+    return audiosegment_to_tensor(wave)
+
+
+def remove_silence_edges(
+    audio: AudioSegment, keep_silence: int = 100, silence_threshold: float = -50
+):
+    """
+    Remove edge silences longer than `keep_silence` ms.
+
+    Parameters:
+        audio: an AudioSegment object.
+        keep_silence: kept silence in the edge.
+        only_edge: If true, only remove edge silences.
+        silence_threshold: the threshold of silence.
+
+    Returns:
+        An AudioSegment object
+    """
+    # Remove leading silence
+    start_idx = detect_leading_silence(audio, silence_threshold=silence_threshold)
+    start_idx = max(0, start_idx - keep_silence)
+    audio = audio[start_idx:]
+
+    # Remove trailing silence
+    audio = audio.reverse()
+    start_idx = detect_leading_silence(audio, silence_threshold=silence_threshold)
+    start_idx = max(0, start_idx - keep_silence)
+    audio = audio[start_idx:]
+    audio = audio.reverse()
+
+    return audio
+
+
+def audiosegment_to_tensor(aseg):
+    """
+    Convert a pydub.AudioSegment to PyTorch audio tensor
+    """
+    audio_data = np.array(aseg.get_array_of_samples())
+
+    # Convert to float32 and normalize to [-1, 1] range
+    audio_data = audio_data.astype(np.float32) / 32768.0
+
+    # Handle channels
+    if aseg.channels == 1:
+        # Mono channel: add channel dimension (T) -> (1, T)
+        tensor_data = torch.from_numpy(audio_data).unsqueeze(0)
+    else:
+        # Multi-channel: reshape to (C, T)
+        tensor_data = torch.from_numpy(audio_data.reshape(-1, aseg.channels).T)
+
+    return tensor_data
+
+
+def tensor_to_audiosegment(tensor, sample_rate):
+    """
+    Convert a PyTorch audio tensor to pydub.AudioSegment
+
+    Parameters:
+        tensor: Tensor with shape (C, T), where C is the number of channels
+            and T is the time steps
+        sample_rate: Audio sample rate
+    """
+    # Convert tensor to numpy array
+    audio_np = tensor.cpu().numpy()
+
+    # Add channel dimension if single channel
+    if audio_np.ndim == 1:
+        audio_np = audio_np[np.newaxis, :]
+
+    # Convert to int16 type (common format for pydub)
+    # Assumes tensor values are in [-1, 1] range as floating point
+    audio_np = (audio_np * 32768.0).clip(-32768, 32767).astype(np.int16)
+
+    # Convert to byte stream
+    # For multi-channel audio, pydub requires interleaved format
+    # (e.g., left-right-left-right)
+    if audio_np.shape[0] > 1:
+        # Convert to interleaved format
+        audio_np = audio_np.transpose(1, 0).flatten()
+    audio_bytes = audio_np.tobytes()
+
+    # Create AudioSegment
+    audio_segment = AudioSegment(
+        data=audio_bytes,
+        sample_width=2,
+        frame_rate=sample_rate,
+        channels=tensor.shape[0],
+    )
+
+    return audio_segment

zipvoice/utils/lr_scheduler.py

@@ -0,0 +1,245 @@
+# Copyright      2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# 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 logging
+from typing import List, Optional, Union
+
+import torch
+from torch.optim import Optimizer
+
+
+class LRScheduler(object):
+    """
+    Base-class for learning rate schedulers where the learning-rate depends on both the
+    batch and the epoch.
+    """
+
+    def __init__(self, optimizer: Optimizer, verbose: bool = False):
+        # Attach optimizer
+        if not isinstance(optimizer, Optimizer):
+            raise TypeError("{} is not an Optimizer".format(type(optimizer).__name__))
+        self.optimizer = optimizer
+        self.verbose = verbose
+
+        for group in optimizer.param_groups:
+            group.setdefault("base_lr", group["lr"])
+
+        self.base_lrs = [group["base_lr"] for group in optimizer.param_groups]
+
+        self.epoch = 0
+        self.batch = 0
+
+    def state_dict(self):
+        """Returns the state of the scheduler as a :class:`dict`.
+
+        It contains an entry for every variable in self.__dict__ which
+        is not the optimizer.
+        """
+        return {
+            # the user might try to override the base_lr, so don't include this in the
+            # state. previously they were included.
+            # "base_lrs": self.base_lrs,
+            "epoch": self.epoch,
+            "batch": self.batch,
+        }
+
+    def load_state_dict(self, state_dict):
+        """Loads the schedulers state.
+
+        Args:
+            state_dict (dict): scheduler state. Should be an object returned
+                from a call to :meth:`state_dict`.
+        """
+        # the things with base_lrs are a work-around for a previous problem
+        # where base_lrs were written with the state dict.
+        base_lrs = self.base_lrs
+        self.__dict__.update(state_dict)
+        self.base_lrs = base_lrs
+
+    def get_last_lr(self) -> List[float]:
+        """Return last computed learning rate by current scheduler.
+        Will be a list of float."""
+        return self._last_lr
+
+    def get_lr(self):
+        # Compute list of learning rates from self.epoch and self.batch and
+        # self.base_lrs; this must be overloaded by the user.
+        # e.g. return [some_formula(self.batch, self.epoch, base_lr)
+        # for base_lr in self.base_lrs ]
+        raise NotImplementedError
+
+    def step_batch(self, batch: Optional[int] = None) -> None:
+        # Step the batch index, or just set it.  If `batch` is specified, it
+        # must be the batch index from the start of training, i.e. summed over
+        # all epochs.
+        # You can call this in any order; if you don't provide 'batch', it should
+        # of course be called once per batch.
+        if batch is not None:
+            self.batch = batch
+        else:
+            self.batch = self.batch + 1
+        self._set_lrs()
+
+    def step_epoch(self, epoch: Optional[int] = None):
+        # Step the epoch index, or just set it.  If you provide the 'epoch' arg, you
+        # should call this at the start of the epoch; if you don't provide the 'epoch'
+        # arg, you should call it at the end of the epoch.
+        if epoch is not None:
+            self.epoch = epoch
+        else:
+            self.epoch = self.epoch + 1
+        self._set_lrs()
+
+    def _set_lrs(self):
+        values = self.get_lr()
+        assert len(values) == len(self.optimizer.param_groups)
+
+        for i, data in enumerate(zip(self.optimizer.param_groups, values)):
+            param_group, lr = data
+            param_group["lr"] = lr
+            self.print_lr(self.verbose, i, lr)
+        self._last_lr = [group["lr"] for group in self.optimizer.param_groups]
+
+    def print_lr(self, is_verbose, group, lr):
+        """Display the current learning rate."""
+        if is_verbose:
+            logging.warning(
+                f"Epoch={self.epoch}, batch={self.batch}: adjusting learning rate"
+                f" of group {group} to {lr:.4e}."
+            )
+
+
+class Eden(LRScheduler):
+    """
+    Eden scheduler.
+    The basic formula (before warmup) is:
+      lr = base_lr * (((batch**2 + lr_batches**2) / lr_batches**2) ** -0.25 *
+                     (((epoch**2 + lr_epochs**2) / lr_epochs**2) ** -0.25)) * warmup
+    where `warmup` increases from linearly 0.5 to 1 over `warmup_batches` batches
+    and then stays constant at 1.
+
+    If you don't have the concept of epochs, or one epoch takes a very long time,
+    you can replace the notion of 'epoch' with some measure of the amount of data
+    processed, e.g. hours of data or frames of data, with 'lr_epochs' being set to
+    some measure representing "quite a lot of data": say, one fifth or one third
+    of an entire training run, but it doesn't matter much.  You could also use
+    Eden2 which has only the notion of batches.
+
+    We suggest base_lr = 0.04 (passed to optimizer) if used with ScaledAdam
+
+    Args:
+        optimizer: the optimizer to change the learning rates on
+        lr_batches: the number of batches after which we start significantly
+              decreasing the learning rate, suggest 5000.
+        lr_epochs: the number of epochs after which we start significantly
+              decreasing the learning rate, suggest 6 if you plan to do e.g.
+              20 to 40 epochs, but may need smaller number if dataset is huge
+              and you will do few epochs.
+    """
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        lr_batches: Union[int, float],
+        lr_epochs: Union[int, float],
+        warmup_batches: Union[int, float] = 500.0,
+        warmup_start: float = 0.5,
+        verbose: bool = False,
+    ):
+        super(Eden, self).__init__(optimizer, verbose)
+        self.lr_batches = lr_batches
+        self.lr_epochs = lr_epochs
+        self.warmup_batches = warmup_batches
+
+        assert 0.0 <= warmup_start <= 1.0, warmup_start
+        self.warmup_start = warmup_start
+
+    def get_lr(self):
+        factor = (
+            (self.batch**2 + self.lr_batches**2) / self.lr_batches**2
+        ) ** -0.25 * (
+            ((self.epoch**2 + self.lr_epochs**2) / self.lr_epochs**2) ** -0.25
+        )
+        warmup_factor = (
+            1.0
+            if self.batch >= self.warmup_batches
+            else self.warmup_start
+            + (1.0 - self.warmup_start) * (self.batch / self.warmup_batches)
+            # else 0.5 + 0.5 * (self.batch / self.warmup_batches)
+        )
+
+        return [x * factor * warmup_factor for x in self.base_lrs]
+
+
+class FixedLRScheduler(LRScheduler):
+    """
+    Fixed learning rate scheduler.
+
+    Args:
+        optimizer: the optimizer to change the learning rates on
+    """
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        verbose: bool = False,
+    ):
+        super(FixedLRScheduler, self).__init__(optimizer, verbose)
+
+    def get_lr(self):
+
+        return [x for x in self.base_lrs]
+
+
+def _test_eden():
+    m = torch.nn.Linear(100, 100)
+    from zipvoice.utils.optim import ScaledAdam
+
+    optim = ScaledAdam(m.parameters(), lr=0.03)
+
+    scheduler = Eden(optim, lr_batches=100, lr_epochs=2, verbose=True)
+
+    for epoch in range(10):
+        scheduler.step_epoch(epoch)  # sets epoch to `epoch`
+
+        for step in range(20):
+            x = torch.randn(200, 100).detach()
+            x.requires_grad = True
+            y = m(x)
+            dy = torch.randn(200, 100).detach()
+            f = (y * dy).sum()
+            f.backward()
+
+            optim.step()
+            scheduler.step_batch()
+            optim.zero_grad()
+
+    logging.info(f"last lr = {scheduler.get_last_lr()}")
+    logging.info(f"state dict = {scheduler.state_dict()}")
+
+
+if __name__ == "__main__":
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+    logging.getLogger().setLevel(logging.INFO)
+    import subprocess
+
+    s = subprocess.check_output(
+        "git status -uno .; git log -1; git diff HEAD .", shell=True
+    )
+    logging.info(s)
+
+    _test_eden()

zipvoice/utils/optim.py

@@ -0,0 +1,868 @@
+# Copyright      2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# 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 contextlib
+import logging
+from collections import defaultdict
+from typing import Dict, List, Tuple
+
+import torch
+from lhotse.utils import fix_random_seed
+from torch import Tensor
+from torch.optim import Optimizer
+
+
+class BatchedOptimizer(Optimizer):
+    """
+    This class adds to class Optimizer the capability to optimize parameters in batches:
+    it will stack the parameters and their grads for you so the optimizer can work
+    on tensors with an extra leading dimension.  This is intended for speed with GPUs,
+    as it reduces the number of kernels launched in the optimizer.
+
+    Args:
+      params:
+    """
+
+    def __init__(self, params, defaults):
+        super(BatchedOptimizer, self).__init__(params, defaults)
+
+    @contextlib.contextmanager
+    def batched_params(self, param_group, group_params_names):
+        """
+        This function returns (technically, yields) a list of
+          of tuples (p, state), where
+        p is a `fake` parameter that is stacked (over axis 0) from real parameters
+        that share the same shape, and its gradient is also stacked;
+        `state` is the state corresponding to this batch of parameters
+        (it will be physically located in the "state" for one of the real
+        parameters, the last one that has any particular shape and dtype).
+
+        This function is decorated as a context manager so that it can
+        write parameters back to their "real" locations.
+
+        The idea is, instead of doing:
+        <code>
+          for p in group["params"]:
+             state = self.state[p]
+             ...
+        </code>
+        you can do:
+        <code>
+          with self.batched_params(group["params"]) as batches:
+             for p, state, p_names in batches:
+                 ...
+        </code>
+
+        Args:
+          group: a parameter group, which is a list of parameters; should be
+                one of self.param_groups.
+          group_params_names: name for each parameter in group,
+                which is List[str].
+        """
+        batches = defaultdict(
+            list
+        )  # `batches` maps from tuple (dtype_as_str,*shape) to list of nn.Parameter
+        batches_names = defaultdict(
+            list
+        )  # `batches` maps from tuple (dtype_as_str,*shape) to list of str
+
+        assert len(param_group) == len(group_params_names)
+        for p, named_p in zip(param_group, group_params_names):
+            key = (str(p.dtype), *p.shape)
+            batches[key].append(p)
+            batches_names[key].append(named_p)
+
+        batches_names_keys = list(batches_names.keys())
+        sorted_idx = sorted(
+            range(len(batches_names)), key=lambda i: batches_names_keys[i]
+        )
+        batches_names = [batches_names[batches_names_keys[idx]] for idx in sorted_idx]
+        batches = [batches[batches_names_keys[idx]] for idx in sorted_idx]
+
+        stacked_params_dict = dict()
+
+        # turn batches into a list, in deterministic order.
+        # tuples will contain tuples of (stacked_param, state, stacked_params_names),
+        # one for each batch in `batches`.
+        tuples = []
+
+        for batch, batch_names in zip(batches, batches_names):
+            p = batch[0]
+            # we arbitrarily store the state in the
+            # state corresponding to the 1st parameter in the
+            # group.  class Optimizer will take care of saving/loading state.
+            state = self.state[p]
+            p_stacked = torch.stack(batch)
+            grad = torch.stack(
+                [torch.zeros_like(p) if p.grad is None else p.grad for p in batch]
+            )
+            p_stacked.grad = grad
+            stacked_params_dict[key] = p_stacked
+            tuples.append((p_stacked, state, batch_names))
+
+        yield tuples  # <-- calling code will do the actual optimization here!
+
+        for (stacked_params, _state, _names), batch in zip(tuples, batches):
+            for i, p in enumerate(batch):  # batch is list of Parameter
+                p.copy_(stacked_params[i])
+
+
+def basic_step(group, p, state, grad):
+    # computes basic Adam update using beta2 (dividing by gradient stddev) only.  no
+    # momentum yet.
+    lr = group["lr"]
+    if p.numel() == p.shape[0]:
+        lr = lr * group["scalar_lr_scale"]
+    beta2 = group["betas"][1]
+    eps = group["eps"]
+    # p shape: (batch_size,) or (batch_size, 1, [1,..])
+    try:
+        exp_avg_sq = state[
+            "exp_avg_sq"
+        ]  # shape: (batch_size,) or (batch_size, 1, [1,..])
+    except KeyError:
+        exp_avg_sq = torch.zeros(*p.shape, device=p.device, dtype=torch.float)
+        state["exp_avg_sq"] = exp_avg_sq
+
+    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+    # bias_correction2 is like in Adam.
+    # slower update at the start will help stability anyway.
+    bias_correction2 = 1 - beta2 ** (state["step"] + 1)
+    if bias_correction2 < 0.99:
+        # note: not in-place.
+        exp_avg_sq = exp_avg_sq * (1.0 / bias_correction2)
+    denom = exp_avg_sq.sqrt().add_(eps)
+
+    return -lr * grad / denom
+
+
+def scaling_step(group, p, state, grad):
+    delta = basic_step(group, p, state, grad)
+    if p.numel() == p.shape[0]:
+        return delta
+    # there is no scaling for scalar parameters.
+    # (p.shape[0] is the batch of parameters.)
+
+    step = state["step"]
+    size_update_period = group["size_update_period"]
+
+    try:
+        param_rms = state["param_rms"]
+        scale_grads = state["scale_grads"]
+        scale_exp_avg_sq = state["scale_exp_avg_sq"]
+    except KeyError:
+        # we know p.ndim > 1 because we'd have returned above if not, so don't worry
+        # about the speial case of dim=[] that pytorch treats inconsistently.
+        param_rms = (p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt()
+        param_rms = param_rms.to(torch.float)
+        scale_exp_avg_sq = torch.zeros_like(param_rms)
+        scale_grads = torch.zeros(
+            size_update_period,
+            *param_rms.shape,
+            dtype=torch.float,
+            device=p.device,
+        )
+        state["param_rms"] = param_rms
+        state["scale_grads"] = scale_grads
+        state["scale_exp_avg_sq"] = scale_exp_avg_sq
+
+    # on every step, update the gradient w.r.t. the scale of the parameter, we
+    # store these as a batch and periodically update the size (for speed only, to
+    # avoid too many operations).
+    scale_grads[step % size_update_period] = (p * grad).sum(
+        dim=list(range(1, p.ndim)), keepdim=True
+    )
+
+    # periodically recompute the value of param_rms.
+    if step % size_update_period == size_update_period - 1:
+        param_rms.copy_((p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt())
+
+    param_min_rms = group["param_min_rms"]
+
+    # scale the step size by param_rms.  This is the most important "scaling" part of
+    # ScaledAdam
+    delta *= param_rms.clamp(min=param_min_rms)
+
+    if step % size_update_period == size_update_period - 1 and step > 0:
+        # This block updates the size of parameter by adding a step ("delta") value in
+        # the direction of either shrinking or growing it.
+        beta2 = group["betas"][1]
+        size_lr = group["lr"] * group["scalar_lr_scale"]
+        param_max_rms = group["param_max_rms"]
+        eps = group["eps"]
+        # correct beta2 for the size update period: we will have
+        # faster decay at this level.
+        beta2_corr = beta2**size_update_period
+        scale_exp_avg_sq.mul_(beta2_corr).add_(
+            (scale_grads**2).mean(dim=0),  # mean over dim `size_update_period`
+            alpha=1 - beta2_corr,
+        )  # shape is (batch_size, 1, 1, ...)
+
+        # The 1st time we reach here is when size_step == 1.
+        size_step = (step + 1) // size_update_period
+        bias_correction2 = 1 - beta2_corr**size_step
+
+        denom = scale_exp_avg_sq.sqrt() + eps
+
+        scale_step = (
+            -size_lr * (bias_correction2**0.5) * scale_grads.sum(dim=0) / denom
+        )
+
+        is_too_small = param_rms < param_min_rms
+
+        # when the param gets too small, just don't shrink it any further.
+        scale_step.masked_fill_(is_too_small, 0.0)
+
+        # The following may help prevent instability: don't allow the scale step to be
+        # too large in either direction.
+        scale_step.clamp_(min=-0.1, max=0.1)
+
+        # and ensure the parameter rms after update never exceeds param_max_rms.
+        # We have to look at the trained model for parameters at or around the
+        # param_max_rms, because sometimes they can indicate a problem with the
+        # topology or settings.
+        scale_step = torch.minimum(scale_step, (param_max_rms - param_rms) / param_rms)
+
+        delta.add_(p * scale_step)
+
+    return delta
+
+
+def momentum_step(group, p, state, grad):
+    delta = scaling_step(group, p, state, grad)
+    beta1 = group["betas"][0]
+    try:
+        stored_delta = state["delta"]
+    except KeyError:
+        stored_delta = torch.zeros(*p.shape, device=p.device, dtype=torch.float)
+        state["delta"] = stored_delta
+    stored_delta.mul_(beta1)
+    stored_delta.add_(delta, alpha=(1 - beta1))
+    # we don't bother doing the "bias correction" part of Adam for beta1 because this is
+    # just an edge effect that affects the first 10 or so batches; and the effect of not
+    # doing it is just to do a slower update for the first few batches, which will help
+    # stability.
+    return stored_delta
+
+
+class ScaledAdam(BatchedOptimizer):
+    """
+     Implements 'Scaled Adam', a variant of Adam where we scale each parameter's update
+     proportional to the norm of that parameter; and also learn the scale of the
+     parameter, in log space, subject to upper and lower limits (as if we had factored
+     each parameter as param = underlying_param * log_scale.exp())
+
+
+     Args:
+          params: The parameters or param_groups to optimize (like other Optimizer
+                    subclasses) Unlike common optimizers, which accept
+                    model.parameters() or groups of parameters(), this optimizer
+                    could accept model.named_parameters() or groups of
+                    named_parameters(). See comments of function
+                    _get_names_of_parameters for its 4 possible cases.
+              lr:  The learning rate.  We will typically use a learning rate schedule
+                    that starts at 0.03 and decreases over time, i.e. much higher
+                    than other common optimizers.
+     clipping_scale: (e.g. 2.0)
+                   A scale for gradient-clipping: if specified, the normalized gradients
+                   over the whole model will be clipped to have 2-norm equal to
+                   `clipping_scale` times the median 2-norm over the most recent period
+                   of `clipping_update_period` minibatches.  By "normalized gradients",
+                   we mean after multiplying by the rms parameter value for this tensor
+                   [for non-scalars]; this is appropriate because our update is scaled
+                   by this quantity.
+            betas: beta1,beta2 are momentum constants for regular momentum, and moving
+                    sum-sq grad. Must satisfy 0 < beta <= beta2 < 1.
+     scalar_lr_scale: A scaling factor on the learning rate, that we use to update the
+                    scale of each parameter tensor and scalar parameters of the mode..
+                    If each parameter were decomposed as p * p_scale.exp(),
+                    where (p**2).mean().sqrt() == 1.0, scalar_lr_scale would be a the
+                    scaling factor on the learning rate of p_scale.
+              eps:  A general-purpose epsilon to prevent division by zero
+    param_min_rms: Minimum root-mean-square value of parameter tensor, for purposes of
+                   learning the scale on the parameters (we'll constrain the rms of
+                   each non-scalar parameter tensor to be >= this value)
+    param_max_rms: Maximum root-mean-square value of parameter tensor, for purposes of
+                   learning the scale on the parameters (we'll constrain the rms of
+                   each non-scalar parameter tensor to be <= this value)
+       scalar_max: Maximum absolute value for scalar parameters (applicable if your
+                   model has any parameters with numel() == 1).
+    size_update_period: The periodicity, in steps, with which we update the size (scale)
+                   of the parameter tensor.  This is provided to save a little time
+                   in the update.
+     clipping_update_period: if clipping_scale is specified, this is the period
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=3e-02,
+        clipping_scale=None,
+        betas=(0.9, 0.98),
+        scalar_lr_scale=0.1,
+        eps=1.0e-08,
+        param_min_rms=1.0e-05,
+        param_max_rms=3.0,
+        scalar_max=10.0,
+        size_update_period=4,
+        clipping_update_period=100,
+    ):
+
+        defaults = dict(
+            lr=lr,
+            clipping_scale=clipping_scale,
+            betas=betas,
+            scalar_lr_scale=scalar_lr_scale,
+            eps=eps,
+            param_min_rms=param_min_rms,
+            param_max_rms=param_max_rms,
+            scalar_max=scalar_max,
+            size_update_period=size_update_period,
+            clipping_update_period=clipping_update_period,
+        )
+
+        # If params only contains parameters or group of parameters,
+        # i.e when parameter names are not given,
+        # this flag will be set to False in funciton _get_names_of_parameters.
+        self.show_dominant_parameters = True
+        param_groups, parameters_names = self._get_names_of_parameters(params)
+        super(ScaledAdam, self).__init__(param_groups, defaults)
+        assert len(self.param_groups) == len(parameters_names)
+        self.parameters_names = parameters_names
+
+    def _get_names_of_parameters(
+        self, params_or_named_params
+    ) -> Tuple[List[Dict], List[List[str]]]:
+        """
+        Args:
+          params_or_named_params: according to the way ScaledAdam is initialized
+            in train.py, this argument could be one of following 4 cases,
+            case 1, a generator of parameter, e.g.:
+              optimizer = ScaledAdam(model.parameters(), lr=params.base_lr,
+                clipping_scale=3.0)
+
+            case 2, a list of parameter groups with different config, e.g.:
+              model_param_groups = [
+                      {'params': model.encoder.parameters(), 'lr': 0.05},
+                      {'params': model.decoder.parameters(), 'lr': 0.01},
+                      {'params': model.joiner.parameters(), 'lr': 0.03},
+                      ]
+              optimizer = ScaledAdam(model_param_groups, lr=params.base_lr,
+                clipping_scale=3.0)
+
+            case 3, a generator of named_parameter, e.g.:
+              optimizer = ScaledAdam(model.named_parameters(), lr=params.base_lr,
+                clipping_scale=3.0)
+
+            case 4, a list of named_parameter groups with different config, e.g.:
+              model_named_param_groups = [
+                      {'named_params': model.encoder.named_parameters(), 'lr': 0.05},
+                      {'named_params': model.decoder.named_parameters(), 'lr': 0.01},
+                      {'named_params': model.joiner.named_parameters(), 'lr': 0.03},
+                      ]
+              optimizer = ScaledAdam(model_named_param_groups, lr=params.base_lr,
+                clipping_scale=3.0)
+
+          For case 1 and case 2, input params is used to initialize the underlying
+            torch.optimizer.
+          For case 3 and case 4, firstly, names and params are extracted from input
+            named_params, then, these extracted params are used to initialize the
+            underlying torch.optimizer, and these extracted names are mainly used by
+            function `_show_gradient_dominating_parameter`
+
+        Returns:
+          Returns a tuple containing 2 elements:
+            - `param_groups` with type List[Dict], each Dict element is a parameter
+                group. An example of `param_groups` could be:
+              [
+                  {'params': `one iterable of Parameter`, 'lr': 0.05},
+                  {'params': `another iterable of Parameter`, 'lr': 0.08},
+                  {'params': `a third iterable of Parameter`, 'lr': 0.1},
+              ]
+            - `param_gruops_names` with type List[List[str]],
+               each `List[str]` is for a group['params'] in param_groups,
+               and each `str` is the name of a parameter.
+               A dummy name "foo" is related to each parameter,
+               if input are params without names, i.e. case 1 or case 2.
+        """
+        # variable naming convention in this function:
+        #   p is short for param.
+        #   np is short for named_param.
+        #   p_or_np is short for param_or_named_param.
+        #   cur is short for current.
+        #   group is a dict,
+        #   e.g. {'params': iterable of parameter, 'lr': 0.05, other fields}.
+        #   groups is a List[group]
+
+        iterable_or_groups = list(params_or_named_params)
+        if len(iterable_or_groups) == 0:
+            raise ValueError("optimizer got an empty parameter list")
+
+        # The first value of returned tuple.  A list of dicts containing at
+        # least 'params' as a key.
+        param_groups = []
+
+        # The second value of returned tuple,
+        # a List[List[str]], each sub-List is for a group.
+        param_groups_names = []
+
+        if not isinstance(iterable_or_groups[0], dict):
+            # case 1 or case 3,
+            # the input is an iterable of parameter or named parameter.
+            param_iterable_cur_group = []
+            param_names_cur_group = []
+            for p_or_np in iterable_or_groups:
+                if isinstance(p_or_np, tuple):
+                    # case 3
+                    name, param = p_or_np
+                else:
+                    # case 1
+                    assert isinstance(p_or_np, torch.Tensor)
+                    param = p_or_np
+                    # Assign a dummy name as a placeholder
+                    name = "foo"
+                    self.show_dominant_parameters = False
+                param_iterable_cur_group.append(param)
+                param_names_cur_group.append(name)
+            param_groups.append({"params": param_iterable_cur_group})
+            param_groups_names.append(param_names_cur_group)
+        else:
+            # case 2 or case 4
+            # the input is groups of parameter or named parameter.
+            for cur_group in iterable_or_groups:
+                if "named_params" in cur_group:
+                    name_list = [x[0] for x in cur_group["named_params"]]
+                    p_list = [x[1] for x in cur_group["named_params"]]
+                    del cur_group["named_params"]
+                    cur_group["params"] = p_list
+                else:
+                    assert "params" in cur_group
+                    name_list = ["foo" for _ in cur_group["params"]]
+                param_groups.append(cur_group)
+                param_groups_names.append(name_list)
+
+        return param_groups, param_groups_names
+
+    def __setstate__(self, state):
+        super(ScaledAdam, self).__setstate__(state)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group, group_params_names in zip(self.param_groups, self.parameters_names):
+
+            with self.batched_params(group["params"], group_params_names) as batches:
+
+                # batches is list of pairs (stacked_param, state).  stacked_param is
+                # like a regular parameter, and will have a .grad, but the 1st dim
+                # corresponds to a stacking dim, it is not a real dim.
+
+                if (
+                    len(batches[0][1]) == 0
+                ):  # if len(first state) == 0: not yet initialized
+                    clipping_scale = 1
+                else:
+                    clipping_scale = self._get_clipping_scale(group, batches)
+
+                for p, state, _ in batches:
+                    # Perform optimization step.
+                    # grad is not going to be None, we handled that when creating the
+                    # batches.
+                    grad = p.grad
+                    if grad.is_sparse:
+                        raise RuntimeError(
+                            "ScaledAdam optimizer does not support sparse gradients"
+                        )
+
+                    try:
+                        cur_step = state["step"]
+                    except KeyError:
+                        state["step"] = 0
+                        cur_step = 0
+
+                    grad = (
+                        p.grad if clipping_scale == 1.0 else p.grad.mul_(clipping_scale)
+                    )
+                    p += momentum_step(group, p.detach(), state, grad)
+
+                    if p.numel() == p.shape[0]:  # scalar parameter
+                        scalar_max = group["scalar_max"]
+                        p.clamp_(min=-scalar_max, max=scalar_max)
+
+                    state["step"] = cur_step + 1
+
+        return loss
+
+    def _get_clipping_scale(
+        self, group: dict, tuples: List[Tuple[Tensor, dict, List[str]]]
+    ) -> float:
+        """
+        Returns a scalar factor <= 1.0 that dictates gradient clipping, i.e. we will
+        scale the gradients by this amount before applying the rest of the update.
+
+        Args:
+           group: the parameter group, an item in self.param_groups
+           tuples: a list of tuples of (param, state, param_names)
+                where param is a batched set of parameters,
+                with a .grad (1st dim is batch dim)
+                and state is the state-dict where optimization parameters are kept.
+                param_names is a List[str] while each str is name for a parameter
+                in batched set of parameters "param".
+        """
+        assert len(tuples) >= 1
+        clipping_scale = group["clipping_scale"]
+        (first_p, first_state, _) = tuples[0]
+        step = first_state["step"]
+        if clipping_scale is None or step == 0:
+            # no clipping.  return early on step == 0 because the other
+            # parameters' state won't have been initialized yet.
+            return 1.0
+        clipping_update_period = group["clipping_update_period"]
+        scalar_lr_scale = group["scalar_lr_scale"]
+
+        tot_sumsq = torch.tensor(0.0, device=first_p.device)
+        for p, state, param_names in tuples:
+            grad = p.grad
+            if grad.is_sparse:
+                raise RuntimeError(
+                    "ScaledAdam optimizer does not support sparse gradients"
+                )
+            if p.numel() == p.shape[0]:  # a batch of scalars
+                tot_sumsq += (grad**2).sum() * (
+                    scalar_lr_scale**2
+                )  # sum() to change shape [1] to []
+            else:
+                tot_sumsq += ((grad * state["param_rms"]) ** 2).sum()
+
+        tot_norm = tot_sumsq.sqrt()
+        if "model_norms" not in first_state:
+            first_state["model_norms"] = torch.zeros(
+                clipping_update_period, device=p.device
+            )
+        first_state["model_norms"][step % clipping_update_period] = tot_norm
+
+        irregular_estimate_steps = [
+            i for i in [10, 20, 40] if i < clipping_update_period
+        ]
+        if step % clipping_update_period == 0 or step in irregular_estimate_steps:
+            # Print some stats.
+            # We don't reach here if step == 0 because we would have returned
+            # above.
+            sorted_norms = first_state["model_norms"].sort()[0].to("cpu")
+            if step in irregular_estimate_steps:
+                sorted_norms = sorted_norms[-step:]
+            num_norms = sorted_norms.numel()
+            quartiles = []
+            for n in range(0, 5):
+                index = min(num_norms - 1, (num_norms // 4) * n)
+                quartiles.append(sorted_norms[index].item())
+
+            median = quartiles[2]
+            if median - median != 0:
+                raise RuntimeError("Too many grads were not finite")
+            threshold = clipping_scale * median
+            if step in irregular_estimate_steps:
+                # use larger thresholds on first few steps of estimating threshold,
+                # as norm may be changing rapidly.
+                threshold = threshold * 2.0
+            first_state["model_norm_threshold"] = threshold
+            percent_clipped = (
+                first_state["num_clipped"] * 100.0 / num_norms
+                if "num_clipped" in first_state
+                else 0.0
+            )
+            first_state["num_clipped"] = 0
+            quartiles = " ".join(["%.3e" % x for x in quartiles])
+            logging.warning(
+                f"Clipping_scale={clipping_scale}, grad-norm quartiles {quartiles}, "
+                f"threshold={threshold:.3e}, percent-clipped={percent_clipped:.1f}"
+            )
+
+        try:
+            model_norm_threshold = first_state["model_norm_threshold"]
+        except KeyError:
+            return 1.0  # threshold has not yet been set.
+
+        ans = min(1.0, (model_norm_threshold / (tot_norm + 1.0e-20)).item())
+        if ans != ans:  # e.g. ans is nan
+            ans = 0.0
+        if ans < 1.0:
+            first_state["num_clipped"] += 1
+        if ans < 0.5:
+            logging.debug(
+                f"Scaling gradients by {ans}, "
+                f"model_norm_threshold={model_norm_threshold}"
+            )
+            if self.show_dominant_parameters:
+                assert p.shape[0] == len(param_names)
+                self._show_gradient_dominating_parameter(
+                    tuples, tot_sumsq, group["scalar_lr_scale"]
+                )
+                self._show_param_with_unusual_grad(tuples)
+
+        if ans == 0.0:
+            for p, state, param_names in tuples:
+                p.grad.zero_()  # get rid of infinity()
+
+        return ans
+
+    def _show_param_with_unusual_grad(
+        self,
+        tuples: List[Tuple[Tensor, dict, List[str]]],
+    ):
+        """
+        Print information about parameter which has the largest ratio of
+        grad-on-this-batch divided by normal grad size.
+           tuples: a list of tuples of (param, state, param_names)
+                where param is a batched set of parameters,
+                with a .grad (1st dim is batch dim)
+                and state is the state-dict where optimization parameters are kept.
+                param_names is a List[str] while each str is name for a parameter
+                in batched set of parameters "param".
+        """
+        # ratios_names is a list of 3-tuples: (grad_ratio, param_name, tensor)
+        ratios_names = []
+        for p, state, batch_param_names in tuples:
+            dims = list(range(1, p.ndim))
+
+            def mean(x):
+                # workaround for bad interface of torch's "mean" for when dims is the
+                # empty list.
+                if len(dims) > 0:
+                    return x.mean(dim=dims)
+                else:
+                    return x
+
+            grad_ratio = (
+                (mean(p.grad**2) / state["exp_avg_sq"].mean(dim=dims))
+                .sqrt()
+                .to("cpu")
+            )
+
+            ratios_names += zip(
+                grad_ratio.tolist(), batch_param_names, p.grad.unbind(dim=0)
+            )
+
+        ratios_names = sorted(ratios_names, reverse=True)
+        ratios_names = ratios_names[:10]
+        ratios_names = [
+            (ratio, name, largest_index(tensor))
+            for (ratio, name, tensor) in ratios_names
+        ]
+
+        logging.debug(
+            f"Parameters with most larger-than-usual grads, with ratios, "
+            f"are: {ratios_names}"
+        )
+
+    def _show_gradient_dominating_parameter(
+        self,
+        tuples: List[Tuple[Tensor, dict, List[str]]],
+        tot_sumsq: Tensor,
+        scalar_lr_scale: float,
+    ):
+        """
+        Show information of parameter which dominates tot_sumsq.
+
+        Args:
+           tuples: a list of tuples of (param, state, param_names)
+                where param is a batched set of parameters,
+                with a .grad (1st dim is batch dim)
+                and state is the state-dict where optimization parameters are kept.
+                param_names is a List[str] while each str is name for a parameter
+                in batched set of parameters "param".
+            tot_sumsq: sumsq of all parameters. Though it's could be calculated
+                from tuples, we still pass it to save some time.
+        """
+        all_sumsq_orig = {}
+        for p, state, batch_param_names in tuples:
+            # p is a stacked batch parameters.
+            batch_grad = p.grad
+            if p.numel() == p.shape[0]:  # a batch of scalars
+                # Dummy values used by following `zip` statement.
+                batch_rms_orig = torch.full(
+                    p.shape, scalar_lr_scale, device=batch_grad.device
+                )
+            else:
+                batch_rms_orig = state["param_rms"]
+            batch_sumsq_orig = (batch_grad * batch_rms_orig) ** 2
+            if batch_grad.ndim > 1:
+                # need to guard it with if-statement because sum() sums over
+                # all dims if dim == ().
+                batch_sumsq_orig = batch_sumsq_orig.sum(
+                    dim=list(range(1, batch_grad.ndim))
+                )
+            for name, sumsq_orig, rms, grad in zip(
+                batch_param_names, batch_sumsq_orig, batch_rms_orig, batch_grad
+            ):
+
+                proportion_orig = sumsq_orig / tot_sumsq
+                all_sumsq_orig[name] = (proportion_orig, sumsq_orig, rms, grad)
+
+        sorted_by_proportion = {
+            k: v
+            for k, v in sorted(
+                all_sumsq_orig.items(),
+                key=lambda item: item[1][0],
+                reverse=True,
+            )
+        }
+        dominant_param_name = next(iter(sorted_by_proportion))
+        (
+            dominant_proportion,
+            dominant_sumsq,
+            dominant_rms,
+            dominant_grad,
+        ) = sorted_by_proportion[dominant_param_name]
+        logging.debug(
+            f"Parameter dominating tot_sumsq {dominant_param_name}"
+            f" with proportion {dominant_proportion:.2f},"
+            f" where dominant_sumsq=(grad_sumsq*orig_rms_sq)"
+            f"={dominant_sumsq:.3e},"
+            f" grad_sumsq={(dominant_grad**2).sum():.3e},"
+            f" orig_rms_sq={(dominant_rms**2).item():.3e}"
+        )
+
+
+def largest_index(x: Tensor):
+    x = x.contiguous()
+    argmax = x.abs().argmax().item()
+    return [(argmax // x.stride(i)) % x.size(i) for i in range(x.ndim)]
+
+
+def _test_scaled_adam(hidden_dim: int):
+    import timeit
+
+    from zipvoice.models.modules.scaling import ScaledLinear
+    from zipvoice.utils.lr_scheduler import Eden
+
+    E = 100
+    B = 4
+    T = 2
+    logging.info("in test_eve_cain")
+    # device = torch.device('cuda')
+    device = torch.device("cpu")
+    dtype = torch.float32
+
+    fix_random_seed(42)
+    # these input_magnitudes and output_magnitudes are to test that
+    # Abel is working as we expect and is able to adjust scales of
+    # different dims differently.
+    input_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
+    output_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
+
+    fix_random_seed(42)
+    Linear = ScaledLinear
+
+    m = torch.nn.Sequential(
+        Linear(E, hidden_dim),
+        torch.nn.PReLU(),
+        Linear(hidden_dim, hidden_dim),
+        torch.nn.PReLU(),
+        Linear(hidden_dim, E),
+    ).to(device)
+
+    train_pairs = [
+        (
+            100.0 * torch.randn(B, T, E, device=device, dtype=dtype) * input_magnitudes,
+            torch.randn(B, T, E, device=device, dtype=dtype) * output_magnitudes,
+        )
+        for _ in range(20)
+    ]
+    optim = ScaledAdam(m.named_parameters(), lr=0.03, clipping_scale=2.0)
+    scheduler = Eden(optim, lr_batches=200, lr_epochs=5, verbose=False)
+
+    start = timeit.default_timer()
+    avg_loss = 0.0
+    for epoch in range(180):
+        scheduler.step_epoch()
+        # if epoch == 100 and iter in [2,3]:
+        #    optim.reset_speedup()  # check it doesn't crash.
+
+        # if epoch == 130:
+        #    opts = diagnostics.TensorDiagnosticOptions(
+        #        512
+        #    )  # allow 4 megabytes per sub-module
+        #    diagnostic = diagnostics.attach_diagnostics(m, opts)
+
+        for n, (x, y) in enumerate(train_pairs):
+            y_out = m(x)
+            loss = ((y_out - y) ** 2).mean() * 100.0
+            if epoch == 0 and n == 0:
+                avg_loss = loss.item()
+            else:
+                avg_loss = 0.98 * avg_loss + 0.02 * loss.item()
+            if n == 0 and epoch % 5 == 0:
+                # norm1 = '%.2e' % (m[0].weight**2).mean().sqrt().item()
+                # norm1b = '%.2e' % (m[0].bias**2).mean().sqrt().item()
+                # norm2 = '%.2e' % (m[2].weight**2).mean().sqrt().item()
+                # norm2b = '%.2e' % (m[2].bias**2).mean().sqrt().item()
+                # scale1 = '%.2e' % (m[0].weight_scale.exp().item())
+                # scale1b = '%.2e' % (m[0].bias_scale.exp().item())
+                # scale2 = '%.2e' % (m[2].weight_scale.exp().item())
+                # scale2b = '%.2e' % (m[2].bias_scale.exp().item())
+                lr = scheduler.get_last_lr()[0]
+                logging.info(
+                    f"Iter {iter}, epoch {epoch}, batch {n}, "
+                    f"avg_loss {avg_loss:.4g}, lr={lr:.4e}"
+                )  # , norms={norm1,norm1b,norm2,norm2b}")
+                # scales={scale1,scale1b,scale2,scale2b}
+            loss.log().backward()
+            optim.step()
+            optim.zero_grad()
+            scheduler.step_batch()
+
+        # diagnostic.print_diagnostics()
+
+        stop = timeit.default_timer()
+        logging.info(f"Iter={iter}, Time taken: {stop - start}")
+
+        logging.info(f"last lr = {scheduler.get_last_lr()}")
+        # logging.info("state dict = ", scheduler.state_dict())
+        # logging.info("optim state_dict = ", optim.state_dict())
+        logging.info(f"input_magnitudes = {input_magnitudes}")
+        logging.info(f"output_magnitudes = {output_magnitudes}")
+
+
+if __name__ == "__main__":
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+    logging.getLogger().setLevel(logging.INFO)
+    import subprocess
+
+    s = subprocess.check_output(
+        "git status -uno .; git log -1; git diff HEAD .", shell=True
+    )
+    logging.info(s)
+    import sys
+
+    if len(sys.argv) > 1:
+        hidden_dim = int(sys.argv[1])
+    else:
+        hidden_dim = 200
+
+    _test_scaled_adam(hidden_dim)

zipvoice/utils/scaling_converter.py

@@ -0,0 +1,105 @@
+# Copyright    2022-2023  Xiaomi Corp.        (authors: Fangjun Kuang,
+#                                                       Zengwei Yao)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This file replaces various modules in a model.
+Specifically, ActivationBalancer is replaced with an identity operator;
+Whiten is also replaced with an identity operator;
+BasicNorm is replaced by a module with `exp` removed.
+"""
+
+import copy
+from typing import List
+
+import torch
+import torch.nn as nn
+
+from zipvoice.models.modules.scaling import (
+    Balancer,
+    Dropout3,
+    SwooshL,
+    SwooshLOnnx,
+    SwooshR,
+    SwooshROnnx,
+    Whiten,
+)
+from zipvoice.models.modules.zipformer import CompactRelPositionalEncoding
+
+
+# Copied from https://pytorch.org/docs/1.9.0/_modules/torch/nn/modules/module.html#Module.get_submodule  # noqa
+# get_submodule was added to nn.Module at v1.9.0
+def get_submodule(model, target):
+    if target == "":
+        return model
+    atoms: List[str] = target.split(".")
+    mod: torch.nn.Module = model
+    for item in atoms:
+        if not hasattr(mod, item):
+            raise AttributeError(
+                mod._get_name() + " has no " "attribute `" + item + "`"
+            )
+        mod = getattr(mod, item)
+        if not isinstance(mod, torch.nn.Module):
+            raise AttributeError("`" + item + "` is not " "an nn.Module")
+    return mod
+
+
+def convert_scaled_to_non_scaled(
+    model: nn.Module,
+    inplace: bool = False,
+    is_pnnx: bool = False,
+    is_onnx: bool = False,
+):
+    """
+    Args:
+      model:
+        The model to be converted.
+      inplace:
+        If True, the input model is modified inplace.
+        If False, the input model is copied and we modify the copied version.
+      is_pnnx:
+        True if we are going to export the model for PNNX.
+      is_onnx:
+        True if we are going to export the model for ONNX.
+    Return:
+      Return a model without scaled layers.
+    """
+    if not inplace:
+        model = copy.deepcopy(model)
+
+    d = {}
+    for name, m in model.named_modules():
+        if isinstance(m, (Balancer, Dropout3, Whiten)):
+            d[name] = nn.Identity()
+        elif is_onnx and isinstance(m, SwooshR):
+            d[name] = SwooshROnnx()
+        elif is_onnx and isinstance(m, SwooshL):
+            d[name] = SwooshLOnnx()
+        elif is_onnx and isinstance(m, CompactRelPositionalEncoding):
+            # We want to recreate the positional encoding vector when
+            # the input changes, so we have to use torch.jit.script()
+            # to replace torch.jit.trace()
+            d[name] = torch.jit.script(m)
+
+    for k, v in d.items():
+        if "." in k:
+            parent, child = k.rsplit(".", maxsplit=1)
+            setattr(get_submodule(model, parent), child, v)
+        else:
+            setattr(model, k, v)
+
+    return model

zipvoice/utils/tensorrt.py

@@ -0,0 +1,143 @@
+# Copyright         2025  Nvidia Corp.        (authors: Yuekai Zhang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script provides utility functions for working with TensorRT in ZipVoice.
+"""
+
+import logging
+import os
+import queue
+from typing import Any, Tuple, Optional
+
+import torch
+import torch.nn as nn
+
+
+class TrtContextWrapper:
+    """A wrapper class for managing TensorRT execution contexts."""
+
+    def __init__(
+        self, trt_engine: Any, trt_concurrent: int = 1, device: str = "cuda:0"
+    ):
+        """
+        Initializes the TrtContextWrapper.
+
+        Args:
+            trt_engine (Any): The TensorRT engine.
+            trt_concurrent (int, optional): The number of concurrent contexts. Defaults to 1.
+            device (str, optional): The device to use. Defaults to 'cuda:0'.
+        """
+        self.trt_context_pool = queue.Queue(maxsize=trt_concurrent)
+        self.trt_engine = trt_engine
+        self.device = device
+        for _ in range(trt_concurrent):
+            trt_context = trt_engine.create_execution_context()
+            trt_stream = torch.cuda.stream(torch.cuda.Stream(torch.device(device)))
+            assert trt_context is not None, 'failed to create trt context, maybe not enough CUDA memory, try reduce current trt concurrent {}'.format(trt_concurrent)
+            self.trt_context_pool.put([trt_context, trt_stream])
+        assert self.trt_context_pool.empty() is False, 'no avaialbe estimator context'
+        self.feat_dim = 100
+
+    def acquire_estimator(self) -> Tuple[list, Any]:
+        """Acquires a TensorRT context from the pool."""
+        return self.trt_context_pool.get(), self.trt_engine
+
+    def release_estimator(self, context: Any, stream: Any):
+        """
+        Releases a TensorRT context back to the pool.
+
+        Args:
+            context (Any): The TensorRT context.
+            stream (Any): The CUDA stream.
+        """
+        self.trt_context_pool.put([context, stream])
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        padding_mask: torch.Tensor,
+        guidance_scale: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Executes the TensorRT engine.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+            t (torch.Tensor): The time tensor.
+            padding_mask (torch.Tensor): The padding mask tensor.
+            guidance_scale (torch.Tensor): The guidance scale tensor.
+
+        Returns:
+            torch.Tensor: The output tensor.
+        """
+        x = x.to(torch.float16)
+        t = t.to(torch.float16)
+        padding_mask = padding_mask.to(torch.float16)
+        if guidance_scale is not None:
+            guidance_scale = guidance_scale.to(torch.float16)
+        [estimator, stream], trt_engine = self.acquire_estimator()
+        # NOTE need to synchronize when switching stream
+        torch.cuda.current_stream().synchronize()
+        batch_size = x.size(0)
+        seq_len = x.size(1)
+        
+        # Create output tensor with shape (N, T, 100)
+        output = torch.empty(batch_size, seq_len, self.feat_dim, dtype=x.dtype, device=x.device)
+        
+        with stream:
+            estimator.set_input_shape('x', (batch_size, x.size(1), x.size(2)))
+            estimator.set_input_shape('t', (batch_size,))
+            estimator.set_input_shape('padding_mask', (batch_size, padding_mask.size(1)))
+            if guidance_scale is not None:
+                estimator.set_input_shape('guidance_scale', (batch_size,))
+            
+            # Set input tensor addresses
+            input_data_ptrs = [x.contiguous().data_ptr(), t.contiguous().data_ptr(), padding_mask.contiguous().data_ptr()]
+            if guidance_scale is not None:
+                input_data_ptrs.append(guidance_scale.contiguous().data_ptr())
+            for i, j in enumerate(input_data_ptrs):
+                estimator.set_tensor_address(trt_engine.get_tensor_name(i), j)
+            
+            # Set output tensor address
+            # The output tensor name should be the last tensor name in the engine
+            num_tensors = trt_engine.num_io_tensors
+            output_tensor_name = trt_engine.get_tensor_name(num_tensors - 1)  # Last tensor is output
+            estimator.set_tensor_address(output_tensor_name, output.contiguous().data_ptr())
+            
+            # run trt engine
+            assert estimator.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True
+            torch.cuda.current_stream().synchronize()
+        self.release_estimator(estimator, stream)
+        return output.to(torch.float32)
+
+def load_trt(model: nn.Module, trt_model: str, trt_concurrent: int = 1):
+    """
+    Loads a TensorRT engine and replaces the model's fm_decoder with a TrtContextWrapper.
+
+    Args:
+        model (nn.Module): The model to modify.
+        trt_model (str): The path to the TensorRT engine file.
+        trt_concurrent (int, optional): The number of concurrent contexts. Defaults to 1.
+    """
+    assert os.path.exists(trt_model), f"Please export trt model first."
+    import tensorrt as trt
+    with open(trt_model, 'rb') as f:
+        estimator_engine = trt.Runtime(trt.Logger(trt.Logger.INFO)).deserialize_cuda_engine(f.read())
+    assert estimator_engine is not None, 'failed to load trt {}'.format(trt_model)
+    del model.fm_decoder
+    model.fm_decoder = TrtContextWrapper(estimator_engine, trt_concurrent=trt_concurrent, device='cuda')