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

Implement weight_scale, set weight_scale=10

Browse Source
This commit is contained in:
Daniel Povey 2023-05-18 15:48:14 +08:00
parent f6c7392430
commit d2c198c072
1 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
egs/libriheavy/LM/zipformer1/subformer.py
View File

@ -805,13 +805,17 @@ class LearnedDownsamplingModule(nn.Module):
downsampling_factor: factor to downsample by, e.g. 2 or 4. There is no downsampling_factor: factor to downsample by, e.g. 2 or 4. There is no
fundamental reason why this has to be an integer, but we make it so fundamental reason why this has to be an integer, but we make it so
anyway. anyway.
weight_scale: constant scaling factor on the weights, introduced to make fp16 training
more stable by reducing gradient magnitudes.
""" """
def __init__(self, def __init__(self,
embed_dim: int, embed_dim: int,
downsampling_factor: int): downsampling_factor: int,
weight_scale: float = 10.0):
super().__init__() super().__init__()
self.weight_scale = weight_scale
self.to_scores = nn.Linear(embed_dim, 1, bias=False) self.to_scores = nn.Linear(embed_dim, 1, bias=False)
# score_balancer is just to keep the magnitudes of the scores in # score_balancer is just to keep the magnitudes of the scores in
# a fixed range and keep them balanced around zero, to stop # a fixed range and keep them balanced around zero, to stop
@ -855,7 +859,7 @@ class LearnedDownsamplingModule(nn.Module):
sscores, indexes = scores.sort(dim=-1, descending=True) sscores, indexes = scores.sort(dim=-1, descending=True)
weights = sscores.clamp(min=0.0, max=1.0) weights = sscores.clamp(min=0.0, max=self.weight_scale)
weights = self.copy_weights1(weights) weights = self.copy_weights1(weights)
if self.training: if self.training:
@ -987,6 +991,9 @@ class LearnedDownsamplingModule(nn.Module):
# unsqueeze at position 1 so the extra cost relates to the source position. # unsqueeze at position 1 so the extra cost relates to the source position.
attn_offset = attn_offset + weights.clamp(min=eps).log().unsqueeze(1) attn_offset = attn_offset + weights.clamp(min=eps).log().unsqueeze(1)
if self.weight_scale != 1.0:
attn_offset = attn_offset - math.log(self.weight_scale)
return attn_offset return attn_offset