mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-09-08 00:24:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove pruned_transducer_stateless4b

Browse Source
This commit is contained in:
Daniel Povey 2022-05-31 12:29:45 +08:00
parent 7011956c6c
commit 8f877efec5
15 changed files with 0 additions and 5238 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
egs/librispeech/ASR/pruned_transducer_stateless4b/__init__.py
View File

@ -1 +0,0 @@
../pruned_transducer_stateless2/__init__.py

1
egs/librispeech/ASR/pruned_transducer_stateless4b/asr_datamodule.py
View File

@ -1 +0,0 @@
../pruned_transducer_stateless2/asr_datamodule.py

1
egs/librispeech/ASR/pruned_transducer_stateless4b/beam_search.py
View File

@ -1 +0,0 @@
../pruned_transducer_stateless2/beam_search.py

1202
egs/librispeech/ASR/pruned_transducer_stateless4b/conformer.py
View File

File diff suppressed because it is too large Load Diff

633
egs/librispeech/ASR/pruned_transducer_stateless4b/decode.py
View File

@ -1,633 +0,0 @@
#!/usr/bin/env python3
#
# Copyright 2021-2022 Xiaomi Corporation (Author: 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.
"""
Usage:
(1) greedy search
./pruned_transducer_stateless4/decode.py \
--epoch 30 \
--avg 15 \
--exp-dir ./pruned_transducer_stateless4/exp \
--max-duration 600 \
--decoding-method greedy_search
(2) beam search (not recommended)
./pruned_transducer_stateless4/decode.py \
--epoch 30 \
--avg 15 \
--exp-dir ./pruned_transducer_stateless4/exp \
--max-duration 600 \
--decoding-method beam_search \
--beam-size 4
(3) modified beam search
./pruned_transducer_stateless4/decode.py \
--epoch 30 \
--avg 15 \
--exp-dir ./pruned_transducer_stateless4/exp \
--max-duration 600 \
--decoding-method modified_beam_search \
--beam-size 4
(4) fast beam search
./pruned_transducer_stateless4/decode.py \
--epoch 30 \
--avg 15 \
--exp-dir ./pruned_transducer_stateless4/exp \
--max-duration 600 \
--decoding-method fast_beam_search \
--beam 4 \
--max-contexts 4 \
--max-states 8
"""
import argparse
import logging
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import k2
import sentencepiece as spm
import torch
import torch.nn as nn
from asr_datamodule import LibriSpeechAsrDataModule
from beam_search import (
beam_search,
fast_beam_search_one_best,
greedy_search,
greedy_search_batch,
modified_beam_search,
)
from train import get_params, get_transducer_model
from icefall.checkpoint import (
average_checkpoints,
average_checkpoints_with_averaged_model,
find_checkpoints,
load_checkpoint,
)
from icefall.utils import (
AttributeDict,
setup_logger,
store_transcripts,
str2bool,
write_error_stats,
)
def get_parser():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument(
"--epoch",
type=int,
default=30,
help="""It specifies the checkpoint to use for decoding.
Note: Epoch counts from 1.
You can specify --avg to use more checkpoints for model averaging.""",
)
parser.add_argument(
"--iter",
type=int,
default=0,
help="""If positive, --epoch is ignored and it
will use the checkpoint exp_dir/checkpoint-iter.pt.
You can specify --avg to use more checkpoints for model averaging.
""",
)
parser.add_argument(
"--avg",
type=int,
default=15,
help="Number of checkpoints to average. Automatically select "
"consecutive checkpoints before the checkpoint specified by "
"'--epoch' and '--iter'",
)
parser.add_argument(
"--use-averaged-model",
type=str2bool,
default=False,
help="Whether to load averaged model. Currently it only supports "
"using --epoch. If True, it would decode with the averaged model "
"over the epoch range from `epoch-avg` (excluded) to `epoch`."
"Actually only the models with epoch number of `epoch-avg` and "
"`epoch` are loaded for averaging. ",
)
parser.add_argument(
"--exp-dir",
type=str,
default="pruned_transducer_stateless4/exp",
help="The experiment dir",
)
parser.add_argument(
"--bpe-model",
type=str,
default="data/lang_bpe_500/bpe.model",
help="Path to the BPE model",
)
parser.add_argument(
"--decoding-method",
type=str,
default="greedy_search",
help="""Possible values are:
- greedy_search
- beam_search
- modified_beam_search
- fast_beam_search
""",
)
parser.add_argument(
"--beam-size",
type=int,
default=4,
help="""An integer indicating how many candidates we will keep for each
frame. Used only when --decoding-method is beam_search or
modified_beam_search.""",
)
parser.add_argument(
"--beam",
type=float,
default=4,
help="""A floating point value to calculate the cutoff score during beam
search (i.e., `cutoff = max-score - beam`), which is the same as the
`beam` in Kaldi.
Used only when --decoding-method is fast_beam_search""",
)
parser.add_argument(
"--max-contexts",
type=int,
default=4,
help="""Used only when --decoding-method is
fast_beam_search""",
)
parser.add_argument(
"--max-states",
type=int,
default=8,
help="""Used only when --decoding-method is
fast_beam_search""",
)
parser.add_argument(
"--context-size",
type=int,
default=2,
help="The context size in the decoder. 1 means bigram; "
"2 means tri-gram",
)
parser.add_argument(
"--max-sym-per-frame",
type=int,
default=1,
help="""Maximum number of symbols per frame.
Used only when --decoding_method is greedy_search""",
)
return parser
def decode_one_batch(
params: AttributeDict,
model: nn.Module,
sp: spm.SentencePieceProcessor,
batch: dict,
decoding_graph: Optional[k2.Fsa] = None,
) -> Dict[str, List[List[str]]]:
"""Decode one batch and return the result in a dict. The dict has the
following format:
- key: It indicates the setting used for decoding. For example,
if greedy_search is used, it would be "greedy_search"
If beam search with a beam size of 7 is used, it would be
"beam_7"
- value: It contains the decoding result. `len(value)` equals to
batch size. `value[i]` is the decoding result for the i-th
utterance in the given batch.
Args:
params:
It's the return value of :func:`get_params`.
model:
The neural model.
sp:
The BPE model.
batch:
It is the return value from iterating
`lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
for the format of the `batch`.
decoding_graph:
The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
only when --decoding_method is fast_beam_search.
Returns:
Return the decoding result. See above description for the format of
the returned dict.
"""
device = next(model.parameters()).device
feature = batch["inputs"]
assert feature.ndim == 3
feature = feature.to(device)
# at entry, feature is (N, T, C)
supervisions = batch["supervisions"]
feature_lens = supervisions["num_frames"].to(device)
encoder_out, encoder_out_lens = model.encoder(
x=feature, x_lens=feature_lens
)
hyps = []
if params.decoding_method == "fast_beam_search":
hyp_tokens = fast_beam_search_one_best(
model=model,
decoding_graph=decoding_graph,
encoder_out=encoder_out,
encoder_out_lens=encoder_out_lens,
beam=params.beam,
max_contexts=params.max_contexts,
max_states=params.max_states,
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
elif (
params.decoding_method == "greedy_search"
and params.max_sym_per_frame == 1
):
hyp_tokens = greedy_search_batch(
model=model,
encoder_out=encoder_out,
encoder_out_lens=encoder_out_lens,
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
elif params.decoding_method == "modified_beam_search":
hyp_tokens = modified_beam_search(
model=model,
encoder_out=encoder_out,
encoder_out_lens=encoder_out_lens,
beam=params.beam_size,
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
else:
batch_size = encoder_out.size(0)
for i in range(batch_size):
# fmt: off
encoder_out_i = encoder_out[i:i + 1, :encoder_out_lens[i]]
# fmt: on
if params.decoding_method == "greedy_search":
hyp = greedy_search(
model=model,
encoder_out=encoder_out_i,
max_sym_per_frame=params.max_sym_per_frame,
)
elif params.decoding_method == "beam_search":
hyp = beam_search(
model=model,
encoder_out=encoder_out_i,
beam=params.beam_size,
)
else:
raise ValueError(
f"Unsupported decoding method: {params.decoding_method}"
)
hyps.append(sp.decode(hyp).split())
if params.decoding_method == "greedy_search":
return {"greedy_search": hyps}
elif params.decoding_method == "fast_beam_search":
return {
(
f"beam_{params.beam}_"
f"max_contexts_{params.max_contexts}_"
f"max_states_{params.max_states}"
): hyps
}
else:
return {f"beam_size_{params.beam_size}": hyps}
def decode_dataset(
dl: torch.utils.data.DataLoader,
params: AttributeDict,
model: nn.Module,
sp: spm.SentencePieceProcessor,
decoding_graph: Optional[k2.Fsa] = None,
) -> Dict[str, List[Tuple[List[str], List[str]]]]:
"""Decode dataset.
Args:
dl:
PyTorch's dataloader containing the dataset to decode.
params:
It is returned by :func:`get_params`.
model:
The neural model.
sp:
The BPE model.
decoding_graph:
The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
only when --decoding_method is fast_beam_search.
Returns:
Return a dict, whose key may be "greedy_search" if greedy search
is used, or it may be "beam_7" if beam size of 7 is used.
Its value is a list of tuples. Each tuple contains two elements:
The first is the reference transcript, and the second is the
predicted result.
"""
num_cuts = 0
try:
num_batches = len(dl)
except TypeError:
num_batches = "?"
if params.decoding_method == "greedy_search":
log_interval = 50
else:
log_interval = 10
results = defaultdict(list)
for batch_idx, batch in enumerate(dl):
texts = batch["supervisions"]["text"]
hyps_dict = decode_one_batch(
params=params,
model=model,
sp=sp,
decoding_graph=decoding_graph,
batch=batch,
)
for name, hyps in hyps_dict.items():
this_batch = []
assert len(hyps) == len(texts)
for hyp_words, ref_text in zip(hyps, texts):
ref_words = ref_text.split()
this_batch.append((ref_words, hyp_words))
results[name].extend(this_batch)
num_cuts += len(texts)
if batch_idx % log_interval == 0:
batch_str = f"{batch_idx}/{num_batches}"
logging.info(
f"batch {batch_str}, cuts processed until now is {num_cuts}"
)
return results
def save_results(
params: AttributeDict,
test_set_name: str,
results_dict: Dict[str, List[Tuple[List[int], List[int]]]],
):
test_set_wers = dict()
for key, results in results_dict.items():
recog_path = (
params.res_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
)
store_transcripts(filename=recog_path, texts=results)
logging.info(f"The transcripts are stored in {recog_path}")
# The following prints out WERs, per-word error statistics and aligned
# ref/hyp pairs.
errs_filename = (
params.res_dir / f"errs-{test_set_name}-{key}-{params.suffix}.txt"
)
with open(errs_filename, "w") as f:
wer = write_error_stats(
f, f"{test_set_name}-{key}", results, enable_log=True
)
test_set_wers[key] = wer
logging.info("Wrote detailed error stats to {}".format(errs_filename))
test_set_wers = sorted(test_set_wers.items(), key=lambda x: x[1])
errs_info = (
params.res_dir
/ f"wer-summary-{test_set_name}-{key}-{params.suffix}.txt"
)
with open(errs_info, "w") as f:
print("settings\tWER", file=f)
for key, val in test_set_wers:
print("{}\t{}".format(key, val), file=f)
s = "\nFor {}, WER of different settings are:\n".format(test_set_name)
note = "\tbest for {}".format(test_set_name)
for key, val in test_set_wers:
s += "{}\t{}{}\n".format(key, val, note)
note = ""
logging.info(s)
@torch.no_grad()
def main():
parser = get_parser()
LibriSpeechAsrDataModule.add_arguments(parser)
args = parser.parse_args()
args.exp_dir = Path(args.exp_dir)
params = get_params()
params.update(vars(args))
assert params.decoding_method in (
"greedy_search",
"beam_search",
"fast_beam_search",
"modified_beam_search",
)
params.res_dir = params.exp_dir / params.decoding_method
if params.iter > 0:
params.suffix = f"iter-{params.iter}-avg-{params.avg}"
else:
params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
if "fast_beam_search" in params.decoding_method:
params.suffix += f"-beam-{params.beam}"
params.suffix += f"-max-contexts-{params.max_contexts}"
params.suffix += f"-max-states-{params.max_states}"
elif "beam_search" in params.decoding_method:
params.suffix += (
f"-{params.decoding_method}-beam-size-{params.beam_size}"
)
else:
params.suffix += f"-context-{params.context_size}"
params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
if params.use_averaged_model:
params.suffix += "-use-averaged-model"
setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
logging.info("Decoding started")
device = torch.device("cpu")
if torch.cuda.is_available():
device = torch.device("cuda", 0)
logging.info(f"Device: {device}")
sp = spm.SentencePieceProcessor()
sp.load(params.bpe_model)
# <blk> and <unk> is defined in local/train_bpe_model.py
params.blank_id = sp.piece_to_id("<blk>")
params.unk_id = sp.piece_to_id("<unk>")
params.vocab_size = sp.get_piece_size()
logging.info(params)
logging.info("About to create model")
model = get_transducer_model(params)
if not params.use_averaged_model:
if params.iter > 0:
filenames = find_checkpoints(
params.exp_dir, iteration=-params.iter
)[: params.avg]
if len(filenames) == 0:
raise ValueError(
f"No checkpoints found for"
f" --iter {params.iter}, --avg {params.avg}"
)
elif len(filenames) < params.avg:
raise ValueError(
f"Not enough checkpoints ({len(filenames)}) found for"
f" --iter {params.iter}, --avg {params.avg}"
)
logging.info(f"averaging {filenames}")
model.to(device)
model.load_state_dict(average_checkpoints(filenames, device=device))
elif params.avg == 1:
load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model)
else:
start = params.epoch - params.avg + 1
filenames = []
for i in range(start, params.epoch + 1):
if i >= 1:
filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
logging.info(f"averaging {filenames}")
model.to(device)
model.load_state_dict(average_checkpoints(filenames, device=device))
else:
if params.iter > 0:
filenames = find_checkpoints(
params.exp_dir, iteration=-params.iter
)[: params.avg + 1]
if len(filenames) == 0:
raise ValueError(
f"No checkpoints found for"
f" --iter {params.iter}, --avg {params.avg}"
)
elif len(filenames) < params.avg + 1:
raise ValueError(
f"Not enough checkpoints ({len(filenames)}) found for"
f" --iter {params.iter}, --avg {params.avg}"
)
filename_start = filenames[-1]
filename_end = filenames[0]
logging.info(
"Calculating the averaged model over iteration checkpoints"
f" from {filename_start} (excluded) to {filename_end}"
)
model.to(device)
model.load_state_dict(
average_checkpoints_with_averaged_model(
filename_start=filename_start,
filename_end=filename_end,
device=device,
)
)
else:
assert params.avg > 0
start = params.epoch - params.avg
assert start >= 1
filename_start = f"{params.exp_dir}/epoch-{start}.pt"
filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
logging.info(
f"Calculating the averaged model over epoch range from "
f"{start} (excluded) to {params.epoch}"
)
model.to(device)
model.load_state_dict(
average_checkpoints_with_averaged_model(
filename_start=filename_start,
filename_end=filename_end,
device=device,
)
)
model.to(device)
model.eval()
if params.decoding_method == "fast_beam_search":
decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
else:
decoding_graph = None
num_param = sum([p.numel() for p in model.parameters()])
logging.info(f"Number of model parameters: {num_param}")
librispeech = LibriSpeechAsrDataModule(args)
test_clean_cuts = librispeech.test_clean_cuts()
test_other_cuts = librispeech.test_other_cuts()
test_clean_dl = librispeech.test_dataloaders(test_clean_cuts)
test_other_dl = librispeech.test_dataloaders(test_other_cuts)
test_sets = ["test-clean", "test-other"]
test_dl = [test_clean_dl, test_other_dl]
for test_set, test_dl in zip(test_sets, test_dl):
results_dict = decode_dataset(
dl=test_dl,
params=params,
model=model,
sp=sp,
decoding_graph=decoding_graph,
)
save_results(
params=params,
test_set_name=test_set,
results_dict=results_dict,
)
logging.info("Done!")
if __name__ == "__main__":
main()

1
egs/librispeech/ASR/pruned_transducer_stateless4b/decoder.py
View File

@ -1 +0,0 @@
../pruned_transducer_stateless2/decoder.py

275
egs/librispeech/ASR/pruned_transducer_stateless4b/diagonalize.py
View File

@ -1,275 +0,0 @@
#!/usr/bin/env python3
# Copyright (c) 2022 Xiaomi Corporation (author: 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 copy
import math
import warnings
from typing import Optional, Tuple
import logging
import torch
from torch import Tensor, nn
import torch.distributed as dist
# some utilities for diagnalizing models (rotating their parameters matrices
# so that large and small parameter values are separated as much as possible).
def _get_normalized_covar(x: Tensor) -> Tensor:
"""
Returns a covariance matrix normalized to have trace==dim, equal to
matmul(x , x.t()) times a constant.
Args:
x: a matrix of shape (i, j)
Returns: a covariance matrix of shape (i, i), equal to matmul(x, x.t())
"""
covar = torch.matmul(x, x.t())
return covar * (x.shape[0] / (covar.trace() + 1.0e-20))
@torch.no_grad()
def get_diag_covar_in(m: nn.Module) -> Tensor:
"""
Returns a covariance matrix that shows, in the input space of
this module, which direction parameter matrices vary in.
"""
if isinstance(m, nn.Linear):
return _get_normalized_covar(m.weight.t());
elif isinstance(m, nn.Conv1d) or isinstance(m, nn.Conv2d):
# m.weight is of size (out_channels, in_channels, kernel_size)
# or (out_channels, in_channels, kernel_dim0, kernel_dim1)
# assert here that groups == 1
w = m.weight
assert m.groups == 1
out_channels = w.shape[0]
in_channels = w.shape[1]
w = w.reshape(out_channels, in_channels, -1)
w = w.permute(1, 0, 2) # (in_channels, out_channels, kernel_size)
w = w.reshape(in_channels, -1)
return _get_normalized_covar(w) # (in_channels, in_channels)
elif isinstance(m, nn.Sequential):
return get_diag_covar_in(m[0], t)
else:
# some modules have this function; if not, at this point, it is an error.
return m.get_diag_covar_in()
@torch.no_grad()
def get_diag_covar_out(m: nn.Module) -> Tensor:
"""
Returns a covariance matrix that shows, in the output space of
this module, which direction parameter matrices vary in.
"""
if isinstance(m, nn.Linear):
return _get_normalized_covar(m.weight);
elif isinstance(m, nn.Conv1d) or isinstance(m, nn.Conv2d):
# m.weight is of size (out_channels, in_channels, kernel_size)
# or (out_channels, in_channels, kernel_dim0, kernel_dim1)
# assert here that groups == 1
w = m.weight
assert m.groups == 1
out_channels = w.shape[0]
in_channels = w.shape[1]
w = w.reshape(out_channels, -1)
return _get_normalized_covar(w) # (out_channels, out_channels)
w = w.permute(1, 0, 2) # (in_channels, out_channels, kernel_size)
w = w.reshape(in_channels, -1)
return _get_normalized_covar(x) # (in_channels, in_channels)
elif isinstance(m, nn.Sequential):
return get_diag_covar_out(m[-1])
else:
# some modules have this function; if not, at this point, it is an error.
return m.get_diag_covar_out()
@torch.no_grad()
def get_diag_covar_inout(m: nn.Module) -> Tensor:
"""
Returns a covariance matrix that shows, in the input and
output space of this module, which are assumed to be the
same (e.g if it is a module intended to be added to a residual/
bypass connection),
which direction parameter matrices vary in.
"""
if isinstance(m, nn.Sequential):
# this is only correct if it's a Sequential of non-residual modules.
return get_diag_covar_in(m[0]) + get_diag_covar_out(m[-1])
else:
# some modules have this function; if not, at this point, it is an error.
return m.get_diag_covar_inout()
@torch.no_grad()
def apply_transformation_in(m: nn.Module, t: Tensor) -> None:
"""
Applies this transformation matrix on the input space of this module.
Args:
m: module to transform on the input space
t: transformation matrix, indexed (new_dim_in, old_dim_in)
"""
if isinstance(m, nn.Linear):
m.weight[:] = torch.matmul(m.weight, t.t())
elif isinstance(m, nn.Conv1d) or isinstance(m, nn.Conv2d):
# m.weight is of size (out_channels, in_channels, kernel_size)
# or (out_channels, in_channels, kernel_dim0, kernel_dim1)
# assert here that groups == 1
w = m.weight
assert m.groups == 1
out_channels = w.shape[0]
in_channels = w.shape[1]
w = w.reshape(out_channels, in_channels, -1)
w = w.permute(1, 0, 2) # (in_channels, out_channels, kernel_size)
w = w.reshape(in_channels, -1)
w = torch.matmul(t, w).reshape(in_channels, out_channels, -1) # (in_channels, out_channels, kernel_size)
w = w.permute(1, 0, 2) # (out_channels, in_channels, kernel_size)
w = w.reshape(m.weight.shape) # (out_channels, in_channels, [1 or 2 kernel dims])
m.weight[:] = w
elif isinstance(m, nn.Sequential):
apply_transformation_in(m[0], t)
else:
# some modules have this function; if not, at this point, it is an error.
m.apply_transformation_in(t)
@torch.no_grad()
def apply_transformation_out(m: nn.Module, t: Tensor) -> None:
"""
Applies this transformation matrix on the output space of this module.
Args:
m: module to transform on the input space
t: transformation matrix, indexed (new_dim_out, old_dim_out)
"""
if isinstance(m, nn.Linear):
m.weight[:] = torch.matmul(t, m.weight)
if m.bias is not None:
m.bias[:] = torch.matmul(t, m.bias)
elif isinstance(m, nn.Conv1d) or isinstance(m, nn.Conv2d):
# m.weight is of size (out_channels, in_channels, kernel_size)
# or (out_channels, in_channels, kernel_dim0, kernel_dim1)
# assert here that groups == 1
w = m.weight
assert m.groups == 1
out_channels = w.shape[0]
in_channels = w.shape[1]
w = w.reshape(out_channels, -1)
w = torch.matmul(t, w)
w = w.reshape(m.weight.shape) # (out_channels, in_channels, [1 or 2 kernel dims])
m.weight[:] = w
if m.bias is not None:
m.bias[:] = torch.matmul(t, m.bias)
elif isinstance(m, nn.Sequential):
apply_transformation_out(m[-1], t)
else:
# some modules have this function; if not, at this point, it is an error.
m.apply_transformation_out(t)
@torch.no_grad()
def apply_transformation_inout(m: nn.Module, t: Tensor) -> None:
if isinstance(m, nn.Sequential):
apply_transformation_in(m, t)
apply_transformation_out(m, t)
else:
# some modules have this function; if not, at this point, it is an error.
m.apply_transformation_inout(t)
def get_transformation(cov: Tensor) -> Tensor:
"""
Returns a covariance-diagonalizing transformation that diagonalizes
the covariance matrix that is passed in.
Args: cov, of shape (dim0, dim0).
Returns: a transformation indexed (new_dim0, old_dim0), i.e. of
shape dim0 by dim0 but 1st index is the newly created indexes.
"""
old_diag_stddev = cov.diag().var().sqrt().item()
l, U = cov.symeig(eigenvectors=True)
new_diag_stddev = l.var().sqrt().item()
logging.info(f"Variance of diag of param-var changed from {old_diag_stddev:.3e} "
f"to {new_diag_stddev:.3e}, max diag elem changed from {cov.diag().max().item():.2e} to {l[-1].item():.2e}")
return U.t() # U.t() is indexed (new_dim, old_dim)
class OrthogonalTransformation(nn.Module):
def __init__(self, num_channels: int):
super(OrthogonalTransformation, self).__init__()
# `weight` is indexed (channel_out, channel_in)
self.register_buffer('weight', torch.eye(num_channels)) # not a parameter
self.register_buffer('feats_cov', torch.eye(num_channels)) # not a parameter
self.step = 0 # just to co-ordinate updating feats_cov every 10 batches; not saved to disk.
self.beta = 0.9 # affects how long we remember the stats. not super critical.
def forward(self, x: Tensor):
"""
Args:
x: Tensor of shape (*, num_channel)
Returns:
Tensor of shape (*, num_channels), x multiplied by orthogonal matrix.
"""
x = torch.matmul(x, self.weight.t())
if self.step % 10 == 0 and self.train():
with torch.no_grad():
# store covariance after input transform.
# Update covariance stats every 10 batches (in training mode)
f = x.reshape(-1, x.shape[-1])
f = f * (f.shape[0] ** -0.5)
cov = torch.matmul(f.t(), f) # channel_dim by channel_dim
self.feats_cov.mul_(self.beta).add_(cov, alpha=(1-self.beta))
self.step += 1
return x
@torch.no_grad()
def apply_transformation_in(self, t: Tensor) -> None:
"""
Rotate only the input feature space with an orthogonal matrix.
t is indexed (new_channel_dim, old_channel_dim)
"""
# note, self.weight is indexed (channel_out, channel_in), interpreted
# initially as (channel_out, old_channel_in), which we multiply
# by t.t() which is (old_channel_in, new_channel_in)
self.weight[:] = torch.matmul(self.weight, t.t())
@torch.no_grad()
def apply_transformation_out(self, t: Tensor) -> None:
"""
Rotate only the output feature space with an orthogonal matrix.
t is indexed (new_channel_dim, old_channel_dim)
We don't bother updating the covariance stats; they will decay.
"""
# note, self.weight is indexed (channel_out, channel_in), interpreted
# initially as (old_channel_out, old_channe), which we pre-multiply
# by t which is (new_channel_out, old_channel_out)
self.weight[:] = torch.matmul(t, self.weight)
self.feats_cov[:] = torch.matmul(t, torch.matmul(self.feats_cov, t.t()))
@torch.no_grad()
def get_transformation_out(self) -> Tensor:
# see also get_transformation() above for notes on this.
cov = 0.5 * (self.feats_cov + self.feats_cov.t()) # make sure symmetric
t = get_transformation(cov)
if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 0:
# make sure all processes in the process group share the same version of `t`.
# this would usually be the case, but if on this batch we modified self.feats_cov,
# it won't be the same among all processes because DDP synchronizes buffers at the
# beginning, not the end, of the forward().
logging.info("Broadcastint transformation")
dist.broadcast(t)
return t

1
egs/librispeech/ASR/pruned_transducer_stateless4b/encoder_interface.py
View File

@ -1 +0,0 @@
../pruned_transducer_stateless2/encoder_interface.py

1
egs/librispeech/ASR/pruned_transducer_stateless4b/export.py
View File

@ -1 +0,0 @@
../pruned_transducer_stateless2/export.py

1
egs/librispeech/ASR/pruned_transducer_stateless4b/joiner.py
View File

@ -1 +0,0 @@
../pruned_transducer_stateless2/joiner.py

270
egs/librispeech/ASR/pruned_transducer_stateless4b/model.py
View File

@ -1,270 +0,0 @@
# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang, Wei Kang)
#
# 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 k2
import torch
from torch import Tensor
import torch.nn as nn
from encoder_interface import EncoderInterface
from scaling import ScaledLinear
from diagonalize import get_diag_covar_in, apply_transformation_in, get_transformation, apply_transformation_in, apply_transformation_out
from icefall.utils import add_sos
class Transducer(nn.Module):
"""It implements https://arxiv.org/pdf/1211.3711.pdf
"Sequence Transduction with Recurrent Neural Networks"
"""
def __init__(
self,
encoder: EncoderInterface,
decoder: nn.Module,
joiner: nn.Module,
encoder_dim: int,
decoder_dim: int,
joiner_dim: int,
vocab_size: int,
):
"""
Args:
encoder:
It is the transcription network in the paper. Its accepts
two inputs: `x` of (N, T, encoder_dim) and `x_lens` of shape (N,).
It returns two tensors: `logits` of shape (N, T, encoder_dm) and
`logit_lens` of shape (N,).
decoder:
It is the prediction network in the paper. Its input shape
is (N, U) and its output shape is (N, U, decoder_dim).
It should contain one attribute: `blank_id`.
joiner:
It has two inputs with shapes: (N, T, encoder_dim) and (N, U, decoder_dim).
Its output shape is (N, T, U, vocab_size). Note that its output contains
unnormalized probs, i.e., not processed by log-softmax.
"""
super().__init__()
assert isinstance(encoder, EncoderInterface), type(encoder)
assert hasattr(decoder, "blank_id")
self.encoder = encoder
self.decoder = decoder
self.joiner = joiner
self.simple_am_proj = ScaledLinear(
encoder_dim, vocab_size, initial_speed=0.5
)
self.simple_lm_proj = ScaledLinear(decoder_dim, vocab_size)
def forward(
self,
x: torch.Tensor,
x_lens: torch.Tensor,
y: k2.RaggedTensor,
prune_range: int = 5,
am_scale: float = 0.0,
lm_scale: float = 0.0,
warmup: float = 1.0,
) -> torch.Tensor:
"""
Args:
x:
A 3-D tensor of shape (N, T, C).
x_lens:
A 1-D tensor of shape (N,). It contains the number of frames in `x`
before padding.
y:
A ragged tensor with 2 axes [utt][label]. It contains labels of each
utterance.
prune_range:
The prune range for rnnt loss, it means how many symbols(context)
we are considering for each frame to compute the loss.
am_scale:
The scale to smooth the loss with am (output of encoder network)
part
lm_scale:
The scale to smooth the loss with lm (output of predictor network)
part
warmup:
A value warmup >= 0 that determines which modules are active, values
warmup > 1 "are fully warmed up" and all modules will be active.
Returns:
Return the transducer loss.
Note:
Regarding am_scale & lm_scale, it will make the loss-function one of
the form:
lm_scale * lm_probs + am_scale * am_probs +
(1-lm_scale-am_scale) * combined_probs
"""
assert x.ndim == 3, x.shape
assert x_lens.ndim == 1, x_lens.shape
assert y.num_axes == 2, y.num_axes
assert x.size(0) == x_lens.size(0) == y.dim0
encoder_out, x_lens = self.encoder(x, x_lens, warmup=warmup)
assert torch.all(x_lens > 0)
# Now for the decoder, i.e., the prediction network
row_splits = y.shape.row_splits(1)
y_lens = row_splits[1:] - row_splits[:-1]
blank_id = self.decoder.blank_id
sos_y = add_sos(y, sos_id=blank_id)
# sos_y_padded: [B, S + 1], start with SOS.
sos_y_padded = sos_y.pad(mode="constant", padding_value=blank_id)
# decoder_out: [B, S + 1, decoder_dim]
decoder_out = self.decoder(sos_y_padded)
# Note: y does not start with SOS
# y_padded : [B, S]
y_padded = y.pad(mode="constant", padding_value=0)
y_padded = y_padded.to(torch.int64)
boundary = torch.zeros(
(x.size(0), 4), dtype=torch.int64, device=x.device
)
boundary[:, 2] = y_lens
boundary[:, 3] = x_lens
lm = self.simple_lm_proj(decoder_out)
am = self.simple_am_proj(encoder_out)
with torch.cuda.amp.autocast(enabled=False):
simple_loss, (px_grad, py_grad) = k2.rnnt_loss_smoothed(
lm=lm.float(),
am=am.float(),
symbols=y_padded,
termination_symbol=blank_id,
lm_only_scale=lm_scale,
am_only_scale=am_scale,
boundary=boundary,
reduction="sum",
return_grad=True,
)
# ranges : [B, T, prune_range]
ranges = k2.get_rnnt_prune_ranges(
px_grad=px_grad,
py_grad=py_grad,
boundary=boundary,
s_range=prune_range,
)
# am_pruned : [B, T, prune_range, encoder_dim]
# lm_pruned : [B, T, prune_range, decoder_dim]
am_pruned, lm_pruned = k2.do_rnnt_pruning(
am=self.joiner.encoder_proj(encoder_out),
lm=self.joiner.decoder_proj(decoder_out),
ranges=ranges,
)
# logits : [B, T, prune_range, vocab_size]
# project_input=False since we applied the decoder's input projections
# prior to do_rnnt_pruning (this is an optimization for speed).
logits = self.joiner(am_pruned, lm_pruned, project_input=False)
with torch.cuda.amp.autocast(enabled=False):
pruned_loss = k2.rnnt_loss_pruned(
logits=logits.float(),
symbols=y_padded,
ranges=ranges,
termination_symbol=blank_id,
boundary=boundary,
reduction="sum",
)
return (simple_loss, pruned_loss)
def diagonalize(self) -> None:
cur_transform = None
for l in self.encoder.encoder.layers:
if cur_transform is not None:
l.apply_transformation_in(cur_transform)
cur_transform = l.get_transformation_out()
l.apply_transformation_out(cur_transform)
self.encoder.diagonalize() # diagonalizes self_attn layers, this is
# purely internal to the self_attn layers.
apply_transformation_in(self.simple_am_proj, cur_transform)
apply_transformation_in(self.joiner.encoder_proj, cur_transform)
def _test_model():
import logging
logging.getLogger().setLevel(logging.INFO)
from conformer import Conformer
from joiner import Joiner
from decoder import Decoder
feature_dim = 40
attention_dim = 256
encoder_dim = 512
decoder_dim = 513
joiner_dim = 514
vocab_size = 1000
encoder = Conformer(num_features=40,
subsampling_factor=4,
d_model=encoder_dim,
nhead=4,
dim_feedforward=512,
num_encoder_layers=4)
decoder = Decoder(
vocab_size=600,
decoder_dim=decoder_dim,
blank_id=0,
context_size=2)
joiner = Joiner(
encoder_dim=encoder_dim,
decoder_dim=decoder_dim,
joiner_dim=joiner_dim,
vocab_size=vocab_size)
model = Transducer(encoder=encoder,
decoder=decoder,
joiner=joiner,
encoder_dim=encoder_dim,
decoder_dim=decoder_dim,
joiner_dim=joiner_dim,
vocab_size=vocab_size)
batch_size = 5
seq_len = 50
feats = torch.randn(batch_size, seq_len, feature_dim)
x_lens = torch.full((batch_size,), seq_len, dtype=torch.int64)
y = k2.ragged.create_ragged_tensor(torch.arange(5, dtype=torch.int32).reshape(1,5).expand(batch_size,5))
model.eval() # eval mode so it's not random.
(simple_loss1, pruned_loss1) = model(feats, x_lens, y)
model.diagonalize()
(simple_loss2, pruned_loss2) = model(feats, x_lens, y)
print(f"simple_loss1 = {simple_loss1.mean().item()}, simple_loss2 = {simple_loss2.mean().item()}")
print(f"pruned_loss1 = {pruned_loss1.mean().item()}, pruned_loss2 = {pruned_loss2.mean().item()}")
model.diagonalize()
if __name__ == '__main__':
_test_model()

821
egs/librispeech/ASR/pruned_transducer_stateless4b/optim.py
View File

@ -1,821 +0,0 @@
# 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.
from typing import List, Optional, Union, Tuple, List
import torch
from torch import Tensor
from torch.optim import Optimizer
def _product(*args):
ans = args[0]
for x in args[1:]:
ans = ans * x
return ans
def _sum(*args):
ans = args[0]
for x in args[1:]:
ans = ans + x
return ans
def _mean_like(x: Tensor, size) -> Tensor:
"""
Returns a mean of x with the shape `size`, summing
over all dimensions in which x.shape[dim] != 1 and size[dim] == 1.
Args:
x: Tensor to compute mean over some dims.
size: a sequence of integers or torch.Size, with
len(size) == len(x.size), that
broadcasts with x and elementwise <= x.size.
"""
ndim = x.ndim
dims_to_sum = [i for i in range(ndim) if x.shape[i] != size[i]]
return x.mean(dims_to_sum, keepdim=True)
def _sum_like(x: Tensor, size) -> Tensor:
"""
Returns a sum of values of x with the shape `size`, summing
over all dimensions in which x.shape[dim] != 1 and size[dim] == 1.
Args:
x: Tensor to compute mean over some dims.
size: a sequence of integers or torch.Size, with
len(size) == len(x.size), that
broadcasts with x and elementwise <= x.size.
"""
ndim = x.ndim
dims_to_sum = [i for i in range(ndim) if x.shape[i] != size[i]]
return x.sum(dims_to_sum, keepdim=True)
def _get_factor_shapes(x_shape: List) -> List[List]:
"""
Returns a list of the shapes of factors of a tensor with shape
x_shape, e.g.:
_get_factor_shapes([3,4]) = [[3,1], [1,4]]
_get_factor_shapes([5]) = [[1]]
"""
base_shape = [1] * len(x_shape)
shapes = []
numel = 1
for n in x_shape:
numel *= n
for dim,size in enumerate(x_shape):
if size != 0 and size != numel:
shape = list(base_shape)
shape[dim] = size
shapes.append(shape)
if len(shapes) == 0:
shapes.append(base_shape)
return shapes
def _update_factorization(x: Tensor, x_factorized: Tensor,
speed: float, eps: float) -> None:
"""
This function is like _update_factors but instead of operating
on the factors directly it operates on their product.
Args:
x: the Tensor to be normalized; for purposes of understanding
what this function does, assume that x is drawn from some fixed
distribution. Require x.ndim >= 2 and that at least 2 dims
have size != 1.
x_factorized: a Tensor with the same shape as x, but it is
a product of factors, one factor for each axis i that has
x.shape[i] != 1 and x.shape[i] != x.numel(); if no axes
remain, we have a single scalar factor. This function
modifies x_factorized to be closer to a model of the stddevs
of elements of x.
speed: update speed, e.g 0.02 (consider this as 1-beta_2
where beta_2 is the 2nd beta from Adam). Must satisfy
speed > 0 and speed * x.ndim < 1.
eps: small value intended to prevent zero values in
x_factorized; you can think of it as a floor on elements
of x_factorized.
"""
x_norm_var = (x / x_factorized) ** 2 + eps*eps
shapes = _get_factor_shapes(list(x.shape))
factors = []
for shape in shapes:
# elements of this_mean will be >1.0 if the squares of elements in x are
# larger than predicted by x_factorized; and <1.0 otherwise.
this_mean = _mean_like(x_norm_var, shape)
f = ((1.0 - speed) + speed * this_mean)
factors.append(f)
x_factorized *= _product(*factors)
def _get_factor_grads(x: Tensor, x_grad: Tensor) -> List[Tensor]:
"""
Get gradients w.r.t. the factors in a factorized representation of x.
Consider the example where
x_norm.shape == (3, 4), x_factor1.shape == (3, 1), x_factor2.shape == (1, 4),
and:
x = x_norm * x_factor1.exp() * x_factor2.exp()
Given x and x_grad, this function will return the (x_factor1_grad, x_factor2_grad),
i.e. the loss-function derivatives w.r.t. x_factor1 and x_factor2.
The specific values of the factors are not needed here; their
product is enough.
"""
shapes = _get_factor_shapes(list(x.shape))
prod = x * x_grad
return [ _sum_like(prod, shape) for shape in shapes ]
def _init_factors(x: Tensor, eps: float = 1.0e-04) -> Tuple[Tensor]:
"""
Initialize some factors which we will use to normalize the variance of x.
For example: for a Tensor of shape (3, 4) we will return factors of shapes
(3, 1) and (1, 4) having the property that x / (factor1 * factor2) has quite
close to unit variance when indexed on either dimension, i.e. that
(x[i,:]**2).mean() is close to 1 and (x[:,j]**2).mean() is close to 1.
"""
assert x.numel() > 1
factors = []
x2 = x**2 + eps*eps
for d in range(x.ndim):
if x.shape[d] != x.numel() and x.shape[d] != 1:
shape = [1] * x.ndim
shape[d] = x.shape[d]
factor = _mean_like(x2, shape)
x2 = x2 / factor
factors.append(factor ** 0.5)
# We need at least one factor, to capture the scalar magnitude.
if len(factors) == 0:
factors.append(x2.mean(dim=tuple(range(x.ndim)),
keepdims=True) ** 0.5)
return factors
def _init_factorization(x: Tensor, eps: float = 1.0e-04) -> Tensor:
return _product(*_init_factors(x, eps))
class Abel(Optimizer):
"""
Implements the Abel algorithm. You can think of this as a modified version
of the Adam algorithm in which the parameters have an implicit factorization
that is maintained in a special way. (Actually this can also be considered
as a refinement of the "Eve" algorithm, see below).
Suppose we have a parameter `x` of
shape (40, 50). Imagine that we factorize x as follows:
x = x_norm * x_factor1.exp() * x_factor2.exp()
where (x_norm, x_factor1, x_factor2) are of shapes: (40, 50), (40, 1), (1, 50),
and all 3 of x_norm, x_factor1 and x_factor2 are learned by Adam. Also imagine
that we frequently reconstruct x and then refactorize it in such a way that
the standard deviation of elements of x, when restricted to any specific rows or column,
equals `target_rms == 0.1`. So it's a bit like we constrain the elements of x_norm
to have a specific standard deviation on any row or column, and instead learn the
magnitudes of the rows and columns in log-space as x_factor1 and x_factor2.
But this is all done inside the optimizer. To simplify things a bit when updating
the factorization, we assume in the explanation below and in most of the code that
target_rms equals 1.0; it then simply becomes a factor of `target_rms` in the
learning rate when we update x (but not the part of the step that comes from
the udpate of x_factor1 and x_factor2).
Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-08).
target_rms (float, optional): target root-mean-square value of
x_norm in factorization (conceptually). Actually this now just becomes
a factor in the learning rate, we may remove it at some point.
.. _Adam\: A Method for Stochastic Optimization:
https://arxiv.org/abs/1412.6980
.. _On the Convergence of Adam and Beyond:
https://openreview.net/forum?id=ryQu7f-RZ
"""
def __init__(
self,
params,
lr=1e-3,
betas=(0.9, 0.98),
eps=1e-04,
target_rms=0.1,
):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= betas[0] < 1.0:
raise ValueError(
"Invalid beta parameter at index 0: {}".format(betas[0])
)
if not 0.0 <= betas[1] < 1.0:
raise ValueError(
"Invalid beta parameter at index 1: {}".format(betas[1])
)
if not 0 < target_rms <= 1.0:
raise ValueError("Invalid target_rms value: {}".format(target_rms))
defaults = dict(
lr=lr,
betas=betas,
eps=eps,
target_rms=target_rms,
)
super(Abel, self).__init__(params, defaults)
def __setstate__(self, state):
super(Abel, 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 in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
# Perform optimization step
grad = p.grad
if grad.is_sparse:
raise RuntimeError(
"AdamW does not support sparse gradients"
)
state = self.state[p]
eps = group["eps"]
# State initialization
if len(state) == 0:
state["step"] = 0
# The change in parameter p from one step to the next; this
# is the moving average of the steps before taking momentum
# (beta) into account. We implement momentum this way,
# instead of the "exp_avg" method as in Adam, to save a
# little memory and complexity because this way the momentum for the
# various fctors doesn't have to be kept track of
# separately.
state["delta"] = torch.zeros_like(
p, memory_format=torch.preserve_format
)
# Exponential moving average of squared gradient values,
# w.r.t. x.
state["exp_avg_sq"] = torch.zeros_like(
p, memory_format=torch.preserve_format
)
if p.numel() > 1:
# exponential moving average of gradients w.r.t
# the factors x_factor1, x_factor2 etc.
state["factors_exp_avg_sq"] = [
torch.zeros(*shape, dtype=p.dtype, device=p.device)
for shape in _get_factor_shapes(list(p.shape))
]
# we call _init_factorization inside step().
state["factorization"] = torch.zeros_like(p)
exp_avg_sq = state["exp_avg_sq"]
delta, exp_avg_sq = state["delta"], state["exp_avg_sq"]
lr = group["lr"]
target_rms = group["target_rms"]
eps = group["eps"]
beta1, beta2 = group["betas"]
state["step"] += 1
step = state["step"]
# we don't bother with bias_correction1, this only affects the first few
# steps and anyway the slower-than-it-should be update probably helps stability.
bias_correction2 = 1 - beta2 ** step
# keep this_exp_avg_sq updated as an average variance of gradients.
def update_exp_avg_sq(this_grad, this_exp_avg_sq):
this_exp_avg_sq.mul_(beta2).addcmul_(this_grad, this_grad,
value=1 - beta2)
update_exp_avg_sq(grad, exp_avg_sq)
if p.numel() == 1:
# simpler way of setting this_delta; this is going to be equivalent
# to standard Adam.
denom = (exp_avg_sq/bias_correction2 + eps*eps).sqrt()
this_delta = grad / denom
# eventually will remove this special case, this is to deal
# with our model that has scaling factors.
lr_eff = lr * (bias_correction2 ** 0.5)
else:
# update that includes factorization-related stuff..
factors_exp_avg_sq, factorization = state["factors_exp_avg_sq"], state["factorization"]
# factor_grads: List[Tensor], len == num-factors; the gradients
# w.r.t. the factors of x (x_factor1, x_factor2..)
factor_grads = _get_factor_grads(p, grad)
if step < 10 or step % 10 == 1:
# note, the first step number we'll see is 1.
# update the factorization this only every 10 steps, to save time.
if step % 1000 == 1:
# Periodically refresh the factorization; this is
# out of concern that after a large number of
# multiplications, roundoff could cause it to drift
# significantly away from being a product of
# independent factors, one per dimension.
factorization[:] = _init_factorization(p, eps)
_update_factorization(p, factorization,
speed=0.1,
eps=eps)
factors_sum = None
numel = p.numel()
for g, e in zip(factor_grads, factors_exp_avg_sq):
this_numel = numel / g.numel()
update_exp_avg_sq(g, e)
# this_numel**0.5 will turn into this_numel**-0.25 when we sqrt() and
# divide by denom. This becomes a factor in the learning rate.
# The technique naturally makes us learn faster by (this_numel**0.5)
# in any of the parameter directions corresponding to rows or columns,
# and because we don't want to be to aggressive and risk instability or
# excess param noise, we reduce this to (this_numel**0.25) by dividing
# the effective LR by this_numel**0.25.
this_denom = (e*(this_numel**0.5)/bias_correction2 + eps*eps).sqrt()
assert g.shape == this_denom.shape
factor_delta = g / this_denom
factors_sum = (factor_delta if factors_sum is None
else factors_sum + factor_delta)
# at this point, factors_sum either has the same shape as p, or,
# if p has only one non-trivial dimension, it has a shape equal
# to (1,) * p.ndim. After multiplying by the learning rate
# and p, it will represent the change in parameter that arises
# from updating the factors (before considering momentum!)
denom = (exp_avg_sq/bias_correction2 +
eps*eps).sqrt()
# `grad * factorization / denom` represents the change in parameters x, only considering
# the change from `x_norm` in the factorization, and
# before multiplying by the learning rate or taking into account
# momentum. We compute this in the original space of `x`,
# simply because this is more convenient, but due to invariants
# of the Adam update this would be exactly the same if computed
# on `x_norm`, except for some small differences relating to
# `eps`.
# `p * factors_sum` is the contribution from changes in x_factor1
# and x_factor2: again, before taking into account the learning
# rate or momentum.
this_delta = ((grad * factorization / denom) + p * factors_sum)
lr_eff = lr * (bias_correction2 ** 0.5) / target_rms
delta.mul_(beta1).add_(this_delta, alpha=-lr_eff*(1-beta1))
p.add_(delta)
return loss
@torch.no_grad()
def reset(self):
"""
Resets parts of the state of the optimizer. This is to be called after
refactorization/orthogonalization of parameters. At these points we
discard any momentum because it is no longer accurate/relevant;
we reconstruct the factorization of x (i.e. x_factor1 and x_factor2
in the intro to this class); and we modify exp_avg_sq to
0.5 (exp_avg_sq + exp_avg_sq.mean()); this ensures that none of the
exp_avg_sq values will be substantially too small and prevents any
too-fast updates.
"""
for s in self.state.values():
s["delta"].zero_() # zero out momentum
s["step"] = 0 # will cause state["factorization"] to be reinitialized
s["exp_avg_sq"].zero_()
if "factors_exp_avg_sq" in s:
for e in s["factors_exp_avg_sq"]:
e.zero_()
class Eve(Optimizer):
"""
Implements Eve algorithm. This is a modified version of AdamW with a special
way of setting the weight-decay / shrinkage-factor, which is designed to make the
rms of the parameters approach a particular target_rms (default: 0.1). This is
for use with networks with 'scaled' versions of modules (see scaling.py), which
will be close to invariant to the absolute scale on the parameter matrix.
The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
Eve is unpublished so far.
Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
weight_decay (float, optional): weight decay coefficient (default: 3e-4;
this value means that the weight would decay significantly after
about 3k minibatches. Is not multiplied by learning rate, but
is conditional on RMS-value of parameter being > target_rms.
target_rms (float, optional): target root-mean-square value of
parameters, if they fall below this we will stop applying weight decay.
.. _Adam\: A Method for Stochastic Optimization:
https://arxiv.org/abs/1412.6980
.. _Decoupled Weight Decay Regularization:
https://arxiv.org/abs/1711.05101
.. _On the Convergence of Adam and Beyond:
https://openreview.net/forum?id=ryQu7f-RZ
"""
def __init__(
self,
params,
lr=1e-3,
betas=(0.9, 0.98),
eps=1e-8,
weight_decay=1e-3,
target_rms=0.1,
):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= betas[0] < 1.0:
raise ValueError(
"Invalid beta parameter at index 0: {}".format(betas[0])
)
if not 0.0 <= betas[1] < 1.0:
raise ValueError(
"Invalid beta parameter at index 1: {}".format(betas[1])
)
if not 0 <= weight_decay <= 0.1:
raise ValueError(
"Invalid weight_decay value: {}".format(weight_decay)
)
if not 0 < target_rms <= 10.0:
raise ValueError("Invalid target_rms value: {}".format(target_rms))
defaults = dict(
lr=lr,
betas=betas,
eps=eps,
weight_decay=weight_decay,
target_rms=target_rms,
)
super(Eve, self).__init__(params, defaults)
def __setstate__(self, state):
super(Eve, 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 in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
# Perform optimization step
grad = p.grad
if grad.is_sparse:
raise RuntimeError(
"AdamW does not support sparse gradients"
)
state = self.state[p]
# State initialization
if len(state) == 0:
state["step"] = 0
# Exponential moving average of gradient values
state["exp_avg"] = torch.zeros_like(
p, memory_format=torch.preserve_format
)
# Exponential moving average of squared gradient values
state["exp_avg_sq"] = torch.zeros_like(
p, memory_format=torch.preserve_format
)
exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
beta1, beta2 = group["betas"]
state["step"] += 1
bias_correction1 = 1 - beta1 ** state["step"]
bias_correction2 = 1 - beta2 ** state["step"]
# Decay the first and second moment running average coefficient
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
denom = (exp_avg_sq.sqrt() * (bias_correction2 ** -0.5)).add_(
group["eps"]
)
step_size = group["lr"] / bias_correction1
target_rms = group["target_rms"]
weight_decay = group["weight_decay"]
if p.numel() > 1:
# avoid applying this weight-decay on "scaling factors"
# (which are scalar).
is_above_target_rms = p.norm() > (
target_rms * (p.numel() ** 0.5)
)
p.mul_(1 - (weight_decay * is_above_target_rms))
p.addcdiv_(exp_avg, denom, value=-step_size)
return loss
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("initial_lr", group["lr"])
self.base_lrs = [
group["initial_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 {
"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`.
"""
self.__dict__.update(state_dict)
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:
print(
f"Epoch={self.epoch}, batch={self.batch}: adjusting learning rate"
f" of group {group} to {lr:.4e}."
)
class Eden(LRScheduler):
"""
Eden scheduler.
lr = initial_lr * (((batch**2 + lr_batches**2) / lr_batches**2) ** -0.25 *
(((epoch**2 + lr_epochs**2) / lr_epochs**2) ** -0.25))
E.g. suggest initial-lr = 0.003 (passed to optimizer).
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],
verbose: bool = False,
):
super(Eden, self).__init__(optimizer, verbose)
self.lr_batches = lr_batches
self.lr_epochs = lr_epochs
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
)
return [x * factor for x in self.base_lrs]
def _test_eden():
m = torch.nn.Linear(100, 100)
optim = Eve(m.parameters(), lr=0.003)
scheduler = Eden(optim, lr_batches=30, 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()
print("last lr = ", scheduler.get_last_lr())
print("state dict = ", scheduler.state_dict())
def _test_abel():
import timeit
from scaling import ScaledLinear
E = 100
B = 4
T = 2
print("in test_abel")
device = torch.device('cuda')
dtype = torch.float32
# 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()
for iter in [0,1]:
Linear = torch.nn.Linear if iter == 0 else ScaledLinear
m = torch.nn.Sequential(Linear(E, 200),
torch.nn.ReLU(),
Linear(200, 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) ]
if iter == 0: optim = Abel(m.parameters(), lr=0.003)
else: optim = Eve(m.parameters(), lr=0.003)
scheduler = Eden(optim, lr_batches=300, lr_epochs=20, verbose=False)
start = timeit.default_timer()
for epoch in range(150):
scheduler.step_epoch()
if isinstance(optim, Abel) and epoch % 5 == 0: # this is every 100 minibatches.
optim.reset() # just test that calling reset() doesn't
# break things, you wouldn't call it every
# epoch like this.
for n, (x,y) in enumerate(train_pairs):
y_out = m(x)
loss = ((y_out - y)**2).mean() * 100.0
if n == 0 and epoch % 10 == 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())
print(f"Iter {iter}, epoch {epoch}, batch {n}, loss {loss.item()}, norms={norm1,norm1b,norm2,norm2b}") # scales={scale1,scale1b,scale2,scale2b}
loss.log().backward()
optim.step()
optim.zero_grad()
scheduler.step_batch()
stop = timeit.default_timer()
print(f"Iter={iter}, Time taken: {stop - start}")
print("last lr = ", scheduler.get_last_lr())
#print("state dict = ", scheduler.state_dict())
#print("optim state_dict = ", optim.state_dict())
print("input_magnitudes = ", input_magnitudes)
print("output_magnitudes = ", output_magnitudes)
def stddev(x):
return ((x-x.mean())**2).mean().sqrt()
print("Un-normalized input col magnitudes log-stddev: ", stddev((m[0].weight**2).sum(dim=0).sqrt().log()))
print("Normalized input col magnitudes log-stddev: ", stddev(((m[0].weight**2).sum(dim=0).sqrt() * input_magnitudes).log()))
print("Un-normalized 0-output row magnitudes log-stddev: ", stddev((m[0].weight**2).sum(dim=1).sqrt().log()))
print("Un-normalized 2-input col magnitudes log-stddev: ", stddev((m[2].weight**2).sum(dim=0).sqrt().log()))
print("Un-normalized 2-output row magnitudes log-stddev: ", stddev((m[2].weight**2).sum(dim=1).sqrt().log()))
print("Normalized output row magnitudes log-stddev: ", stddev(((m[2].weight**2).sum(dim=1).sqrt() / output_magnitudes).log()))
if __name__ == "__main__":
_test_abel()
#_test_eden()

971
egs/librispeech/ASR/pruned_transducer_stateless4b/optim2.py
View File

@ -1,971 +0,0 @@
# 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.
from typing import List, Optional, Union, Tuple, List
import torch
from torch import Tensor
from torch.optim import Optimizer
def _product(*args):
ans = args[0]
for x in args[1:]:
ans = ans * x
return ans
def _sum(*args):
ans = args[0]
for x in args[1:]:
ans = ans + x
return ans
def _mean_like(x: Tensor, size) -> Tensor:
"""
Returns a mean of x with the shape `size`, summing
over all dimensions in which x.shape[dim] != 1 and size[dim] == 1.
Args:
x: Tensor to compute mean over some dims.
size: a sequence of integers or torch.Size, with
len(size) == len(x.size), that
broadcasts with x and elementwise <= x.size.
"""
ndim = x.ndim
dims_to_sum = [i for i in range(ndim) if x.shape[i] != size[i]]
return x.mean(dims_to_sum, keepdim=True)
def _sum_like(x: Tensor, size) -> Tensor:
"""
Returns a sum of values of x with the shape `size`, summing
over all dimensions in which x.shape[dim] != 1 and size[dim] == 1.
Args:
x: Tensor to compute mean over some dims.
size: a sequence of integers or torch.Size, with
len(size) == len(x.size), that
broadcasts with x and elementwise <= x.size.
"""
ndim = x.ndim
dims_to_sum = [i for i in range(ndim) if x.shape[i] != size[i]]
return x.sum(dims_to_sum, keepdim=True)
def _get_factor_shapes(x_shape: List) -> List[List]:
"""
Returns a list of the shapes of factors of a tensor with shape
x_shape, e.g.:
_get_factor_shapes([3,4]) = [[3,1], [1,4]]
_get_factor_shapes([5]) = [[1]]
"""
base_shape = [1] * len(x_shape)
shapes = []
numel = 1
for n in x_shape:
numel *= n
for dim,size in enumerate(x_shape):
if size != 0 and size != numel:
shape = list(base_shape)
shape[dim] = size
shapes.append(shape)
if len(shapes) == 0:
shapes.append(base_shape)
return shapes
def _update_factorization(x: Tensor, x_factorized: Tensor,
speed: float, eps: float) -> None:
"""
This function is like _update_factors but instead of operating
on the factors directly it operates on their product.
Args:
x: the Tensor to be normalized; for purposes of understanding
what this function does, assume that x is drawn from some fixed
distribution. Require x.ndim >= 2 and that at least 2 dims
have size != 1.
x_factorized: a Tensor with the same shape as x, but it is
a product of factors, one factor for each axis i that has
x.shape[i] != 1 and x.shape[i] != x.numel(); if no axes
remain, we have a single scalar factor. This function
modifies x_factorized to be closer to a model of the stddevs
of elements of x.
speed: update speed, e.g 0.02 (consider this as 1-beta_2
where beta_2 is the 2nd beta from Adam). Must satisfy
speed > 0 and speed * x.ndim < 1.
eps: small value intended to prevent zero values in
x_factorized; you can think of it as a floor on elements
of x_factorized.
"""
x_norm_var = (x / x_factorized) ** 2 + eps*eps
shapes = _get_factor_shapes(list(x.shape))
factors = []
for shape in shapes:
# elements of this_mean will be >1.0 if the squares of elements in x are
# larger than predicted by x_factorized; and <1.0 otherwise.
this_mean = _mean_like(x_norm_var, shape)
f = ((1.0 - speed) + speed * this_mean)
factors.append(f)
x_factorized *= _product(*factors)
def _get_factor_grads(x: Tensor, x_grad: Tensor) -> List[Tensor]:
"""
Get gradients w.r.t. the factors in a factorized representation of x.
Consider the example where
x_norm.shape == (3, 4), x_factor1.shape == (3, 1), x_factor2.shape == (1, 4),
and:
x = x_norm * x_factor1.exp() * x_factor2.exp()
Given x and x_grad, this function will return the (x_factor1_grad, x_factor2_grad),
i.e. the loss-function derivatives w.r.t. x_factor1 and x_factor2.
The specific values of the factors are not needed here; their
product is enough.
"""
shapes = _get_factor_shapes(list(x.shape))
prod = x * x_grad
return [ _sum_like(prod, shape) for shape in shapes ]
def _init_factors(x: Tensor, eps: float = 1.0e-04) -> Tuple[Tensor]:
"""
Initialize some factors which we will use to normalize the variance of x.
For example: for a Tensor of shape (3, 4) we will return factors of shapes
(3, 1) and (1, 4) having the property that x / (factor1 * factor2) has quite
close to unit variance when indexed on either dimension, i.e. that
(x[i,:]**2).mean() is close to 1 and (x[:,j]**2).mean() is close to 1.
"""
assert x.numel() > 1
factors = []
x2 = x**2 + eps*eps
for d in range(x.ndim):
if x.shape[d] != x.numel() and x.shape[d] != 1:
shape = [1] * x.ndim
shape[d] = x.shape[d]
factor = _mean_like(x2, shape)
x2 = x2 / factor
factors.append(factor ** 0.5)
# We need at least one factor, to capture the scalar magnitude.
if len(factors) == 0:
factors.append(x2.mean(dim=tuple(range(x.ndim)),
keepdims=True) ** 0.5)
return factors
def _init_factorization(x: Tensor, eps: float = 1.0e-04) -> Tensor:
return _product(*_init_factors(x, eps))
class Abel(Optimizer):
"""
Implements the Abel algorithm. You can think of this as a modified version
of the Adam algorithm in which the parameters have an implicit factorization
that is maintained in a special way. (Actually this can also be considered
as a refinement of the "Eve" algorithm, see below).
Suppose we have a parameter `x` of
shape (40, 50). Imagine that we factorize x as follows:
x = x_norm * x_factor1.exp() * x_factor2.exp()
where (x_norm, x_factor1, x_factor2) are of shapes: (40, 50), (40, 1), (1, 50),
and all 3 of x_norm, x_factor1 and x_factor2 are learned by Adam. Also imagine
that we frequently reconstruct x and then refactorize it in such a way that
the standard deviation of elements of x, when restricted to any specific rows or column,
equals `target_rms == 0.1`. So it's a bit like we constrain the elements of x_norm
to have a specific standard deviation on any row or column, and instead learn the
magnitudes of the rows and columns in log-space as x_factor1 and x_factor2.
But this is all done inside the optimizer. To simplify things a bit when updating
the factorization, we assume in the explanation below and in most of the code that
target_rms equals 1.0; it then simply becomes a factor of `target_rms` in the
learning rate when we update x (but not the part of the step that comes from
the udpate of x_factor1 and x_factor2).
Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-08).
target_rms (float, optional): target root-mean-square value of
x_norm in factorization (conceptually). Actually this now just becomes
a factor in the learning rate, we may remove it at some point.
.. _Adam\: A Method for Stochastic Optimization:
https://arxiv.org/abs/1412.6980
.. _On the Convergence of Adam and Beyond:
https://openreview.net/forum?id=ryQu7f-RZ
"""
def __init__(
self,
params,
lr=1e-3,
betas=(0.9, 0.98),
eps=1e-04,
target_rms=0.1,
):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= betas[0] < 1.0:
raise ValueError(
"Invalid beta parameter at index 0: {}".format(betas[0])
)
if not 0.0 <= betas[1] < 1.0:
raise ValueError(
"Invalid beta parameter at index 1: {}".format(betas[1])
)
if not 0 < target_rms <= 1.0:
raise ValueError("Invalid target_rms value: {}".format(target_rms))
defaults = dict(
lr=lr,
betas=betas,
eps=eps,
target_rms=target_rms,
)
super(Abel, self).__init__(params, defaults)
def __setstate__(self, state):
super(Abel, 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 in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
# Perform optimization step
grad = p.grad
if grad.is_sparse:
raise RuntimeError(
"AdamW does not support sparse gradients"
)
state = self.state[p]
eps = group["eps"]
# State initialization
if len(state) == 0:
state["step"] = 0
# The change in parameter p from one step to the next; this
# is the moving average of the steps before taking momentum
# (beta) into account. We implement momentum this way,
# instead of the "exp_avg" method as in Adam, to save a
# little memory and complexity because this way the momentum for the
# various fctors doesn't have to be kept track of
# separately.
state["delta"] = torch.zeros_like(
p, memory_format=torch.preserve_format
)
# Exponential moving average of squared gradient values,
# w.r.t. x.
state["exp_avg_sq"] = torch.zeros_like(
p, memory_format=torch.preserve_format
)
if p.numel() > 1:
# exponential moving average of gradients w.r.t
# the factors x_factor1, x_factor2 etc.
state["factors_exp_avg_sq"] = [
torch.zeros(*shape, dtype=p.dtype, device=p.device)
for shape in _get_factor_shapes(list(p.shape))
]
# we call _init_factorization inside step().
state["factorization"] = torch.zeros_like(p)
exp_avg_sq = state["exp_avg_sq"]
delta, exp_avg_sq = state["delta"], state["exp_avg_sq"]
lr = group["lr"]
target_rms = group["target_rms"]
eps = group["eps"]
beta1, beta2 = group["betas"]
state["step"] += 1
step = state["step"]
# we don't bother with bias_correction1, this only affects the first few
# steps and anyway the slower-than-it-should be update probably helps stability.
bias_correction2 = 1 - beta2 ** step
# keep this_exp_avg_sq updated as an average variance of gradients.
def update_exp_avg_sq(this_grad, this_exp_avg_sq):
this_exp_avg_sq.mul_(beta2).addcmul_(this_grad, this_grad,
value=1 - beta2)
update_exp_avg_sq(grad, exp_avg_sq)
if p.numel() == 1:
# simpler way of setting this_delta; this is going to be equivalent
# to standard Adam.
denom = (exp_avg_sq/bias_correction2 + eps*eps).sqrt()
this_delta = grad / denom
# eventually will remove this special case, this is to deal
# with our model that has scaling factors.
lr_eff = lr * (bias_correction2 ** 0.5)
else:
# update that includes factorization-related stuff..
factors_exp_avg_sq, factorization = state["factors_exp_avg_sq"], state["factorization"]
# factor_grads: List[Tensor], len == num-factors; the gradients
# w.r.t. the factors of x (x_factor1, x_factor2..)
factor_grads = _get_factor_grads(p, grad)
if step < 10 or step % 10 == 1:
# note, the first step number we'll see is 1.
# update the factorization this only every 10 steps, to save time.
if step % 1000 == 1:
# Periodically refresh the factorization; this is
# out of concern that after a large number of
# multiplications, roundoff could cause it to drift
# significantly away from being a product of
# independent factors, one per dimension.
factorization[:] = _init_factorization(p, eps)
_update_factorization(p, factorization,
speed=0.1,
eps=eps)
factors_sum = None
numel = p.numel()
for g, e in zip(factor_grads, factors_exp_avg_sq):
this_numel = numel / g.numel()
update_exp_avg_sq(g, e)
# this_numel**0.5 will turn into this_numel**-0.25 when we sqrt() and
# divide by denom. This becomes a factor in the learning rate.
# The technique naturally makes us learn faster by (this_numel**0.5)
# in any of the parameter directions corresponding to rows or columns,
# and because we don't want to be to aggressive and risk instability or
# excess param noise, we reduce this to (this_numel**0.25) by dividing
# the effective LR by this_numel**0.25.
this_denom = (e*(this_numel**0.5)/bias_correction2 + eps*eps).sqrt()
assert g.shape == this_denom.shape
factor_delta = g / this_denom
factors_sum = (factor_delta if factors_sum is None
else factors_sum + factor_delta)
# at this point, factors_sum either has the same shape as p, or,
# if p has only one non-trivial dimension, it has a shape equal
# to (1,) * p.ndim. After multiplying by the learning rate
# and p, it will represent the change in parameter that arises
# from updating the factors (before considering momentum!)
denom = (exp_avg_sq/bias_correction2 +
eps*eps).sqrt()
# `grad * factorization / denom` represents the change in parameters x, only considering
# the change from `x_norm` in the factorization, and
# before multiplying by the learning rate or taking into account
# momentum. We compute this in the original space of `x`,
# simply because this is more convenient, but due to invariants
# of the Adam update this would be exactly the same if computed
# on `x_norm`, except for some small differences relating to
# `eps`.
# `p * factors_sum` is the contribution from changes in x_factor1
# and x_factor2: again, before taking into account the learning
# rate or momentum.
this_delta = ((grad * factorization / denom) + p * factors_sum)
lr_eff = lr * (bias_correction2 ** 0.5) / target_rms
delta.mul_(beta1).add_(this_delta, alpha=-lr_eff*(1-beta1))
p.add_(delta)
return loss
@torch.no_grad()
def reset(self):
"""
Resets parts of the state of the optimizer. This is to be called after
refactorization/orthogonalization of parameters. At these points we
discard any momentum because it is no longer accurate/relevant;
we reconstruct the factorization of x (i.e. x_factor1 and x_factor2
in the intro to this class); and we modify exp_avg_sq to
0.5 (exp_avg_sq + exp_avg_sq.mean()); this ensures that none of the
exp_avg_sq values will be substantially too small and prevents any
too-fast updates.
"""
for s in self.state.values():
s["delta"].zero_() # zero out momentum
s["step"] = 0 # will cause state["factorization"] to be reinitialized
s["exp_avg_sq"].zero_()
if "factors_exp_avg_sq" in s:
for e in s["factors_exp_avg_sq"]:
e.zero_()
class Eve(Optimizer):
"""
Implements Eve algorithm. This is a modified version of AdamW with a special
way of setting the weight-decay / shrinkage-factor, which is designed to make the
rms of the parameters approach a particular target_rms (default: 0.1). This is
for use with networks with 'scaled' versions of modules (see scaling.py), which
will be close to invariant to the absolute scale on the parameter matrix.
The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
Eve is unpublished so far.
Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
weight_decay (float, optional): weight decay coefficient (default: 3e-4;
this value means that the weight would decay significantly after
about 3k minibatches. Is not multiplied by learning rate, but
is conditional on RMS-value of parameter being > target_rms.
target_rms (float, optional): target root-mean-square value of
parameters, if they fall below this we will stop applying weight decay.
.. _Adam\: A Method for Stochastic Optimization:
https://arxiv.org/abs/1412.6980
.. _Decoupled Weight Decay Regularization:
https://arxiv.org/abs/1711.05101
.. _On the Convergence of Adam and Beyond:
https://openreview.net/forum?id=ryQu7f-RZ
"""
def __init__(
self,
params,
lr=1e-3,
betas=(0.9, 0.98),
eps=1e-8,
weight_decay=1e-3,
target_rms=0.1,
):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= betas[0] < 1.0:
raise ValueError(
"Invalid beta parameter at index 0: {}".format(betas[0])
)
if not 0.0 <= betas[1] < 1.0:
raise ValueError(
"Invalid beta parameter at index 1: {}".format(betas[1])
)
if not 0 <= weight_decay <= 0.1:
raise ValueError(
"Invalid weight_decay value: {}".format(weight_decay)
)
if not 0 < target_rms <= 10.0:
raise ValueError("Invalid target_rms value: {}".format(target_rms))
defaults = dict(
lr=lr,
betas=betas,
eps=eps,
weight_decay=weight_decay,
target_rms=target_rms,
)
super(Eve, self).__init__(params, defaults)
def __setstate__(self, state):
super(Eve, 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 in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
# Perform optimization step
grad = p.grad
if grad.is_sparse:
raise RuntimeError(
"AdamW does not support sparse gradients"
)
state = self.state[p]
# State initialization
if len(state) == 0:
state["step"] = 0
# Exponential moving average of gradient values
state["exp_avg"] = torch.zeros_like(
p, memory_format=torch.preserve_format
)
# Exponential moving average of squared gradient values
state["exp_avg_sq"] = torch.zeros_like(
p, memory_format=torch.preserve_format
)
exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
beta1, beta2 = group["betas"]
# forget bias_correction1. We use normal momentum. exp_avg really
# just stores the moving-average gradient step.
bias_correction2 = 1 - beta2 ** self._step_for_bias_correction(state["step"])
grad = self._change_coordinates(grad, state, forward=True)
# Decay the second moment running average coefficient
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
denom = (exp_avg_sq.sqrt()).add_(group["eps"])
step_size = group["lr"] # forget bias_correction1
target_rms = group["target_rms"]
weight_decay = group["weight_decay"]
this_delta = grad / denom
this_delta = self._change_coordinates(this_delta, state, forward=False)
# treat/repurpose exp_avg as moving-average of `this_delta`
exp_avg.mul_(beta1).add_(this_delta, alpha=-step_size*(1-beta1)*(bias_correction2 ** 0.5))
if p.numel() > 1:
# avoid applying this weight-decay on "scaling factors"
# (which are scalar).
is_above_target_rms = p.norm() > (
target_rms * (p.numel() ** 0.5)
)
p.mul_(1 - (weight_decay * is_above_target_rms))
p.add_(exp_avg)
state["step"] += 1
return loss
def _change_coordinates(self,
grad: Tensor,
state: dict,
forward: bool):
"""
Possibly performs some magnitude-preserving changes of co-ordinates on `grad`
"""
ndim = grad.ndim
numel = grad.numel()
for dim in range(grad.ndim):
# see for each dim in turn whether we want to perform any changes in co-ordinates,
# or store any stats.
size = grad.shape[dim]
if size <= 3 or size > 1024 or size == numel:
# we don't do any such co-ordinate changes in dims that are too small (no point)
# or large (too slow)
continue
if dim == 0:
# FOR NOW: don't do such co-ordinate changes on output
# dims, which will generally be dimension zero. We can revisit this later.
continue
grad = self._change_coordinates_for_dim(grad, state, dim, forward)
return grad
def _change_coordinates_for_dim(self,
grad: Tensor,
state: dict,
dim: int,
forward: bool,
orig_dim: int = -1):
assert grad.ndim > 1
if not (grad.ndim == 2 and dim == 1):
# normalize the format and recurse.
dims_order = []
rev_dims_order = []
ndim = grad.ndim
for i in range(dim):
dims_order.append(i)
rev_dims_order.append(i)
rev_dims_order.append(ndim-1)
for i in range(dim+1, ndim):
dims_order.append(i)
rev_dims_order.append(i-1)
dims_order.append(dim)
# e.g. ndim=4, dim=1, dims_order=(0,2,3,1), rev_dims_order=(0,3,1,2)
new_grad = grad.permute(*dims_order)
assert new_grad.shape[-1] == grad.shape[dim]
new_shape = new_grad.shape
new_grad = new_grad.reshape(-1, new_grad.shape[-1])
new_grad = self._change_coordinates_for_dim(new_grad, state, 1, forward,
orig_dim=dim)
return new_grad.reshape(new_shape).permute(*rev_dims_order)
# OK: grad.ndim == 2 and dim == 1
size = grad.shape[1]
if orig_dim == -1:
orig_dim = dim
step = state["step"]
must_store_stats, must_zero_stats = self._must_store_stats(step)
if must_store_stats:
# store stats for 20 iters preceding when we estimate the
# transform.
stats_name = f"cov_{orig_dim}"
if not stats_name in state:
state[stats_name] = torch.zeros(size, size, dtype=grad.dtype,
device=grad.device)
cov = state[stats_name]
if must_zero_stats:
#print("zero")
cov.zero_()
cov += torch.matmul(grad.t(), grad) * (1/grad.shape[0])
transform_name = f"t_{orig_dim}"
if transform_name not in state:
# make sure they're all allocated at the start, may be better for
# fragmention and debugging reasons (although may not).
state[transform_name] = torch.eye(size, device=grad.device,
dtype=grad.dtype)
must_estimate_transform = self._must_estimate_transform(step)
if must_estimate_transform:
stats_name = f"cov_{orig_dim}"
cov = state[stats_name]
l, U = cov.symeig(eigenvectors=True)
#print("estimate, l = ", l[::20])
t = U.t() # t is indexed (new_dim, old_dim) a.k.a. (transformed_dim, real_dim)
state[transform_name][:] = t
state["exp_avg_sq"].zero_() # zero exp-avg-sq stats, we'll restart these from scratch.
t = state[transform_name]
if forward:
return torch.matmul(grad, t.t())
else:
return torch.matmul(grad, t)
def _must_store_stats(self, step: int) -> Tuple[bool, bool]:
"""
Returns (must_store_stats, must_zero_stats), where
must_store_stats: bool, is True for the 20 steps preceding
a step on which we will estimate the transform.
must_zero_stats: bool, is True for only the 1st of the 20
steps mentioned above.
"""
if step < 4000:
if step < 2000:
return (step % 500 >= 480, step % 500 == 480)
else:
return (step % 1000 >= 980, step % 1000 == 980)
else:
return (step % 2000 >= 1980, step % 2000 == 1980)
def _must_estimate_transform(self, step: int) -> bool:
"""
Returns True if we will re-estimate the transform on this step
"""
if step < 4000:
if step < 2000:
return step % 500 == 0 and step > 0
else:
return step % 1000 == 0 and step > 0
else:
return step % 2000 == 0
def _step_for_bias_correction(self, step: int) -> bool:
"""
Return a modified step for bias-correction (actually just to give us the
right denominator for the exp-avg-sq stats).. this needs to take into
account how long it has been since we re-estimated the transforms, because
we zero the exp-avg-sq stats at those times.
The + 1 is for the "current step". We post-increment the step so it
starts from 0.
"""
if step < 4000:
if step < 2000:
return step % 500 + 1
else:
return step % 1000 + 1
else:
return step % 2000 + 1
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("initial_lr", group["lr"])
self.base_lrs = [
group["initial_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 {
"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`.
"""
self.__dict__.update(state_dict)
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:
print(
f"Epoch={self.epoch}, batch={self.batch}: adjusting learning rate"
f" of group {group} to {lr:.4e}."
)
class Eden(LRScheduler):
"""
Eden scheduler.
lr = initial_lr * (((batch**2 + lr_batches**2) / lr_batches**2) ** -0.25 *
(((epoch**2 + lr_epochs**2) / lr_epochs**2) ** -0.25))
E.g. suggest initial-lr = 0.003 (passed to optimizer).
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],
verbose: bool = False,
):
super(Eden, self).__init__(optimizer, verbose)
self.lr_batches = lr_batches
self.lr_epochs = lr_epochs
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
)
return [x * factor for x in self.base_lrs]
def _test_eden():
m = torch.nn.Linear(100, 100)
optim = Eve(m.parameters(), lr=0.003)
scheduler = Eden(optim, lr_batches=30, 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()
print("last lr = ", scheduler.get_last_lr())
print("state dict = ", scheduler.state_dict())
def _test_abel():
import timeit
from scaling import ScaledLinear
E = 100
B = 4
T = 2
print("in test_abel")
device = torch.device('cuda')
dtype = torch.float32
# 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()
for iter in [1,0]:
Linear = torch.nn.Linear if iter == 0 else ScaledLinear
m = torch.nn.Sequential(Linear(E, 200),
torch.nn.ReLU(),
Linear(200, 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) ]
if iter == 0: optim = Abel(m.parameters(), lr=0.003)
else: optim = Eve(m.parameters(), lr=0.003)
scheduler = Eden(optim, lr_batches=300, lr_epochs=20, verbose=False)
start = timeit.default_timer()
for epoch in range(150):
scheduler.step_epoch()
if isinstance(optim, Abel) and epoch % 5 == 0: # this is every 100 minibatches.
optim.reset() # just test that calling reset() doesn't
# break things, you wouldn't call it every
# epoch like this.
for n, (x,y) in enumerate(train_pairs):
y_out = m(x)
loss = ((y_out - y)**2).mean() * 100.0
if n == 0 and epoch % 10 == 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())
print(f"Iter {iter}, epoch {epoch}, batch {n}, loss {loss.item()}, norms={norm1,norm1b,norm2,norm2b}") # scales={scale1,scale1b,scale2,scale2b}
loss.log().backward()
optim.step()
optim.zero_grad()
scheduler.step_batch()
stop = timeit.default_timer()
print(f"Iter={iter}, Time taken: {stop - start}")
print("last lr = ", scheduler.get_last_lr())
#print("state dict = ", scheduler.state_dict())
#print("optim state_dict = ", optim.state_dict())
print("input_magnitudes = ", input_magnitudes)
print("output_magnitudes = ", output_magnitudes)
def stddev(x):
return ((x-x.mean())**2).mean().sqrt()
print("Un-normalized input col magnitudes log-stddev: ", stddev((m[0].weight**2).sum(dim=0).sqrt().log()))
print("Normalized input col magnitudes log-stddev: ", stddev(((m[0].weight**2).sum(dim=0).sqrt() * input_magnitudes).log()))
print("Un-normalized 0-output row magnitudes log-stddev: ", stddev((m[0].weight**2).sum(dim=1).sqrt().log()))
print("Un-normalized 2-input col magnitudes log-stddev: ", stddev((m[2].weight**2).sum(dim=0).sqrt().log()))
print("Un-normalized 2-output row magnitudes log-stddev: ", stddev((m[2].weight**2).sum(dim=1).sqrt().log()))
print("Normalized output row magnitudes log-stddev: ", stddev(((m[2].weight**2).sum(dim=1).sqrt() / output_magnitudes).log()))
if __name__ == "__main__":
_test_abel()
#_test_eden()

1
egs/librispeech/ASR/pruned_transducer_stateless4b/scaling.py
View File

@ -1 +0,0 @@
../pruned_transducer_stateless2/scaling.py

1058
egs/librispeech/ASR/pruned_transducer_stateless4b/train.py
View File

File diff suppressed because it is too large Load Diff