mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-09-07 08:04:18 +00:00
Draft towards 2nd orthogonalization
This commit is contained in:
Daniel Povey
parent
8aeaf1421a
commit
67f916e599
204
egs/librispeech/ASR/pruned_transducer_stateless4b/diagonalize.py
Normal file
204
egs/librispeech/ASR/pruned_transducer_stateless4b/diagonalize.py
Normal file
@ -0,0 +1,204 @@
|
||||
#!/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
|
||||
|
||||
# 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])
|
||||
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])
|
||||
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])
|
||||
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.
|
||||
"""
|
||||
cov = get_normalized_covar(args[0])
|
||||
for a in args[1:]:
|
||||
cov += get_normalized_covar(a)
|
||||
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}")
|
||||
return U.t() # U.t() is indexed (new_dim, old_dim)
|
||||
@ -1 +0,0 @@
|
||||
../pruned_transducer_stateless2/model.py
|
||||
201
egs/librispeech/ASR/pruned_transducer_stateless4b/model.py
Normal file
201
egs/librispeech/ASR/pruned_transducer_stateless4b/model.py
Normal file
@ -0,0 +1,201 @@
|
||||
# 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
|
||||
|
||||
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 get_diag_covar_in(self) -> Tensor:
|
||||
return (get_diag_covar_in(self.simple_am_proj) +
|
||||
get_diag_covar_in(joiner.encoder_proj) +
|
||||
self.encoder.get_diag_covar_out())
|
||||
Loading…
x
Reference in New Issue
Block a user