File size: 4,578 Bytes
7934b29 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
#
# 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 pytest
import torch
from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder
class TestStochasticDepth:
"""Testing stochastic depth functionality."""
def test_stochastic_depth_model_creation(self):
"""Testing basic model creation and the drop probs are correctly assigned."""
n_layers = 4
model = ConformerEncoder(feat_in=10, n_layers=n_layers, d_model=4, feat_out=8)
# checking that by default SD is disabled
assert model.layer_drop_probs == [0.0] * n_layers
# linear mode
for drop_prob in [0.3, 0.5, 0.9]:
for start_layer in [1, 3]:
model = ConformerEncoder(
feat_in=10,
n_layers=n_layers,
d_model=4,
feat_out=8,
stochastic_depth_drop_prob=drop_prob,
stochastic_depth_start_layer=start_layer,
)
L = n_layers - start_layer
assert model.layer_drop_probs == [0.0] * start_layer + [drop_prob * l / L for l in range(1, L + 1)]
# uniform mode
for drop_prob in [0.3, 0.5, 0.9]:
model = ConformerEncoder(
feat_in=10,
n_layers=n_layers,
d_model=4,
feat_out=8,
stochastic_depth_drop_prob=drop_prob,
stochastic_depth_mode="uniform",
stochastic_depth_start_layer=start_layer,
)
L = n_layers - start_layer
assert model.layer_drop_probs == [0.0] * start_layer + [drop_prob] * L
# checking for errors
for drop_prob in [-1.0, 1.0]:
with pytest.raises(ValueError, match="stochastic_depth_drop_prob has to be in"):
ConformerEncoder(
feat_in=10,
n_layers=n_layers,
d_model=4,
feat_out=8,
stochastic_depth_drop_prob=drop_prob,
stochastic_depth_mode="uniform",
)
with pytest.raises(ValueError, match="stochastic_depth_mode has to be one of"):
ConformerEncoder(feat_in=10, n_layers=n_layers, d_model=4, feat_out=8, stochastic_depth_mode="weird")
for start_layer in [-1, 0, 5]:
with pytest.raises(ValueError, match="stochastic_depth_start_layer has to be in"):
ConformerEncoder(
feat_in=10, n_layers=n_layers, d_model=4, feat_out=8, stochastic_depth_start_layer=start_layer,
)
def test_stochastic_depth_forward(self):
"""Testing that forward works and we get randomness during training, but not during eval."""
random_input = torch.rand((1, 2, 2))
random_length = torch.tensor([2, 2], dtype=torch.int64)
model = ConformerEncoder(
feat_in=2,
n_layers=3,
d_model=4,
feat_out=4,
stochastic_depth_drop_prob=0.8,
dropout=0.0,
dropout_pre_encoder=0.0,
dropout_emb=0.0,
conv_norm_type="layer_norm",
conv_kernel_size=3,
)
model.train()
outputs = [None] * 5
for i in range(5):
outputs[i] = model(audio_signal=random_input, length=random_length)[0]
# checking that not all outputs are the same
num_diff = 0
for i in range(1, 5):
if not torch.allclose(outputs[i], outputs[0]):
num_diff += 1
assert num_diff > 0
model.eval()
outputs = [None] * 5
for i in range(5):
outputs[i] = model(audio_signal=random_input, length=random_length)[0]
# checking that not all outputs are the same
num_diff = 0
for i in range(1, 5):
if not torch.allclose(outputs[i], outputs[0]):
num_diff += 1
assert num_diff == 0
|