mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-09-07 08:04:18 +00:00
202 lines
7.9 KiB
Python
202 lines
7.9 KiB
Python
#!/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], 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)
|