mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-12-11 06:55:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'scaled_adam_exp466' into scaled_adam_exp472.

Browse Source 
Below is a more complete list of the changes I am making, although some of
these may be counted in the last

  numbers XXX below correspond to branches numbered scaled_adam_expXXX.
    - from 412/413 (cherry-picked): dropout for attention in attention_squeeze and nonlin_attention modules,
      but simplified this a little to use the same dropout schedule and drop them out all together
      also have all 3 submodules use separate heads.
    - from 460->461, which is in the history of 464, revert the part about balancing output out attention_squeeze module.
    - merge from 462->467, about using TanSwish not tanh.
    - merge 462->465, remove whitening in self-attention module
    - merge the part of 465->466  that was about diagnostics (name in Whiten module)
This commit is contained in:
Daniel Povey 2022-11-23 14:39:08 +08:00
commit 1d0252d420
3 changed files with 22 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py
View File

@ -561,11 +561,13 @@ class WhiteningPenaltyFunction(torch.autograd.Function):
x: Tensor,
num_groups: int,
whitening_limit: float,
grad_scale: float) -> Tensor:
grad_scale: float,
name: Optional[str]) -> Tensor:
ctx.save_for_backward(x)
ctx.num_groups = num_groups
ctx.whitening_limit = whitening_limit
ctx.grad_scale = grad_scale
ctx.name = name
return x
@staticmethod
@ -580,7 +582,7 @@ class WhiteningPenaltyFunction(torch.autograd.Function):
metric = _whitening_metric(x_detached, ctx.num_groups)
if random.random() < 0.005 or __name__ == "__main__":
logging.info(f"Whitening: num_groups={ctx.num_groups}, num_channels={x_orig.shape[-1]}, "
logging.info(f"Whitening: name={ctx.name}, num_groups={ctx.num_groups}, num_channels={x_orig.shape[-1]}, "
f"metric={metric.item():.2f} vs. limit={ctx.whitening_limit}")
(metric - ctx.whitening_limit).relu().backward()
@ -588,7 +590,7 @@ class WhiteningPenaltyFunction(torch.autograd.Function):
scale = ctx.grad_scale * (x_grad.to(torch.float32).norm() /
(penalty_grad.norm() + 1.0e-20))
penalty_grad = penalty_grad * scale
return x_grad + penalty_grad.to(x_grad.dtype), None, None, None
return x_grad + penalty_grad.to(x_grad.dtype), None, None, None, None
@ -630,7 +632,7 @@ class Whiten(nn.Module):
(self.min_prob, self.max_prob) = prob
assert 0 < self.min_prob < self.max_prob <= 1
self.prob = self.max_prob
self.name = None # will be set in training loop
self.grad_scale = grad_scale
def forward(self,
@ -666,7 +668,8 @@ class Whiten(nn.Module):
return WhiteningPenaltyFunction.apply(x,
self.num_groups,
self.whitening_limit,
self.grad_scale)
self.grad_scale,
self.name)
class WithLoss(torch.autograd.Function):

28
egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
View File

@ -367,7 +367,7 @@ class ZipformerEncoderLayer(nn.Module):
# to work correctly.
layer_skip_rate: FloatLike = ScheduledFloat((0.0, 0.5), (4000.0, 0.05), default=0),
dynamic_skip_rate: FloatLike = ScheduledFloat((0.0, 0.2), (4000.0, 0.0), default=0),
squeeze_const_attention_rate: FloatLike = ScheduledFloat((0.0, 0.5), (4000.0, 0.05), default=0),
const_attention_rate: FloatLike = ScheduledFloat((0.0, 0.25), (4000.0, 0.025), default=0),
bypass_min: FloatLike = ScheduledFloat((0.0, 0.75), (20000.0, 0.25), default=0),
bypass_max: FloatLike = 1.0,
) -> None:
@ -382,7 +382,7 @@ class ZipformerEncoderLayer(nn.Module):
# ever becoming zero.
self.bypass_min = copy.deepcopy(bypass_min)
self.bypass_max = copy.deepcopy(bypass_max)
self.squeeze_const_attention_rate = copy.deepcopy(squeeze_const_attention_rate)
self.const_attention_rate = copy.deepcopy(const_attention_rate)
self.self_attn_weights = RelPositionMultiheadAttentionWeights(
embed_dim, pos_dim=pos_dim, num_heads=num_heads,
@ -480,27 +480,23 @@ class ZipformerEncoderLayer(nn.Module):
key_padding_mask=src_key_padding_mask,
)
squeeze_weights = attn_weights[1:2]
if random.random() < float(self.squeeze_const_attention_rate):
# this form of dropout makes the attention-weights used for the
# squeeze-excite modules constant wherever they are not masked. The intention
# is to encourage these modules to do something similar to an averaging-over-time
# operation.
squeeze_weights = (squeeze_weights > 0.0).to(squeeze_weights.dtype)
# make sure they sum to 1 over the last axis.
squeeze_weights = squeeze_weights * (1.0 / squeeze_weights.sum(dim=-1, keepdim=True))
first_attn_weights = attn_weights[0:3]
if random.random() < float(self.const_attention_rate):
# Make attention weights constant. The intention is to
# encourage these modules to do something similar to an
# averaging-over-time operation.
first_attn_weights = (first_attn_weights > 0.0).to(first_attn_weights.dtype)
first_attn_weights = first_attn_weights * (1.0 / first_attn_weights.sum(dim=-1, keepdim=True))
if torch.jit.is_scripting() or use_self_attn:
src = src + self.nonlin_attention_module(src,
attn_weights[0:1])
first_attn_weights[0:1])
src = src + self.feed_forward1(src)
# pooling module
if torch.jit.is_scripting() or use_self_attn:
src = src + self.attention_squeeze1(src, squeeze_weights)
src = src + self.attention_squeeze1(src, first_attn_weights[1:2])
if torch.jit.is_scripting() or use_self_attn:
src = src + self.self_attn(
@ -513,7 +509,7 @@ class ZipformerEncoderLayer(nn.Module):
# pooling module
if torch.jit.is_scripting() or use_self_attn:
src = src + self.attention_squeeze2(src, squeeze_weights)
src = src + self.attention_squeeze2(src, first_attn_weights[2:3])
src = self.norm_final(self.balancer(src))