From f941991331b8927ebd49b05d52984d488e8c1faa Mon Sep 17 00:00:00 2001 From: Daniel Povey Date: Mon, 10 Oct 2022 13:38:36 +0800 Subject: [PATCH] Fix bug in choosing layers to drop --- .../ASR/pruned_transducer_stateless7/conformer.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/conformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/conformer.py index add243d55..9a78c6838 100644 --- a/egs/librispeech/ASR/pruned_transducer_stateless7/conformer.py +++ b/egs/librispeech/ASR/pruned_transducer_stateless7/conformer.py @@ -471,16 +471,16 @@ class ConformerEncoder(nn.Module): # with their specified probs while reaching this exact target. num_to_drop = int(tot) + int(shared_rng.random() < (tot - int(tot))) - layers = list(range(num_layers)) independent_rng.shuffle(layers) # go through the shuffled layers until we get the required number of samples. - for layer in itertools.cycle(layers): - if independent_rng.random() < layerdrop_probs[layer]: - ans.add(layer) - if len(ans) == num_to_drop: - break + if num_to_drop > 0: + for layer in itertools.cycle(layers): + if independent_rng.random() < layerdrop_probs[layer]: + ans.add(layer) + if len(ans) == num_to_drop: + break if shared_rng.random() < 0.005 or __name__ == "__main__": logging.info(f"warmup_begin={warmup_begin}, warmup_end={warmup_end}, warmup_count={warmup_count:.1f}, num_to_drop={num_to_drop}, layers_to_drop={ans}") return ans