Initial commit
This commit is contained in:
Yatharth Sharma
670
zipvoice/utils/common.py
Normal file
670
zipvoice/utils/common.py
Normal file
@@ -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()]
|
||||
Reference in New Issue
Block a user