storytold/storyteller-ml
2
0
Fork 0
mirror of https://github.com/storytold/storyteller-ml.git synced 2026-10-28 09:08:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Introduce code from Rust monorepo. This is a mess

Browse Source
This commit is contained in:
Brandon Thomas
2022-02-28 05:43:54 -05:00
parent 921fb49f9a
commit ba326bfcfd
1058 changed files with 192192 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tts/NeMo/tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py
+256
View File
@@ -0,0 +1,256 @@
# Copyright (c) 2021, 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 numpy as np
import pytest
import torch
from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_numpy import RNNTLoss as RNNTLoss_Numpy
from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import RNNTLossNumba
from nemo.core.utils import numba_utils
from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
DEVICES = ['cpu']
if torch.cuda.is_available():
DEVICES.append('cuda')
def wrap_and_call(fn, acts, labels, device):
acts = torch.FloatTensor(acts)
if 'cuda' in device:
acts = acts.cuda()
acts.requires_grad = True
lengths = [acts.shape[1]] * acts.shape[0]
label_lengths = [len(l) for l in labels]
labels = torch.IntTensor(labels)
lengths = torch.IntTensor(lengths)
label_lengths = torch.IntTensor(label_lengths)
if 'cuda' in device:
labels = labels.cuda()
lengths = lengths.cuda()
label_lengths = label_lengths.cuda()
costs = fn(acts, labels, lengths, label_lengths)
cost = torch.sum(costs)
cost.backward()
return costs.data.cpu().numpy(), acts.grad.data.cpu().numpy()
class TestRNNTLossPytorch:
@pytest.mark.unit
@pytest.mark.parametrize('device', DEVICES)
def test_case_small(self, device):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
acts = np.array(
[
[
[[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.2, 0.8, 0.1]],
[[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.2, 0.1, 0.1], [0.7, 0.1, 0.2, 0.1, 0.1]],
]
]
)
labels = [[1, 2]]
fn_pt = RNNTLossNumba(blank=0, reduction='sum')
pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
fn_np = RNNTLoss_Numpy()
np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device)
expected_cost = 4.495666
expected_grads = np.array(
[
[
[
[-0.13116688, -0.3999269, 0.17703125, 0.17703125, 0.17703125],
[-0.18572757, 0.12247056, -0.18168412, 0.12247056, 0.12247056],
[-0.32091254, 0.06269141, 0.06928472, 0.12624499, 0.06269141],
],
[
[0.05456069, -0.21824276, 0.05456069, 0.05456069, 0.05456069],
[0.12073959, 0.12073959, -0.48295835, 0.12073959, 0.12073959],
[-0.6925882, 0.16871116, 0.18645467, 0.16871116, 0.16871116],
],
]
]
)
assert np.allclose(pt_cost, expected_cost, rtol=1e-6), "small_test costs mismatch."
assert np.allclose(pt_grads, expected_grads), "small_test gradient mismatch."
assert np.allclose(pt_cost, np_cost, rtol=1e-6), "small_test costs mismatch."
assert np.allclose(pt_grads, np_grads), "small_test gradient mismatch."
@pytest.mark.unit
@pytest.mark.parametrize('device', DEVICES)
def test_case_small_random(self, device):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
rng = np.random.RandomState(0)
acts = rng.randn(1, 4, 3, 3)
labels = [[1, 2]]
fn_pt = RNNTLossNumba(blank=0, reduction='sum')
pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
fn_np = RNNTLoss_Numpy()
np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device)
assert np.allclose(pt_cost, np_cost, rtol=1e-6), "small_random_test costs mismatch."
assert np.allclose(pt_grads, np_grads), "small_random_test gradient mismatch."
@pytest.mark.unit
@pytest.mark.parametrize('device', DEVICES)
def big_test(self, device):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# minibatch x T x U x alphabet_size
activations = [
[
[
[0.06535690384862791, 0.7875301411923206, 0.08159176605666074],
[0.5297155426466327, 0.7506749639230854, 0.7541348379087998],
[0.6097641124736383, 0.8681404965673826, 0.6225318186056529],
],
[
[0.6685222872103057, 0.8580392805336061, 0.16453892311765583],
[0.989779515236694, 0.944298460961015, 0.6031678586829663],
[0.9467833543605416, 0.666202507295747, 0.28688179752461884],
],
[
[0.09418426230195986, 0.3666735970751962, 0.736168049462793],
[0.1666804425271342, 0.7141542198635192, 0.3993997272216727],
[0.5359823524146038, 0.29182076440286386, 0.6126422611507932],
],
[
[0.3242405528768486, 0.8007644367291621, 0.5241057606558068],
[0.779194617063042, 0.18331417220174862, 0.113745182072432],
[0.24022162381327106, 0.3394695622533106, 0.1341595066017014],
],
],
[
[
[0.5055615569388828, 0.051597282072282646, 0.6402903936686337],
[0.43073311517251, 0.8294731834714112, 0.1774668847323424],
[0.3207001991262245, 0.04288308912457006, 0.30280282975568984],
],
[
[0.6751777088333762, 0.569537369330242, 0.5584738347504452],
[0.08313242153985256, 0.06016544344162322, 0.10795752845152584],
[0.7486153608562472, 0.943918041459349, 0.4863558118797222],
],
[
[0.4181986264486809, 0.6524078485043804, 0.024242983423721887],
[0.13458171554507403, 0.3663418070512402, 0.2958297395361563],
[0.9236695822497084, 0.6899291482654177, 0.7418981733448822],
],
[
[0.25000547599982104, 0.6034295486281007, 0.9872887878887768],
[0.5926057265215715, 0.8846724004467684, 0.5434495396894328],
[0.6607698886038497, 0.3771277082495921, 0.3580209022231813],
],
],
]
expected_costs = [4.2806528590890736, 3.9384369822503591]
expected_grads = [
[
[
[-1.86843902e-01, -6.25548810e-02, 2.49398798e-01],
[-2.03376666e-01, 2.02399328e-01, 9.77333169e-04],
[-1.41016081e-01, 7.91234672e-02, 6.18926100e-02],
],
[
[-1.15517676e-02, -8.12802389e-02, 9.28319991e-02],
[-1.54257029e-01, 2.29432687e-01, -7.51756504e-02],
[-2.46593088e-01, 1.46404594e-01, 1.00188486e-01],
],
[
[-1.29182907e-02, -6.15932420e-02, 7.45115355e-02],
[-5.59857301e-02, 2.19830811e-01, -1.63845062e-01],
[-4.97626871e-01, 2.09239945e-01, 2.88386941e-01],
],
[
[1.36048580e-02, -3.02196294e-02, 1.66147724e-02],
[1.13924511e-01, 6.27811998e-02, -1.76705718e-01],
[-6.67078257e-01, 3.67658824e-01, 2.99419403e-01],
],
],
[
[
[-3.56343776e-01, -5.53474613e-02, 4.11691219e-01],
[-9.69219357e-02, 2.94591039e-02, 6.74628317e-02],
[-6.35175705e-02, 2.76544970e-02, 3.58630717e-02],
],
[
[-1.54499024e-01, -7.39420280e-02, 2.28441030e-01],
[-1.66789949e-01, -8.78955179e-05, 1.66877866e-01],
[-1.72369644e-01, 1.05565332e-01, 6.68043196e-02],
],
[
[2.38748826e-02, -1.18255816e-01, 9.43809375e-02],
[-1.04707085e-01, -1.08934477e-01, 2.13641584e-01],
[-3.69844258e-01, 1.80118099e-01, 1.89726159e-01],
],
[
[2.57137045e-02, -7.94617534e-02, 5.37480488e-02],
[1.22328237e-01, -2.38788679e-01, 1.16460443e-01],
[-5.98686993e-01, 3.02203178e-01, 2.96483815e-01],
],
],
]
activations = np.array(activations)
labels = [[1, 2], [1, 1]]
fn_pt = RNNTLossNumba(blank=0, reduction='sum')
pt_costs, pt_grads = wrap_and_call(fn_pt, activations, labels, device)
fn_np = RNNTLoss_Numpy()
np_costs, np_grads = wrap_and_call(fn_np, activations, labels, device)
assert np.allclose(pt_costs, sum(expected_costs)), "big_test average costs mismatch."
assert np.allclose(pt_grads, expected_grads, rtol=1e-3), "big_test grads for average cost mismatch."
assert np.allclose(pt_costs, np_costs), "big_test average costs mismatch."
assert np.allclose(pt_grads, np_grads, rtol=1e-3), "big_test grads for average cost mismatch."
@pytest.mark.unit
@pytest.mark.parametrize('device', DEVICES)
def test_case_large_random(self, device):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
rng = np.random.RandomState(0)
acts = rng.randn(4, 8, 11, 5)
labels = [
[1, 2, 4, 3, 2, 2, 1, 1, 1, 1],
[3, 2, 2, 3, 4, 1, 1, 1, 1, 1],
[4, 4, 1, 2, 1, 3, 4, 3, 1, 2],
[1, 1, 2, 1, 2, 3, 3, 1, 1, 1],
]
fn_pt = RNNTLossNumba(blank=0, reduction='sum')
pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device)
fn_np = RNNTLoss_Numpy()
np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device)
assert np.allclose(pt_cost, np_cost, atol=1e-5, rtol=1e-3), "large_random_test costs mismatch."
assert np.allclose(pt_grads, np_grads, atol=1e-5, rtol=1e-3), "large_random_test gradient mismatch."
if __name__ == "__main__":
pytest.main([__file__])
tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_gpu_rnnt_kernel.py
+169
View File
@@ -0,0 +1,169 @@
# Copyright (c) 2021, 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 numpy as np
import pytest
import torch
from numba import cuda
from nemo.collections.asr.parts.numba.rnnt_loss import rnnt_numpy
from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import certify_inputs
from nemo.collections.asr.parts.numba.rnnt_loss.utils.cuda_utils import gpu_rnnt_kernel, reduce
from nemo.core.utils import numba_utils
from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
def log_softmax(x, axis=-1):
x = torch.from_numpy(x) # zero-copy
x = torch.log_softmax(x, axis)
x = x.numpy()
return x
class TestRNNTCUDAKernels:
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
@pytest.mark.unit
def test_compute_alphas_kernel(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
random = np.random.RandomState(0)
original_shape = [1, 5, 11, 3]
B, T, U, V = original_shape
# Numpy kernel
x = random.randn(*original_shape)
labels = np.array([[1, 1, 1, 2, 2, 2, 1, 2, 2, 1]]) # [1, 10]
label_len = len(labels[0]) + 1
blank_idx = 0
x_np = log_softmax(x, axis=-1)
ground_alphas, ground_log_likelihood = rnnt_numpy.forward_pass(
x_np[0, :, :label_len, :], labels[0, : label_len - 1], blank_idx
)
# Pytorch kernel
device = torch.device('cuda')
if hasattr(cuda, 'external_stream'):
stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream)
else:
stream = cuda.default_stream()
x_c = torch.tensor(x, device=device, dtype=torch.float32)
labels_c = torch.tensor(labels, device=device, dtype=torch.int32)
# Allocate workspace memory
denom = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
alphas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
llForward = torch.zeros(B, device=device, dtype=x_c.dtype)
input_lengths = torch.tensor([T], dtype=torch.int32, device=device)
label_lengths = torch.tensor([len(labels[0])], dtype=torch.int32, device=device)
# certify input data
certify_inputs(x_c, labels_c, input_lengths, label_lengths)
# flatten activation tensor (for pointer based indexing)
x_c = x_c.view([-1])
# call kernel
# log softmax reduction
reduce.reduce_max(x_c, denom, rows=V, cols=B * T * U, minus=False, stream=stream)
reduce.reduce_exp(x_c, denom, rows=V, cols=B * T * U, minus=True, stream=stream)
# alpha kernel
gpu_rnnt_kernel.compute_alphas_kernel[B, U, stream, 0](
x_c, denom, alphas, llForward, input_lengths, label_lengths, labels_c, B, T, U, V, blank_idx,
)
# sync kernel
stream.synchronize()
# reshape alphas
alphas = alphas.view([B, T, U])
diff = ground_alphas - alphas[0].cpu().numpy()
assert np.abs(diff).mean() <= 1e-5
assert np.square(diff).mean() <= 1e-10
ll_diff = ground_log_likelihood - llForward[0].cpu().numpy()
assert np.abs(ll_diff).mean() <= 1e-5
assert np.square(ll_diff).mean() <= 1e-10
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
@pytest.mark.unit
def test_compute_betas_kernel(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
random = np.random.RandomState(0)
original_shape = [1, 5, 11, 3]
B, T, U, V = original_shape
# Numpy kernel
x = random.randn(*original_shape)
labels = np.array([[1, 1, 1, 2, 2, 2, 1, 2, 2, 1]]) # [1, 10]
label_len = len(labels[0]) + 1
blank_idx = 0
x_np = log_softmax(x, axis=-1)
ground_alphas, ground_log_likelihood = rnnt_numpy.backward_pass(
x_np[0, :, :label_len, :], labels[0, : label_len - 1], blank_idx
)
# Pytorch kernel
device = torch.device('cuda')
if hasattr(cuda, 'external_stream'):
stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream)
else:
stream = cuda.default_stream()
x_c = torch.tensor(x, device=device, dtype=torch.float32)
labels_c = torch.tensor(labels, device=device, dtype=torch.int32)
# Allocate workspace memory
denom = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
betas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype)
llBackward = torch.zeros(B, device=device, dtype=x_c.dtype)
input_lengths = torch.tensor([T], dtype=torch.int32, device=device)
label_lengths = torch.tensor([len(labels[0])], dtype=torch.int32, device=device)
# certify input data
certify_inputs(x_c, labels_c, input_lengths, label_lengths)
# flatten activation tensor (for pointer based indexing)
x_c = x_c.view([-1])
# call kernel
# log softmax reduction
reduce.reduce_max(x_c, denom, rows=V, cols=B * T * U, minus=False, stream=stream)
reduce.reduce_exp(x_c, denom, rows=V, cols=B * T * U, minus=True, stream=stream)
# beta kernel
gpu_rnnt_kernel.compute_betas_kernel[B, U, stream, 0](
x_c, denom, betas, llBackward, input_lengths, label_lengths, labels_c, B, T, U, V, blank_idx,
)
# sync kernel
stream.synchronize()
# reshape alphas
betas = betas.view([B, T, U])
diff = ground_alphas - betas[0].cpu().numpy()
assert np.abs(diff).mean() <= 1e-5
assert np.square(diff).mean() <= 1e-10
ll_diff = ground_log_likelihood - llBackward[0].cpu().numpy()
assert np.abs(ll_diff).mean() <= 1e-5
assert np.square(ll_diff).mean() <= 1e-10
tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_reduce.py
+86
View File
@@ -0,0 +1,86 @@
# Copyright (c) 2021, 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 numpy as np
import pytest
from numba import cuda
from nemo.collections.asr.parts.numba.rnnt_loss.utils.cuda_utils import reduce
from nemo.core.utils import numba_utils
from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
class TestRNNTCUDAReductions:
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
@pytest.mark.unit
def test_reduce_max(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
random = np.random.RandomState(0)
original_shape = [1, 5, 4, 3]
x = random.randn(*original_shape).reshape([-1])
dx = random.randn(*x.shape)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
dx_c = cuda.to_device(dx, stream=stream)
# call kernel
cols = np.prod(original_shape[:3])
reduce.reduce_max(x_c, dx_c, rows=original_shape[-1], cols=cols, minus=False, stream=stream)
# sync kernel
stream.synchronize()
dx_result = dx_c.copy_to_host(stream=stream)
del x_c, dx_c
# collect results in first [B * T * U] values; for all V
assert np.abs(dx_result[cols:] - dx[cols:]).sum() <= 1e-7
# make sure dx_result updates the [B * T * U] values
assert np.abs(dx_result[:cols] - dx[:cols]).sum() > 0
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available")
@pytest.mark.unit
def test_reduce_exp(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
random = np.random.RandomState(0)
original_shape = [1, 5, 4, 2]
x = random.randn(*original_shape).reshape([-1])
dx = np.zeros_like(x)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
dx_c = cuda.to_device(dx, stream=stream)
# call kernel
cols = np.prod(original_shape[:3])
reduce.reduce_exp(x_c, dx_c, rows=original_shape[-1], cols=cols, minus=False, stream=stream)
# sync kernel
stream.synchronize()
dx_result = dx_c.copy_to_host(stream=stream)
del x_c, dx_c
# collect results in first [B * T * U] values; for all V
assert (dx_result[cols:] - dx[cols:]).sum() <= 1e-7
# make sure dx_result updates the [B * T * U] values
assert np.abs(dx_result[:cols] - dx[:cols]).sum() > 0
if __name__ == '__main__':
pytest.main([__file__])
tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_rnnt_helper.py
+319
View File
@@ -0,0 +1,319 @@
# Copyright (c) 2021, 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 numpy as np
import pytest
from numba import cuda
from nemo.collections.asr.parts.numba.rnnt_loss.utils import global_constants, rnnt_helper
from nemo.core.utils import numba_utils
from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
class TestRNNTHelper:
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
@pytest.mark.unit
def test_log_sum_exp(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# wrapper kernel for device function that is tested
@cuda.jit
def _kernel(x, y):
x_pos = cuda.grid(1)
if x_pos < x.shape[0] and x_pos < y.shape[0]:
x[x_pos] = rnnt_helper.log_sum_exp(x[x_pos], y[x_pos])
x = np.zeros([8]) # np.random.rand(8192)
y = np.ones([8]) # np.random.rand(8192)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
y_c = cuda.to_device(y, stream=stream)
# call kernel
threads_per_block = global_constants.threads_per_block()
blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
_kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
# sync kernel
stream.synchronize()
x_new = x_c.copy_to_host(stream=stream)
del x_c, y_c
assert (x_new.sum() - 10.506093500145782) <= 1e-5
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
@pytest.mark.unit
def test_log_sum_exp_neg_inf(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# wrapper kernel for device function that is tested
@cuda.jit
def _kernel(x, y):
x_pos = cuda.grid(1)
if x_pos < x.shape[0] and x_pos < y.shape[0]:
x[x_pos] = rnnt_helper.log_sum_exp(x[x_pos], y[x_pos])
x = np.asarray([global_constants.FP32_NEG_INF] * 8)
y = np.ones([len(x)])
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
y_c = cuda.to_device(y, stream=stream)
# call kernel
threads_per_block = global_constants.threads_per_block()
blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
_kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
# sync kernel
stream.synchronize()
x_new = x_c.copy_to_host(stream=stream)
del x_c, y_c
assert np.allclose(x_new, np.ones_like(x_new), atol=1e-5)
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
@pytest.mark.unit
def test_div_up(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# wrapper kernel for device function that is tested
@cuda.jit
def _kernel(x, y):
x_pos = cuda.grid(1)
if x_pos < x.shape[0] and x_pos < y.shape[0]:
x[x_pos] = rnnt_helper.div_up(x[x_pos], y[x_pos])
x = np.full([8], fill_value=10) # np.random.rand(8192)
y = np.full([8], fill_value=2) # np.random.rand(8192)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
y_c = cuda.to_device(y, stream=stream)
# call kernel
threads_per_block = global_constants.threads_per_block()
blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
_kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
# sync kernel
stream.synchronize()
x_new = x_c.copy_to_host(stream=stream)
del x_c, y_c
for i in range(len(x_new)):
assert x_new[i] == ((10 + 2 - 1) // 2)
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
@pytest.mark.unit
def test_add(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# wrapper kernel for device function that is tested
@cuda.jit
def _kernel(x, y):
x_pos = cuda.grid(1)
if x_pos < x.shape[0] and x_pos < y.shape[0]:
x[x_pos] = rnnt_helper.add(x[x_pos], y[x_pos])
x = np.full([8], fill_value=10) # np.random.rand(8192)
y = np.full([8], fill_value=2) # np.random.rand(8192)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
y_c = cuda.to_device(y, stream=stream)
# call kernel
threads_per_block = global_constants.threads_per_block()
blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
_kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
# sync kernel
stream.synchronize()
x_new = x_c.copy_to_host(stream=stream)
del x_c, y_c
for i in range(len(x_new)):
assert x_new[i] == 12
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
@pytest.mark.unit
def test_maximum(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# wrapper kernel for device function that is tested
@cuda.jit
def _kernel(x, y):
x_pos = cuda.grid(1)
if x_pos < x.shape[0] and x_pos < y.shape[0]:
x[x_pos] = rnnt_helper.maximum(x[x_pos], y[x_pos])
x = np.full([8], fill_value=10) # np.random.rand(8192)
y = np.full([8], fill_value=2) # np.random.rand(8192)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
y_c = cuda.to_device(y, stream=stream)
# call kernel
threads_per_block = global_constants.threads_per_block()
blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
_kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
# sync kernel
stream.synchronize()
x_new = x_c.copy_to_host(stream=stream)
del x_c, y_c
for i in range(len(x_new)):
assert x_new[i] == 10
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
@pytest.mark.unit
def test_identity(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# wrapper kernel for device function that is tested
@cuda.jit
def _kernel(x):
x_pos = cuda.grid(1)
if x_pos < x.shape[0]:
x[x_pos] = rnnt_helper.identity(x[x_pos])
x = np.full([8], fill_value=10) # np.random.rand(8192)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
# call kernel
threads_per_block = global_constants.threads_per_block()
blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
_kernel[blocks_per_grid, threads_per_block, stream](x_c)
# sync kernel
stream.synchronize()
x_new = x_c.copy_to_host(stream=stream)
del x_c
for i in range(len(x_new)):
assert x_new[i] == x[i]
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
@pytest.mark.unit
def test_negate(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# wrapper kernel for device function that is tested
@cuda.jit
def _kernel(x):
x_pos = cuda.grid(1)
if x_pos < x.shape[0]:
x[x_pos] = rnnt_helper.negate(x[x_pos])
x = np.full([8], fill_value=10) # np.random.rand(8192)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
# call kernel
threads_per_block = global_constants.threads_per_block()
blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
_kernel[blocks_per_grid, threads_per_block, stream](x_c)
# sync kernel
stream.synchronize()
x_new = x_c.copy_to_host(stream=stream)
del x_c
for i in range(len(x_new)):
assert x_new[i] == -x[i]
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
@pytest.mark.unit
def test_exponential(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# wrapper kernel for device function that is tested
@cuda.jit
def _kernel(x):
x_pos = cuda.grid(1)
if x_pos < x.shape[0]:
x[x_pos] = rnnt_helper.exponential(x[x_pos])
x = np.random.rand(8)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
# call kernel
threads_per_block = global_constants.threads_per_block()
blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
_kernel[blocks_per_grid, threads_per_block, stream](x_c)
# sync kernel
stream.synchronize()
x_new = x_c.copy_to_host(stream=stream)
del x_c
y = np.exp(x)
for i in range(len(x_new)):
assert (x_new[i] - y[i]) < 1e-4
@pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available")
@pytest.mark.unit
def test_log_plus(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
# wrapper kernel for device function that is tested
@cuda.jit
def _kernel(x, y):
x_pos = cuda.grid(1)
if x_pos < x.shape[0] and x_pos < y.shape[0]:
x[x_pos] = rnnt_helper.log_plus(x[x_pos], y[x_pos])
x = np.full([8], fill_value=10.0) # np.random.rand(8192)
y = np.full([8], fill_value=2.0) # np.random.rand(8192)
stream = cuda.stream()
x_c = cuda.to_device(x, stream=stream)
y_c = cuda.to_device(y, stream=stream)
# call kernel
threads_per_block = global_constants.threads_per_block()
blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block
_kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c)
# sync kernel
stream.synchronize()
x_new = x_c.copy_to_host(stream=stream)
del x_c, y_c
z = np.log1p(np.exp(-np.fabs(x - y))) + np.maximum(x, y)
for i in range(len(x_new)):
assert x_new[i] == z[i]
if __name__ == '__main__':
pytest.main([__file__])
tts/NeMo/tests/collections/asr/numba/spec_augment/test_spec_aug_numba.py
+252
View File
@@ -0,0 +1,252 @@
# Copyright (c) 2021, 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 omegaconf import OmegaConf
from nemo.collections.asr.parts.numba.spec_augment import spec_aug_numba
from nemo.core.utils import numba_utils
from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
def get_cfg(seed=0, dtype='float32'):
# fmt: off
cfg = OmegaConf.create(
dict(b=2, f=80, t=20, device='cuda',
freq_masks=2, time_masks=2, freq_width=27, time_width=0.05, mask_value=0.0,
seed=seed, dtype=dtype)
)
# fmt: on
return cfg
# fmt: off
def prepare_data(b, f, t, device='cuda', freq_masks=0, time_masks=0, freq_width=10, time_width=0.1,
seed=0, dtype='float32',
**kwargs):
torch.manual_seed(seed)
if dtype == 'float16':
dtype = torch.float16
else:
dtype = torch.float
x = torch.randn([b, f, t], dtype=dtype, device=device)
x_len = torch.randint(t, size=[b], device=x.device)
sh = x.shape
bs = sh[0]
if isinstance(time_width, int):
adaptive_temporal_width = False
else:
if time_width > 1.0 or time_width < 0.0:
raise ValueError('If `time_width` is a float value, must be in range [0, 1]')
adaptive_temporal_width = True
if adaptive_temporal_width:
time_width = max(1, int(sh[2] * time_width))
else:
time_width = time_width
# Construct the freq and time masks as well as start positions
if freq_masks > 0:
freq_starts = torch.randint(0, sh[1] - freq_width + 1, size=[bs, freq_masks], device=x.device)
freq_lengths = torch.randint(0, freq_width + 1, size=[bs, freq_masks], device=x.device)
else:
freq_starts = torch.zeros([bs, 1], dtype=torch.int64, device=x.device)
freq_lengths = torch.zeros([bs, 1], dtype=torch.int64, device=x.device)
if time_masks > 0:
time_starts = torch.randint(0, sh[2] - time_width + 1, size=[bs, time_masks], device=x.device)
time_lengths = torch.randint(0, time_width + 1, size=[bs, time_masks], device=x.device)
else:
time_starts = torch.zeros([bs, 1], dtype=torch.int64, device=x.device)
time_lengths = torch.zeros([bs, 1], dtype=torch.int64, device=x.device)
output = dict(
x=x,
x_len=x_len,
freq_starts=freq_starts,
freq_lengths=freq_lengths,
time_starts=time_starts,
time_lengths=time_lengths,
sh=sh,
)
return output
# fmt: on
def launch_kernel(data, cfg):
# Launch CUDA kernel
# fmt: off
data['x'] = spec_aug_numba.launch_spec_augment_kernel(
x=data['x'], x_len=data['x_len'],
freq_starts=data['freq_starts'], freq_lengths=data['freq_lengths'],
time_starts=data['time_starts'], time_lengths=data['time_lengths'],
freq_masks=cfg.freq_masks, time_masks=cfg.time_masks, mask_value=cfg.mask_value
)
# fmt: on
def freq_mask_check(x, x_len, f_start, f_len, mask_value, bidx):
check_result = True
for fidx in range(f_start, f_start + f_len):
if not (x[bidx, fidx, :] == mask_value).all():
check_result = False
break
assert check_result
def time_mask_check(x, x_len, t_start, t_len, mask_value, bidx):
check_result = True
for tidx in range(t_start, t_start + t_len):
if tidx >= x_len[bidx]:
# this sample has smaller length than the time index of mask, ignore
continue
if not (x[bidx, :, tidx] == mask_value).all():
check_result = False
break
assert check_result
class TestSpecAugmentNumba:
@pytest.mark.unit
@pytest.mark.run_only_on('GPU')
@pytest.mark.parametrize('dtype', ['float16', 'float32'])
def test_spec_aug_kernel(self, dtype):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
cfg = get_cfg(seed=0, dtype=dtype)
cfg.freq_masks = 2
cfg.time_masks = 10
data = prepare_data(**cfg)
launch_kernel(data, cfg)
x, x_len, sh = data['x'], data['x_len'], data['sh']
# Assert freq masks are correct
for bidx in range(sh[0]):
for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]):
freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx)
# Assert time masks are correct
for bidx in range(sh[0]):
for t_start, t_len in zip(data['time_starts'][bidx], data['time_lengths'][bidx]):
time_mask_check(x, x_len, t_start, t_len, mask_value=cfg.mask_value, bidx=bidx)
@pytest.mark.unit
@pytest.mark.run_only_on('GPU')
def test_spec_aug_kernel_mask_value(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
cfg = get_cfg(seed=0)
cfg.freq_masks = 2
cfg.time_masks = 10
cfg.mask_value = -1.0
data = prepare_data(**cfg)
launch_kernel(data, cfg)
x, x_len, sh = data['x'], data['x_len'], data['sh']
# Assert freq masks are correct
for bidx in range(sh[0]):
for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]):
freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx)
# Assert time masks are correct
for bidx in range(sh[0]):
for t_start, t_len in zip(data['time_starts'][bidx], data['time_lengths'][bidx]):
time_mask_check(x, x_len, t_start, t_len, mask_value=cfg.mask_value, bidx=bidx)
@pytest.mark.unit
@pytest.mark.run_only_on('GPU')
def test_spec_aug_kernel_grad(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
cfg = get_cfg(seed=0)
cfg.freq_masks = 2
cfg.time_masks = 10
data = prepare_data(**cfg)
launch_kernel(data, cfg)
result = data['x'] # inplace modification via kernel
y = torch.ones_like(result, requires_grad=True)
z = y + result
z.mean().backward()
assert y.grad is not None
@pytest.mark.unit
@pytest.mark.run_only_on('GPU')
def test_spec_aug_kernel_no_freq_mask(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
cfg = get_cfg(seed=0)
cfg.freq_masks = 0
cfg.time_masks = 10
data = prepare_data(**cfg)
launch_kernel(data, cfg)
x, x_len, sh = data['x'], data['x_len'], data['sh']
# Assert time masks are correct
for bidx in range(sh[0]):
for t_start, t_len in zip(data['time_starts'][bidx], data['time_lengths'][bidx]):
time_mask_check(x, x_len, t_start, t_len, mask_value=cfg.mask_value, bidx=bidx)
@pytest.mark.unit
@pytest.mark.run_only_on('GPU')
def test_spec_aug_kernel_no_time_mask(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
cfg = get_cfg(seed=0)
cfg.freq_masks = 2
cfg.time_masks = 0
data = prepare_data(**cfg)
launch_kernel(data, cfg)
x, x_len, sh = data['x'], data['x_len'], data['sh']
# Assert freq masks are correct
for bidx in range(sh[0]):
for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]):
freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx)
@pytest.mark.unit
@pytest.mark.run_only_on('GPU')
def test_spec_aug_kernel_no_freq_time_mask(self):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
cfg = get_cfg(seed=0)
cfg.freq_masks = 0
cfg.time_masks = 0
data = prepare_data(**cfg)
x, x_len, sh = data['x'], data['x_len'], data['sh']
x_copy = x.clone()
launch_kernel(data, cfg)
# Assert no data edits occured
assert (data['x'] - x_copy).abs().mean() <= 1e-9
tts/NeMo/tests/collections/asr/test_asr_classification_model.py
+208
View File
@@ -0,0 +1,208 @@
# Copyright (c) 2020, 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 copy
import os
import pytest
import torch
from omegaconf import DictConfig, ListConfig
from nemo.collections.asr.data import audio_to_label
from nemo.collections.asr.models import EncDecClassificationModel, configs
from nemo.utils.config_utils import assert_dataclass_signature_match
@pytest.fixture()
def speech_classification_model():
preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
encoder = {
'cls': 'nemo.collections.asr.modules.ConvASREncoder',
'params': {
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': 32,
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': True,
'se': True,
'se_context_size': -1,
}
],
},
}
decoder = {
'cls': 'nemo.collections.asr.modules.ConvASRDecoderClassification',
'params': {'feat_in': 32, 'num_classes': 30,},
}
modelConfig = DictConfig(
{
'preprocessor': DictConfig(preprocessor),
'encoder': DictConfig(encoder),
'decoder': DictConfig(decoder),
'labels': ListConfig(["dummy_cls_{}".format(i + 1) for i in range(30)]),
}
)
model = EncDecClassificationModel(cfg=modelConfig)
return model
class TestEncDecClassificationModel:
@pytest.mark.unit
def test_constructor(self, speech_classification_model):
asr_model = speech_classification_model.train()
conv_cnt = (64 * 32 * 1 + 32) + (64 * 1 * 1 + 32) # separable kernel + bias + pointwise kernel + bias
bn_cnt = (4 * 32) * 2 # 2 * moving averages
dec_cnt = 32 * 30 + 30 # fc + bias
param_count = conv_cnt + bn_cnt + dec_cnt
assert asr_model.num_weights == param_count
# Check to/from config_dict:
confdict = asr_model.to_config_dict()
instance2 = EncDecClassificationModel.from_config_dict(confdict)
assert isinstance(instance2, EncDecClassificationModel)
@pytest.mark.unit
def test_forward(self, speech_classification_model):
asr_model = speech_classification_model.eval()
asr_model.preprocessor.featurizer.dither = 0.0
asr_model.preprocessor.featurizer.pad_to = 0
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
with torch.no_grad():
# batch size 1
logprobs_instance = []
for i in range(input_signal.size(0)):
logprobs_ins = asr_model.forward(
input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
)
logprobs_instance.append(logprobs_ins)
logprobs_instance = torch.cat(logprobs_instance, 0)
# batch size 4
logprobs_batch = asr_model.forward(input_signal=input_signal, input_signal_length=length)
assert logprobs_instance.shape == logprobs_batch.shape
diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch))
assert diff <= 1e-6
diff = torch.max(torch.abs(logprobs_instance - logprobs_batch))
assert diff <= 1e-6
@pytest.mark.unit
def test_vocab_change(self, speech_classification_model):
asr_model = speech_classification_model.train()
old_labels = copy.deepcopy(asr_model._cfg.labels)
nw1 = asr_model.num_weights
asr_model.change_labels(new_labels=old_labels)
# No change
assert nw1 == asr_model.num_weights
new_labels = copy.deepcopy(old_labels)
new_labels.append('dummy_cls_31')
new_labels.append('dummy_cls_32')
new_labels.append('dummy_cls_33')
asr_model.change_labels(new_labels=new_labels)
# fully connected + bias
assert asr_model.num_weights == nw1 + 3 * (asr_model.decoder._feat_in + 1)
@pytest.mark.unit
def test_transcription(self, speech_classification_model, test_data_dir):
# Ground truth labels = ["yes", "no"]
audio_filenames = ['an22-flrp-b.wav', 'an90-fbbh-b.wav']
audio_paths = [os.path.join(test_data_dir, "asr", "train", "an4", "wav", fp) for fp in audio_filenames]
model = speech_classification_model.eval()
# Test Top 1 classification transcription
results = model.transcribe(audio_paths, batch_size=2)
assert len(results) == 2
assert results[0].shape == torch.Size([1])
# Test Top 5 classification transcription
model._accuracy.top_k = [5] # set top k to 5 for accuracy calculation
results = model.transcribe(audio_paths, batch_size=2)
assert len(results) == 2
assert results[0].shape == torch.Size([5])
# Test Top 1 and Top 5 classification transcription
model._accuracy.top_k = [1, 5]
results = model.transcribe(audio_paths, batch_size=2)
assert len(results) == 2
assert results[0].shape == torch.Size([2, 1])
assert results[1].shape == torch.Size([2, 5])
assert model._accuracy.top_k == [1, 5]
# Test log probs extraction
model._accuracy.top_k = [1]
results = model.transcribe(audio_paths, batch_size=2, logprobs=True)
assert len(results) == 2
assert results[0].shape == torch.Size([len(model.cfg.labels)])
# Test log probs extraction remains same for any top_k
model._accuracy.top_k = [5]
results = model.transcribe(audio_paths, batch_size=2, logprobs=True)
assert len(results) == 2
assert results[0].shape == torch.Size([len(model.cfg.labels)])
@pytest.mark.unit
def test_EncDecClassificationDatasetConfig_for_AudioToSpeechLabelDataset(self):
# ignore some additional arguments as dataclass is generic
IGNORE_ARGS = [
'is_tarred',
'num_workers',
'batch_size',
'tarred_audio_filepaths',
'shuffle',
'pin_memory',
'drop_last',
'tarred_shard_strategy',
'shuffle_n',
# `featurizer` is supplied at runtime
'featurizer',
# additional ignored arguments
'vad_stream',
'int_values',
'sample_rate',
'normalize_audio',
'augmentor',
]
REMAP_ARGS = {'trim_silence': 'trim'}
result = assert_dataclass_signature_match(
audio_to_label.AudioToSpeechLabelDataset,
configs.EncDecClassificationDatasetConfig,
ignore_args=IGNORE_ARGS,
remap_args=REMAP_ARGS,
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
tts/NeMo/tests/collections/asr/test_asr_ctc_encoder_model_bpe.py
+313
View File
@@ -0,0 +1,313 @@
# Copyright (c) 2020, 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 copy
import os
import shutil
import tempfile
import pytest
import torch
from omegaconf import DictConfig
from nemo.collections.asr.data import audio_to_text
from nemo.collections.asr.models import configs
from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE
from nemo.collections.common import tokenizers
from nemo.utils.config_utils import assert_dataclass_signature_match
@pytest.fixture()
def asr_model(test_data_dir):
preprocessor = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'}
encoder = {
'_target_': 'nemo.collections.asr.modules.ConvASREncoder',
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': 1024,
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': True,
'se': True,
'se_context_size': -1,
}
],
}
decoder = {
'_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
'feat_in': 1024,
'num_classes': -1,
'vocabulary': None,
}
tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'}
modelConfig = DictConfig(
{
'preprocessor': DictConfig(preprocessor),
'encoder': DictConfig(encoder),
'decoder': DictConfig(decoder),
'tokenizer': DictConfig(tokenizer),
}
)
model_instance = EncDecCTCModelBPE(cfg=modelConfig)
return model_instance
class TestEncDecCTCModel:
@pytest.mark.unit
def test_constructor(self, asr_model):
asr_model.train()
# TODO: make proper config and assert correct number of weights
# Check to/from config_dict:
confdict = asr_model.to_config_dict()
instance2 = EncDecCTCModelBPE.from_config_dict(confdict)
assert isinstance(instance2, EncDecCTCModelBPE)
@pytest.mark.unit
def test_forward(self, asr_model):
asr_model = asr_model.eval()
asr_model.preprocessor.featurizer.dither = 0.0
asr_model.preprocessor.featurizer.pad_to = 0
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
with torch.no_grad():
# batch size 1
logprobs_instance = []
for i in range(input_signal.size(0)):
logprobs_ins, _, _ = asr_model.forward(
input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
)
logprobs_instance.append(logprobs_ins)
print(len(logprobs_ins))
logprobs_instance = torch.cat(logprobs_instance, 0)
# batch size 4
logprobs_batch, _, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length)
assert logprobs_instance.shape == logprobs_batch.shape
diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch))
assert diff <= 1e-6
diff = torch.max(torch.abs(logprobs_instance - logprobs_batch))
assert diff <= 1e-6
@pytest.mark.unit
def test_save_restore_artifact(self, asr_model):
with tempfile.TemporaryDirectory() as tmpdir:
save_path = os.path.join(tmpdir, 'ctc_bpe.nemo')
asr_model.train()
asr_model.save_to(save_path)
new_model = EncDecCTCModelBPE.restore_from(save_path)
assert isinstance(new_model, type(asr_model))
assert new_model.vocab_path.endswith('_vocab.txt')
assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128
@pytest.mark.unit
def test_save_restore_artifact_spe(self, asr_model, test_data_dir):
with tempfile.TemporaryDirectory() as tmpdir:
tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
asr_model.change_vocabulary(new_tokenizer_dir=tokenizer_dir, new_tokenizer_type='bpe')
save_path = os.path.join(tmpdir, 'ctc_bpe.nemo')
asr_model.train()
asr_model.save_to(save_path)
new_model = EncDecCTCModelBPE.restore_from(save_path)
assert isinstance(new_model, type(asr_model))
assert isinstance(new_model.tokenizer, tokenizers.SentencePieceTokenizer)
assert new_model.model_path.endswith('_tokenizer.model')
assert new_model.vocab_path.endswith('_vocab.txt')
assert new_model.spe_vocab_path.endswith('_tokenizer.vocab')
assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128
@pytest.mark.unit
def test_vocab_change(self, test_data_dir, asr_model):
old_vocab = copy.deepcopy(asr_model.decoder.vocabulary)
with tempfile.TemporaryDirectory() as save_dir:
save_path = os.path.join(save_dir, 'temp.nemo')
with tempfile.TemporaryDirectory() as tmpdir:
old_tmpdir_path = tmpdir
old_tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
new_tokenizer_dir = os.path.join(tmpdir, 'tokenizer')
os.makedirs(new_tokenizer_dir, exist_ok=True)
shutil.copy2(old_tokenizer_dir, new_tokenizer_dir)
nw1 = asr_model.num_weights
asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
# No change
assert nw1 == asr_model.num_weights
with open(os.path.join(new_tokenizer_dir, 'vocab.txt'), 'a+') as f:
f.write("!\n")
f.write('$\n')
f.write('@\n')
asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
# fully connected + bias
assert asr_model.num_weights == nw1 + 3 * (asr_model.decoder._feat_in + 1)
new_vocab = copy.deepcopy(asr_model.decoder.vocabulary)
assert len(old_vocab) != len(new_vocab)
# save the model (after change of vocabulary)
asr_model.save_to(save_path)
assert os.path.exists(save_path)
# delete copied version of the vocabulary from nested tmpdir (by scope)
# assert copied vocab no longer exists
assert not os.path.exists(os.path.join(old_tmpdir_path, 'tokenizer', 'vocab.txt'))
# make a copy of the tokenizer before renaming
try:
os.rename(old_tokenizer_dir, old_tokenizer_dir + '.bkp')
assert not os.path.exists(old_tokenizer_dir)
# restore model from .nemo
asr_model2 = EncDecCTCModelBPE.restore_from(save_path)
assert isinstance(asr_model2, EncDecCTCModelBPE)
# Check if vocabulary size is same
assert asr_model.tokenizer.tokenizer.vocab_size == asr_model2.tokenizer.tokenizer.vocab_size
# Make a copy of the tokenizer
new_tokenizer_dir = os.path.join(save_dir, 'tokenizer')
os.makedirs(new_tokenizer_dir, exist_ok=True)
new_tokenizer_path = os.path.join(new_tokenizer_dir, 'vocab.txt')
with open(new_tokenizer_path, 'w') as f:
for v in asr_model2.tokenizer.tokenizer.get_vocab():
f.write(f"{v}\n")
# Add some new tokens too
f.write("^\n")
f.write("^^\n")
f.write("^^^\n")
assert os.path.exists(new_tokenizer_path)
# change vocabulary
asr_model2.change_vocabulary(new_tokenizer_dir, new_tokenizer_type='wpe')
assert asr_model.tokenizer.vocab_size != asr_model2.tokenizer.vocab_size
new_save_path = os.path.join(save_dir, 'temp2.nemo')
asr_model2.save_to(new_save_path)
asr_model3 = EncDecCTCModelBPE.restore_from(new_save_path)
assert isinstance(asr_model3, EncDecCTCModelBPE)
# Check if vocabulary size is same
assert asr_model2.tokenizer.tokenizer.vocab_size == asr_model3.tokenizer.tokenizer.vocab_size
assert asr_model2.vocab_path != asr_model3.vocab_path
# Model PT level checks
assert len(asr_model2.artifacts) == 1
finally:
os.rename(old_tokenizer_dir + '.bkp', old_tokenizer_dir)
@pytest.mark.unit
def test_EncDecCTCDatasetConfig_for_AudioToBPEDataset(self):
# ignore some additional arguments as dataclass is generic
IGNORE_ARGS = [
'is_tarred',
'num_workers',
'batch_size',
'tarred_audio_filepaths',
'shuffle',
'pin_memory',
'drop_last',
'tarred_shard_strategy',
'shuffle_n',
'parser',
'normalize',
'unk_index',
'pad_id',
'bos_id',
'eos_id',
'blank_index',
]
REMAP_ARGS = {'trim_silence': 'trim', 'labels': 'tokenizer'}
result = assert_dataclass_signature_match(
audio_to_text.AudioToBPEDataset,
configs.EncDecCTCDatasetConfig,
ignore_args=IGNORE_ARGS,
remap_args=REMAP_ARGS,
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
@pytest.mark.unit
def test_EncDecCTCDatasetConfig_for_TarredAudioToBPEDataset(self):
# ignore some additional arguments as dataclass is generic
IGNORE_ARGS = [
'is_tarred',
'num_workers',
'batch_size',
'shuffle',
'pin_memory',
'drop_last',
'parser',
'normalize',
'unk_index',
'pad_id',
'bos_id',
'eos_id',
'blank_index',
'global_rank',
'world_size',
]
REMAP_ARGS = {
'trim_silence': 'trim',
'tarred_audio_filepaths': 'audio_tar_filepaths',
'tarred_shard_strategy': 'shard_strategy',
'shuffle_n': 'shuffle',
'labels': 'tokenizer',
}
result = assert_dataclass_signature_match(
audio_to_text.TarredAudioToBPEDataset,
configs.EncDecCTCDatasetConfig,
ignore_args=IGNORE_ARGS,
remap_args=REMAP_ARGS,
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
tts/NeMo/tests/collections/asr/test_asr_ctcencdec_model.py
+305
View File
@@ -0,0 +1,305 @@
# Copyright (c) 2020, 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 copy
import pytest
import torch
from omegaconf import DictConfig, OmegaConf, open_dict
import nemo.collections.asr as nemo_asr
from nemo.collections.asr.data import audio_to_text
from nemo.collections.asr.models import EncDecCTCModel, configs
from nemo.utils.config_utils import assert_dataclass_signature_match, update_model_config
@pytest.fixture()
def asr_model():
preprocessor = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'}
encoder = {
'_target_': 'nemo.collections.asr.modules.ConvASREncoder',
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': 1024,
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': True,
'se': True,
'se_context_size': -1,
}
],
}
decoder = {
'_target_': 'nemo.collections.asr.modules.ConvASRDecoder',
'feat_in': 1024,
'num_classes': 28,
'vocabulary': [
' ',
'a',
'b',
'c',
'd',
'e',
'f',
'g',
'h',
'i',
'j',
'k',
'l',
'm',
'n',
'o',
'p',
'q',
'r',
's',
't',
'u',
'v',
'w',
'x',
'y',
'z',
"'",
],
}
modelConfig = DictConfig(
{'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
)
model_instance = EncDecCTCModel(cfg=modelConfig)
return model_instance
class TestEncDecCTCModel:
@pytest.mark.unit
def test_constructor(self, asr_model):
asr_model.train()
# TODO: make proper config and assert correct number of weights
# Check to/from config_dict:
confdict = asr_model.to_config_dict()
instance2 = EncDecCTCModel.from_config_dict(confdict)
assert isinstance(instance2, EncDecCTCModel)
@pytest.mark.unit
def test_forward(self, asr_model):
asr_model = asr_model.eval()
asr_model.preprocessor.featurizer.dither = 0.0
asr_model.preprocessor.featurizer.pad_to = 0
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
with torch.no_grad():
# batch size 1
logprobs_instance = []
for i in range(input_signal.size(0)):
logprobs_ins, _, _ = asr_model.forward(
input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
)
logprobs_instance.append(logprobs_ins)
print(len(logprobs_ins))
logprobs_instance = torch.cat(logprobs_instance, 0)
# batch size 4
logprobs_batch, _, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length)
assert logprobs_instance.shape == logprobs_batch.shape
diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch))
assert diff <= 1e-6
diff = torch.max(torch.abs(logprobs_instance - logprobs_batch))
assert diff <= 1e-6
@pytest.mark.unit
def test_vocab_change(self, asr_model):
old_vocab = copy.deepcopy(asr_model.decoder.vocabulary)
nw1 = asr_model.num_weights
asr_model.change_vocabulary(new_vocabulary=old_vocab)
# No change
assert nw1 == asr_model.num_weights
new_vocab = copy.deepcopy(old_vocab)
new_vocab.append('!')
new_vocab.append('$')
new_vocab.append('@')
asr_model.change_vocabulary(new_vocabulary=new_vocab)
# fully connected + bias
assert asr_model.num_weights == nw1 + 3 * (asr_model.decoder._feat_in + 1)
@pytest.mark.unit
def test_change_conv_asr_se_context_window(self, asr_model):
old_cfg = copy.deepcopy(asr_model.cfg)
asr_model.change_conv_asr_se_context_window(context_window=32) # 32 * 0.01s context
new_config = asr_model.cfg
assert old_cfg.encoder.jasper[0].se_context_size == -1
assert new_config.encoder.jasper[0].se_context_size == 32
for name, m in asr_model.encoder.named_modules():
if type(m).__class__.__name__ == 'SqueezeExcite':
assert m.context_window == 32
@pytest.mark.unit
def test_change_conv_asr_se_context_window_no_config_update(self, asr_model):
old_cfg = copy.deepcopy(asr_model.cfg)
asr_model.change_conv_asr_se_context_window(context_window=32, update_config=False) # 32 * 0.01s context
new_config = asr_model.cfg
assert old_cfg.encoder.jasper[0].se_context_size == -1
assert new_config.encoder.jasper[0].se_context_size == -1 # no change
for name, m in asr_model.encoder.named_modules():
if type(m).__class__.__name__ == 'SqueezeExcite':
assert m.context_window == 32
@pytest.mark.unit
def test_dataclass_instantiation(self, asr_model):
model_cfg = configs.EncDecCTCModelConfig()
# Update mandatory values
vocabulary = asr_model.decoder.vocabulary
model_cfg.model.labels = vocabulary
# Update encoder
model_cfg.model.encoder.activation = 'relu'
model_cfg.model.encoder.feat_in = 64
model_cfg.model.encoder.jasper = [
nemo_asr.modules.conv_asr.JasperEncoderConfig(
filters=1024,
repeat=1,
kernel=[1],
stride=[1],
dilation=[1],
dropout=0.0,
residual=False,
se=True,
se_context_size=-1,
)
]
# Update decoder
model_cfg.model.decoder.feat_in = 1024
model_cfg.model.decoder.num_classes = 28
model_cfg.model.decoder.vocabulary = vocabulary
# Construct the model
asr_cfg = OmegaConf.create({'model': asr_model.cfg})
model_cfg_v1 = update_model_config(model_cfg, asr_cfg)
new_model = EncDecCTCModel(cfg=model_cfg_v1.model)
assert new_model.num_weights == asr_model.num_weights
# trainer and exp manager should be there
# assert 'trainer' in model_cfg_v1
# assert 'exp_manager' in model_cfg_v1
# datasets and optim/sched should not be there after ModelPT.update_model_dataclass()
assert 'train_ds' not in model_cfg_v1.model
assert 'validation_ds' not in model_cfg_v1.model
assert 'test_ds' not in model_cfg_v1.model
assert 'optim' not in model_cfg_v1.model
# Construct the model, without dropping additional keys
asr_cfg = OmegaConf.create({'model': asr_model.cfg})
model_cfg_v2 = update_model_config(model_cfg, asr_cfg, drop_missing_subconfigs=False)
# Assert all components are in config
# assert 'trainer' in model_cfg_v2
# assert 'exp_manager' in model_cfg_v2
assert 'train_ds' in model_cfg_v2.model
assert 'validation_ds' in model_cfg_v2.model
assert 'test_ds' in model_cfg_v2.model
assert 'optim' in model_cfg_v2.model
# Remove extra components (optim and sched can be kept without issue)
with open_dict(model_cfg_v2.model):
model_cfg_v2.model.pop('train_ds')
model_cfg_v2.model.pop('validation_ds')
model_cfg_v2.model.pop('test_ds')
new_model = EncDecCTCModel(cfg=model_cfg_v2.model)
assert new_model.num_weights == asr_model.num_weights
# trainer and exp manager should be there
@pytest.mark.unit
def test_EncDecCTCDatasetConfig_for_AudioToCharDataset(self):
# ignore some additional arguments as dataclass is generic
IGNORE_ARGS = [
'is_tarred',
'num_workers',
'batch_size',
'tarred_audio_filepaths',
'shuffle',
'pin_memory',
'drop_last',
'tarred_shard_strategy',
'shuffle_n',
'use_start_end_token',
'use_start_end_token',
]
REMAP_ARGS = {'trim_silence': 'trim'}
result = assert_dataclass_signature_match(
audio_to_text.AudioToCharDataset,
configs.EncDecCTCDatasetConfig,
ignore_args=IGNORE_ARGS,
remap_args=REMAP_ARGS,
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
@pytest.mark.unit
def test_EncDecCTCDatasetConfig_for_TarredAudioToCharDataset(self):
# ignore some additional arguments as dataclass is generic
IGNORE_ARGS = [
'is_tarred',
'num_workers',
'batch_size',
'shuffle',
'pin_memory',
'drop_last',
'global_rank',
'world_size',
'use_start_end_token',
]
REMAP_ARGS = {
'trim_silence': 'trim',
'tarred_audio_filepaths': 'audio_tar_filepaths',
'tarred_shard_strategy': 'shard_strategy',
'shuffle_n': 'shuffle',
}
result = assert_dataclass_signature_match(
audio_to_text.TarredAudioToCharDataset,
configs.EncDecCTCDatasetConfig,
ignore_args=IGNORE_ARGS,
remap_args=REMAP_ARGS,
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
tts/NeMo/tests/collections/asr/test_asr_datasets.py
+127
View File
@@ -0,0 +1,127 @@
# Copyright (c) 2020, 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 copy
import os
import pytest
from omegaconf import OmegaConf
from nemo.collections.asr.data.audio_to_text import TarredAudioToBPEDataset, TarredAudioToCharDataset
from nemo.collections.asr.data.audio_to_text_dataset import inject_dataloader_value_from_model_config
from nemo.collections.common import tokenizers
from nemo.utils import logging
class TestASRDatasets:
labels = [
" ",
"a",
"b",
"c",
"d",
"e",
"f",
"g",
"h",
"i",
"j",
"k",
"l",
"m",
"n",
"o",
"p",
"q",
"r",
"s",
"t",
"u",
"v",
"w",
"x",
"y",
"z",
"'",
]
@pytest.mark.unit
def test_tarred_dataset(self, test_data_dir):
manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json'))
# Test braceexpand loading
tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
ds_braceexpand = TarredAudioToCharDataset(
audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, sample_rate=16000
)
assert len(ds_braceexpand) == 32
count = 0
for _ in ds_braceexpand:
count += 1
assert count == 32
# Test loading via list
tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)]
ds_list_load = TarredAudioToCharDataset(
audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, sample_rate=16000
)
count = 0
for _ in ds_list_load:
count += 1
assert count == 32
@pytest.mark.unit
def test_mismatch_in_model_dataloader_config(self, caplog):
logging._logger.propagate = True
caplog.set_level(logging.WARNING)
model_cfg = OmegaConf.create(dict(labels=OmegaConf.create(["a", "b", "c"])))
dataloader_cfg = OmegaConf.create(dict(labels=copy.deepcopy(self.labels)))
inject_dataloader_value_from_model_config(model_cfg, dataloader_cfg, key='labels')
assert (
"""`labels` is explicitly provided to the data loader, and is different from the `labels` provided at the model level config."""
in caplog.text
)
logging._logger.propagate = False
@pytest.mark.unit
def test_tarred_bpe_dataset(self, test_data_dir):
manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json'))
tokenizer_path = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
tokenizer = tokenizers.AutoTokenizer(pretrained_model_name='bert-base-cased', vocab_file=tokenizer_path)
# Test braceexpand loading
tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
ds_braceexpand = TarredAudioToBPEDataset(
audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, tokenizer=tokenizer, sample_rate=16000
)
assert len(ds_braceexpand) == 32
count = 0
for _ in ds_braceexpand:
count += 1
assert count == 32
# Test loading via list
tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)]
ds_list_load = TarredAudioToBPEDataset(
audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, tokenizer=tokenizer, sample_rate=16000
)
count = 0
for _ in ds_list_load:
count += 1
assert count == 32
tts/NeMo/tests/collections/asr/test_asr_exportables.py
+301
View File
@@ -0,0 +1,301 @@
# Copyright (c) 2020, 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 os
import tempfile
import onnx
import pytest
from omegaconf import DictConfig, ListConfig
from nemo.collections.asr.models import EncDecClassificationModel, EncDecCTCModel, EncDecSpeakerLabelModel
from nemo.collections.asr.modules import ConvASRDecoder, ConvASREncoder
class TestExportable:
@pytest.mark.run_only_on('GPU')
@pytest.mark.unit
def test_EncDecCTCModel_export_to_onnx(self):
model_config = DictConfig(
{
'preprocessor': DictConfig(self.preprocessor),
'encoder': DictConfig(self.encoder_dict),
'decoder': DictConfig(self.decoder_dict),
}
)
model = EncDecCTCModel(cfg=model_config)
with tempfile.TemporaryDirectory() as tmpdir:
filename = os.path.join(tmpdir, 'qn.onnx')
model.export(output=filename)
onnx_model = onnx.load(filename)
onnx.checker.check_model(onnx_model, full_check=True) # throws when failed
assert onnx_model.graph.input[0].name == 'audio_signal'
assert onnx_model.graph.output[0].name == 'logprobs'
@pytest.mark.run_only_on('GPU')
@pytest.mark.unit
def test_EncDecClassificationModel_export_to_onnx(self, speech_classification_model):
model = speech_classification_model.train()
with tempfile.TemporaryDirectory() as tmpdir:
filename = os.path.join(tmpdir, 'edc.onnx')
model.export(output=filename)
onnx_model = onnx.load(filename)
onnx.checker.check_model(onnx_model, full_check=True) # throws when failed
assert onnx_model.graph.input[0].name == 'audio_signal'
assert onnx_model.graph.output[0].name == 'logits'
@pytest.mark.run_only_on('GPU')
@pytest.mark.unit
def test_EncDecSpeakerLabelModel_export_to_onnx(self, speaker_label_model):
model = speaker_label_model.train()
with tempfile.TemporaryDirectory() as tmpdir:
filename = os.path.join(tmpdir, 'sl.onnx')
model.export(output=filename)
onnx_model = onnx.load(filename)
onnx.checker.check_model(onnx_model, full_check=True) # throws when failed
assert onnx_model.graph.input[0].name == 'audio_signal'
assert onnx_model.graph.output[0].name == 'logits'
@pytest.mark.run_only_on('GPU')
@pytest.mark.unit
def test_EncDecCitrinetModel_export_to_onnx(self, citrinet_model):
model = citrinet_model.train()
with tempfile.TemporaryDirectory() as tmpdir:
filename = os.path.join('.', 'citri.onnx')
model.export(output=filename)
onnx_model = onnx.load(filename)
onnx.checker.check_model(onnx_model, full_check=True) # throws when failed
assert onnx_model.graph.input[0].name == 'audio_signal'
assert onnx_model.graph.output[0].name == 'logprobs'
def setup_method(self):
self.preprocessor = {
'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',
'params': dict({}),
}
self.encoder_dict = {
'cls': 'nemo.collections.asr.modules.ConvASREncoder',
'params': {
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': 1024,
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': True,
'se': True,
'se_context_size': -1,
}
],
},
}
self.decoder_dict = {
'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
'params': {
'feat_in': 1024,
'num_classes': 28,
'vocabulary': [
' ',
'a',
'b',
'c',
'd',
'e',
'f',
'g',
'h',
'i',
'j',
'k',
'l',
'm',
'n',
'o',
'p',
'q',
'r',
's',
't',
'u',
'v',
'w',
'x',
'y',
'z',
"'",
],
},
}
@pytest.fixture()
def speech_classification_model():
preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
encoder = {
'cls': 'nemo.collections.asr.modules.ConvASREncoder',
'params': {
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': 32,
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': True,
'se': True,
'se_context_size': -1,
}
],
},
}
decoder = {
'cls': 'nemo.collections.asr.modules.ConvASRDecoderClassification',
'params': {'feat_in': 32, 'num_classes': 30,},
}
modelConfig = DictConfig(
{
'preprocessor': DictConfig(preprocessor),
'encoder': DictConfig(encoder),
'decoder': DictConfig(decoder),
'labels': ListConfig(["dummy_cls_{}".format(i + 1) for i in range(30)]),
}
)
model = EncDecClassificationModel(cfg=modelConfig)
return model
@pytest.fixture()
def speaker_label_model():
preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
encoder = {
'cls': 'nemo.collections.asr.modules.ConvASREncoder',
'params': {
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': 512,
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': False,
}
],
},
}
decoder = {
'cls': 'nemo.collections.asr.modules.SpeakerDecoder',
'params': {'feat_in': 512, 'num_classes': 2, 'pool_mode': 'xvector', 'emb_sizes': [1024]},
}
modelConfig = DictConfig(
{'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
)
speaker_model = EncDecSpeakerLabelModel(cfg=modelConfig)
return speaker_model
@pytest.fixture()
def citrinet_model():
preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
encoder = {
'cls': 'nemo.collections.asr.modules.ConvASREncoder',
'params': {
'feat_in': 80,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': 512,
'repeat': 1,
'kernel': [5],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': True,
'se': True,
'se_context_size': -1,
},
{
'filters': 512,
'repeat': 5,
'kernel': [11],
'stride': [2],
'dilation': [1],
'dropout': 0.1,
'residual': True,
'separable': True,
'se': True,
'se_context_size': -1,
'stride_last': True,
'residual_mode': 'stride_add',
},
{
'filters': 512,
'repeat': 5,
'kernel': [13],
'stride': [1],
'dilation': [1],
'dropout': 0.1,
'residual': True,
'separable': True,
'se': True,
'se_context_size': -1,
},
{
'filters': 640,
'repeat': 1,
'kernel': [41],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': True,
'separable': True,
'se': True,
'se_context_size': -1,
},
],
},
}
decoder = {
'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
'params': {'feat_in': 640, 'num_classes': 1024, 'vocabulary': list(chr(i % 28) for i in range(0, 1024))},
}
modelConfig = DictConfig(
{'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
)
citri_model = EncDecSpeakerLabelModel(cfg=modelConfig)
return citri_model
tts/NeMo/tests/collections/asr/test_asr_filterbankfeatures_seq_len.py
+99
View File
@@ -0,0 +1,99 @@
# Copyright (c) 2020, 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 librosa
import numpy as np
import pytest
import torch
from nemo.collections.asr.parts.preprocessing.features import FilterbankFeatures
class TestFilterbankFeatures:
@pytest.mark.unit
def test_seq_len(self):
fb_module = FilterbankFeatures(exact_pad=False, pad_to=1)
test_1 = torch.randn(1, 800)
test_1_len = torch.tensor([800])
fb_spec, fb_len = fb_module(test_1, test_1_len)
assert fb_spec.shape[2] == fb_len[0], f"{fb_spec.shape} != {fb_len}"
librosa_spec = librosa.stft(test_1.cpu().detach().numpy().squeeze(), n_fft=512, hop_length=160, win_length=320)
assert librosa_spec.shape[1] == fb_spec.shape[2], f"{librosa_spec.shape} != {fb_spec.shape}"
@pytest.mark.unit
def test_random_stft_sizes(self):
for _ in range(5):
nfft = 2 ** np.random.randint(7, 12)
window_size = np.random.randint(100, nfft)
hop_size = np.random.randint(64, window_size)
fb_module = FilterbankFeatures(
exact_pad=False,
pad_to=1,
n_fft=nfft,
n_window_size=window_size,
n_window_stride=hop_size,
normalize=False,
)
audio_length = np.random.randint(nfft, 2 ** 16)
test_1 = torch.randn(1, audio_length)
test_1_len = torch.tensor([audio_length])
fb_spec, fb_len = fb_module(test_1, test_1_len)
assert (
fb_spec.shape[2] == fb_len[0]
), f"{fb_spec.shape} != {fb_len}: {nfft}, {window_size}, {hop_size}, {audio_length}"
librosa_spec = librosa.stft(
test_1.cpu().detach().numpy().squeeze(), n_fft=nfft, hop_length=hop_size, win_length=window_size
)
assert (
librosa_spec.shape[1] == fb_spec.shape[2]
), f"{librosa_spec.shape} != {fb_spec.shape}: {nfft}, {window_size}, {hop_size}, {audio_length}"
@pytest.mark.unit
def test_random_stft_sizes_exact_pad(self):
for _ in range(5):
nfft = 2 ** np.random.randint(7, 12)
window_size = np.random.randint(100, nfft)
hop_size = np.random.randint(64, window_size)
if hop_size % 2 == 1:
hop_size = hop_size - 1
fb_module = FilterbankFeatures(
exact_pad=True,
pad_to=1,
n_fft=nfft,
n_window_size=window_size,
n_window_stride=hop_size,
normalize=False,
)
audio_length = np.random.randint(nfft, 2 ** 16)
test_1 = torch.randn(1, audio_length)
test_1_len = torch.tensor([audio_length])
fb_spec, fb_len = fb_module(test_1, test_1_len)
assert (
fb_spec.shape[2] == fb_len[0]
), f"{fb_spec.shape} != {fb_len}: {nfft}, {window_size}, {hop_size}, {audio_length}"
test_2 = test_1.cpu().detach().numpy().squeeze()
test_2 = np.pad(test_2, int((nfft - hop_size) // 2), mode="reflect")
librosa_spec = librosa.stft(test_2, n_fft=nfft, hop_length=hop_size, win_length=window_size, center=False,)
assert (
fb_spec.shape[2] == librosa_spec.shape[1]
), f"{fb_spec.shape} != {librosa_spec.shape}: {nfft}, {window_size}, {hop_size}, {audio_length}"
assert (
fb_spec.shape[2] == audio_length // hop_size
), f"{fb_spec.shape}, {nfft}, {window_size}, {hop_size}, {audio_length}, {audio_length // hop_size}"
tts/NeMo/tests/collections/asr/test_asr_metrics.py
+172
View File
@@ -0,0 +1,172 @@
# Copyright (c) 2020, 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 random
import pytest
import torch
from nemo.collections.asr.metrics.wer import WER, word_error_rate
from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
class TestWordErrorRate:
vocabulary = [
" ",
"a",
"b",
"c",
"d",
"e",
"f",
"g",
"h",
"i",
"j",
"k",
"l",
"m",
"n",
"o",
"p",
"q",
"r",
"s",
"t",
"u",
"v",
"w",
"x",
"y",
"z",
"'",
]
def __string_to_ctc_tensor(self, txt: str) -> torch.Tensor:
# This function emulates how CTC output could like for txt
blank_id = len(self.vocabulary)
char_to_ind = dict([(self.vocabulary[i], i) for i in range(len(self.vocabulary))])
string_in_id_form = [char_to_ind[c] for c in txt]
ctc_list = []
prev_id = -1
for c in string_in_id_form:
# when character is repeated we need to insert CTC blank symbol
if c != prev_id:
ctc_list.append(c)
else:
ctc_list.append(blank_id)
ctc_list.append(c)
prev_id = c
return torch.Tensor(ctc_list).unsqueeze(0)
def __reference_string_to_tensor(self, txt: str) -> torch.Tensor:
# Reference tensors aren't produced by CTC logic
char_to_ind = dict([(self.vocabulary[i], i) for i in range(len(self.vocabulary))])
string_in_id_form = [char_to_ind[c] for c in txt]
return torch.Tensor(string_in_id_form).unsqueeze(0)
def get_wer(self, wer, prediction: str, reference: str):
wer(
predictions=self.__string_to_ctc_tensor(prediction),
targets=self.__reference_string_to_tensor(reference),
target_lengths=torch.tensor([len(reference)]),
)
res, _, _ = wer.compute()
res = res.detach().cpu()
# return res[0] / res[1]
return res.item()
@pytest.mark.unit
def test_wer_function(self):
assert word_error_rate(hypotheses=['cat'], references=['cot']) == 1.0
assert word_error_rate(hypotheses=['GPU'], references=['G P U']) == 1.0
assert word_error_rate(hypotheses=['G P U'], references=['GPU']) == 3.0
assert word_error_rate(hypotheses=['ducati motorcycle'], references=['motorcycle']) == 1.0
assert word_error_rate(hypotheses=['ducati motorcycle'], references=['ducuti motorcycle']) == 0.5
assert word_error_rate(hypotheses=['a B c'], references=['a b c']) == 1.0 / 3.0
@pytest.mark.unit
def test_wer_metric_simple(self):
wer = WER(vocabulary=self.vocabulary, batch_dim_index=0, use_cer=False, ctc_decode=True)
assert self.get_wer(wer, 'cat', 'cot') == 1.0
assert self.get_wer(wer, 'gpu', 'g p u') == 1.0
assert self.get_wer(wer, 'g p u', 'gpu') == 3.0
assert self.get_wer(wer, 'ducati motorcycle', 'motorcycle') == 1.0
assert self.get_wer(wer, 'ducati motorcycle', 'ducuti motorcycle') == 0.5
assert abs(self.get_wer(wer, 'a f c', 'a b c') - 1.0 / 3.0) < 1e-6
@pytest.mark.unit
def test_wer_metric_randomized(self):
"""This test relies on correctness of word_error_rate function."""
def __randomString(N):
return ''.join(random.choice(''.join(self.vocabulary)) for i in range(N))
wer = WER(vocabulary=self.vocabulary, batch_dim_index=0, use_cer=False, ctc_decode=True)
for test_id in range(256):
n1 = random.randint(1, 512)
n2 = random.randint(1, 512)
s1 = __randomString(n1)
s2 = __randomString(n2)
# skip empty strings as reference
if s2.strip():
assert (
abs(
self.get_wer(wer, prediction=s1, reference=s2)
- word_error_rate(hypotheses=[s1], references=[s2])
)
< 1e-6
)
@pytest.mark.unit
def test_wer_metric_decode(self):
wer = WER(vocabulary=self.vocabulary, batch_dim_index=0, use_cer=False, ctc_decode=True)
tokens = self.__string_to_ctc_tensor('cat')[0].int().numpy().tolist()
assert tokens == [3, 1, 20]
tokens_decoded = wer.decode_ids_to_tokens(tokens)
assert tokens_decoded == ['c', 'a', 't']
str_decoded = wer.decode_tokens_to_str(tokens)
assert str_decoded == 'cat'
@pytest.mark.unit
def test_wer_metric_return_hypothesis(self):
wer = WER(vocabulary=self.vocabulary, batch_dim_index=0, use_cer=False, ctc_decode=True)
tensor = self.__string_to_ctc_tensor('cat').int()
# pass batchsize 1 tensor, get back list of length 1 Hypothesis
hyp = wer.ctc_decoder_predictions_tensor(tensor, return_hypotheses=True)
hyp = hyp[0]
assert isinstance(hyp, Hypothesis)
assert hyp.y_sequence is None
assert hyp.score == -1.0
assert hyp.text == 'cat'
assert hyp.alignments == [3, 1, 20]
assert hyp.length == 0
length = torch.tensor([tensor.shape[-1]], dtype=torch.long)
# pass batchsize 1 tensor, get back list of length 1 Hypothesis [add length info]
hyp = wer.ctc_decoder_predictions_tensor(tensor, predictions_len=length, return_hypotheses=True)
hyp = hyp[0]
assert isinstance(hyp, Hypothesis)
assert hyp.length == 3
tts/NeMo/tests/collections/asr/test_asr_modules.py
+325
View File
@@ -0,0 +1,325 @@
# Copyright (c) 2020, 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 omegaconf import OmegaConf
from nemo.collections.asr import modules
from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
from nemo.core.utils import numba_utils
from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
from nemo.utils import config_utils, logging
class TestASRModulesBasicTests:
@pytest.mark.unit
def test_AudioToMelSpectrogramPreprocessor1(self):
# Test 1 that should test the pure stft implementation as much as possible
instance1 = modules.AudioToMelSpectrogramPreprocessor(
dither=0, stft_conv=False, mag_power=1.0, normalize=False, preemph=0.0, log=False, pad_to=0
)
instance2 = modules.AudioToMelSpectrogramPreprocessor(
dither=0, stft_conv=True, mag_power=1.0, normalize=False, preemph=0.0, log=False, pad_to=0
)
# Ensure that the two functions behave similarily
for _ in range(10):
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
res1, length1 = instance1(input_signal=input_signal, length=length)
res2, length2 = instance2(input_signal=input_signal, length=length)
for len1, len2 in zip(length1, length2):
assert len1 == len2
assert res1.shape == res2.shape
diff = torch.mean(torch.abs(res1 - res2))
assert diff <= 1e-3
diff = torch.max(torch.abs(res1 - res2))
assert diff <= 1e-2
@pytest.mark.unit
def test_AudioToMelSpectrogramPreprocessor_config(self):
# Test that dataclass matches signature of module
result = config_utils.assert_dataclass_signature_match(
modules.AudioToMelSpectrogramPreprocessor,
modules.audio_preprocessing.AudioToMelSpectrogramPreprocessorConfig,
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
@pytest.mark.unit
def test_AudioToMelSpectrogramPreprocessor_batch(self):
# Test 1 that should test the pure stft implementation as much as possible
instance1 = modules.AudioToMelSpectrogramPreprocessor(normalize="per_feature", dither=0, pad_to=0)
# Ensure that the two functions behave similarily
for _ in range(10):
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
with torch.no_grad():
# batch size 1
res_instance, length_instance = [], []
for i in range(input_signal.size(0)):
res_ins, length_ins = instance1(input_signal=input_signal[i : i + 1], length=length[i : i + 1])
res_instance.append(res_ins)
length_instance.append(length_ins)
res_instance = torch.cat(res_instance, 0)
length_instance = torch.cat(length_instance, 0)
# batch size 4
res_batch, length_batch = instance1(input_signal=input_signal, length=length)
assert res_instance.shape == res_batch.shape
assert length_instance.shape == length_batch.shape
diff = torch.mean(torch.abs(res_instance - res_batch))
assert diff <= 1e-3
diff = torch.max(torch.abs(res_instance - res_batch))
assert diff <= 1e-3
@pytest.mark.unit
def test_AudioToMelSpectrogramPreprocessor2(self):
# Test 2 that should test the stft implementation as used in ASR models
instance1 = modules.AudioToMelSpectrogramPreprocessor(dither=0, stft_conv=False)
instance2 = modules.AudioToMelSpectrogramPreprocessor(dither=0, stft_conv=True)
# Ensure that the two functions behave similarily
for _ in range(5):
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
res1, length1 = instance1(input_signal=input_signal, length=length)
res2, length2 = instance2(input_signal=input_signal, length=length)
for len1, len2 in zip(length1, length2):
assert len1 == len2
assert res1.shape == res2.shape
diff = torch.mean(torch.abs(res1 - res2))
assert diff <= 3e-3
diff = torch.max(torch.abs(res1 - res2))
assert diff <= 3
@pytest.mark.unit
def test_SpectrogramAugmentationr(self):
# Make sure constructor works
instance1 = modules.SpectrogramAugmentation(
freq_masks=10, time_masks=3, rect_masks=3, use_numba_spec_augment=False
)
assert isinstance(instance1, modules.SpectrogramAugmentation)
# Make sure forward doesn't throw with expected input
instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0)
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
res0 = instance0(input_signal=input_signal, length=length)
res = instance1(input_spec=res0[0])
assert res.shape == res0[0].shape
@pytest.mark.unit
@pytest.mark.run_only_on('GPU')
def test_SpectrogramAugmentationr_numba_kernel(self, caplog):
numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)
logging._logger.propagate = True
original_verbosity = logging.get_verbosity()
logging.set_verbosity(logging.DEBUG)
caplog.set_level(logging.DEBUG)
# Make sure constructor works
instance1 = modules.SpectrogramAugmentation(
freq_masks=10, time_masks=3, rect_masks=3, use_numba_spec_augment=True
)
assert isinstance(instance1, modules.SpectrogramAugmentation)
# Make sure forward doesn't throw with expected input
instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0)
input_signal = torch.randn(size=(8, 512))
length = torch.randint(low=161, high=500, size=[8])
res0 = instance0(input_signal=input_signal, length=length)
res = instance1(input_spec=res0[0], length=length)
assert res.shape == res0[0].shape
# check tha numba kernel debug message indicates that it is available for use
assert """Numba SpecAugment kernel is available""" in caplog.text
logging._logger.propagate = False
logging.set_verbosity(original_verbosity)
@pytest.mark.unit
def test_SpectrogramAugmentationr_config(self):
# Test that dataclass matches signature of module
result = config_utils.assert_dataclass_signature_match(
modules.SpectrogramAugmentation, modules.audio_preprocessing.SpectrogramAugmentationConfig,
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
@pytest.mark.unit
def test_CropOrPadSpectrogramAugmentation(self):
# Make sure constructor works
audio_length = 128
instance1 = modules.CropOrPadSpectrogramAugmentation(audio_length=audio_length)
assert isinstance(instance1, modules.CropOrPadSpectrogramAugmentation)
# Make sure forward doesn't throw with expected input
instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0)
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
res0 = instance0(input_signal=input_signal, length=length)
res, new_length = instance1(input_signal=res0[0], length=length)
assert res.shape == torch.Size([4, 64, audio_length])
assert all(new_length == torch.tensor([128] * 4))
@pytest.mark.unit
def test_CropOrPadSpectrogramAugmentation_config(self):
# Test that dataclass matches signature of module
result = config_utils.assert_dataclass_signature_match(
modules.CropOrPadSpectrogramAugmentation,
modules.audio_preprocessing.CropOrPadSpectrogramAugmentationConfig,
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
@pytest.mark.unit
def test_RNNTDecoder(self):
vocab = list(range(10))
vocab = [str(x) for x in vocab]
vocab_size = len(vocab)
pred_config = OmegaConf.create(
{
'_target_': 'nemo.collections.asr.modules.RNNTDecoder',
'prednet': {'pred_hidden': 32, 'pred_rnn_layers': 1,},
'vocab_size': vocab_size,
'blank_as_pad': True,
}
)
prednet = modules.RNNTDecoder.from_config_dict(pred_config)
# num params
pred_hidden = pred_config.prednet.pred_hidden
embed = (vocab_size + 1) * pred_hidden # embedding with blank
rnn = (
2 * 4 * (pred_hidden * pred_hidden + pred_hidden)
) # (ih + hh) * (ifco gates) * (indim * hiddendim + bias)
assert prednet.num_weights == (embed + rnn)
# State initialization
x_ = torch.zeros(4, dtype=torch.float32)
states = prednet.initialize_state(x_)
for state_i in states:
assert state_i.dtype == x_.dtype
assert state_i.device == x_.device
assert state_i.shape[1] == len(x_)
# Blank hypotheses test
blank = vocab_size
hyp = Hypothesis(score=0.0, y_sequence=[blank])
cache = {}
pred, states, _ = prednet.score_hypothesis(hyp, cache)
assert pred.shape == torch.Size([1, 1, pred_hidden])
assert len(states) == 2
for state_i in states:
assert state_i.dtype == pred.dtype
assert state_i.device == pred.device
assert state_i.shape[1] == len(pred)
# Blank stateless predict
g, states = prednet.predict(y=None, state=None, add_sos=False, batch_size=1)
assert g.shape == torch.Size([1, 1, pred_hidden])
assert len(states) == 2
for state_i in states:
assert state_i.dtype == g.dtype
assert state_i.device == g.device
assert state_i.shape[1] == len(g)
# Blank stateful predict
g, states2 = prednet.predict(y=None, state=states, add_sos=False, batch_size=1)
assert g.shape == torch.Size([1, 1, pred_hidden])
assert len(states2) == 2
for state_i, state_j in zip(states, states2):
assert (state_i - state_j).square().sum().sqrt() > 0.0
# Predict with token and state
token = torch.full([1, 1], fill_value=0, dtype=torch.long)
g, states = prednet.predict(y=token, state=states2, add_sos=False, batch_size=None)
assert g.shape == torch.Size([1, 1, pred_hidden])
assert len(states) == 2
# Predict with blank token and no state
token = torch.full([1, 1], fill_value=blank, dtype=torch.long)
g, states = prednet.predict(y=token, state=None, add_sos=False, batch_size=None)
assert g.shape == torch.Size([1, 1, pred_hidden])
assert len(states) == 2
@pytest.mark.unit
def test_RNNTJoint(self):
vocab = list(range(10))
vocab = [str(x) for x in vocab]
vocab_size = len(vocab)
batchsize = 4
encoder_hidden = 64
pred_hidden = 32
joint_hidden = 16
joint_cfg = OmegaConf.create(
{
'_target_': 'nemo.collections.asr.modules.RNNTJoint',
'num_classes': vocab_size,
'vocabulary': vocab,
'jointnet': {
'encoder_hidden': encoder_hidden,
'pred_hidden': pred_hidden,
'joint_hidden': joint_hidden,
'activation': 'relu',
},
}
)
jointnet = modules.RNNTJoint.from_config_dict(joint_cfg)
enc = torch.zeros(batchsize, encoder_hidden, 48) # [B, D1, T]
dec = torch.zeros(batchsize, pred_hidden, 24) # [B, D2, U]
# forward call test
out = jointnet(encoder_outputs=enc, decoder_outputs=dec)
assert out.shape == torch.Size([batchsize, 48, 24, vocab_size + 1]) # [B, T, U, V + 1]
# joint() step test
enc2 = enc.transpose(1, 2) # [B, T, D1]
dec2 = dec.transpose(1, 2) # [B, U, D2]
out2 = jointnet.joint(enc2, dec2) # [B, T, U, V + 1]
assert (out - out2).abs().sum() <= 1e-5
# assert vocab size
assert jointnet.num_classes_with_blank == vocab_size + 1
tts/NeMo/tests/collections/asr/test_asr_regression_model.py
+95
View File
@@ -0,0 +1,95 @@
# Copyright (c) 2021, 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 os
import pytest
from omegaconf import DictConfig
from nemo.collections.asr.models.classification_models import EncDecRegressionModel
@pytest.fixture()
def speech_regression_model():
preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
encoder = {
'cls': 'nemo.collections.asr.modules.ConvASREncoder',
'params': {
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': 32,
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': True,
'se': True,
'se_context_size': -1,
}
],
},
}
decoder = {
'cls': 'nemo.collections.asr.modules.conv_asr.ConvASRDecoderClassification',
'params': {'feat_in': 32, 'return_logits': True, 'num_classes': 1},
}
modelConfig = DictConfig(
{
'preprocessor': DictConfig(preprocessor),
'encoder': DictConfig(encoder),
'decoder': DictConfig(decoder),
'labels': None,
'is_regression_task': True,
}
)
model = EncDecRegressionModel(cfg=modelConfig)
return model
class TestEncDecRegressionModel:
@pytest.mark.unit
def test_constructor(self, speech_regression_model):
asr_model = speech_regression_model.train()
conv_cnt = (64 * 32 * 1 + 32) + (64 * 1 * 1 + 32) # separable kernel + bias + pointwise kernel + bias
bn_cnt = (4 * 32) * 2 # 2 * moving averages
dec_cnt = 32 * 1 + 1 # fc + bias
param_count = conv_cnt + bn_cnt + dec_cnt
assert asr_model.num_weights == param_count
# Check to/from config_dict:
confdict = asr_model.to_config_dict()
instance2 = EncDecRegressionModel.from_config_dict(confdict)
assert isinstance(instance2, EncDecRegressionModel)
@pytest.mark.unit
def test_transcription(self, speech_regression_model, test_data_dir):
audio_filenames = ['an22-flrp-b.wav', 'an90-fbbh-b.wav']
audio_paths = [os.path.join(test_data_dir, "asr", "train", "an4", "wav", fp) for fp in audio_filenames]
model = speech_regression_model.eval()
# Test Top 1 classification transcription
results = model.transcribe(audio_paths, batch_size=2)
assert len(results) == 2
tts/NeMo/tests/collections/asr/test_asr_rnnt_encdec_model.py
+276
View File
@@ -0,0 +1,276 @@
# Copyright (c) 2020, 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 copy
import pytest
import torch
from omegaconf import DictConfig, ListConfig
from nemo.collections.asr.models import EncDecRNNTModel
from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode
from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
from nemo.core.utils import numba_utils
from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
from nemo.utils.config_utils import assert_dataclass_signature_match
NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
@pytest.fixture()
def asr_model():
preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
# fmt: off
labels = [' ', 'a', 'b', 'c', 'd', 'e', 'f',
'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o',
'p', 'q', 'r', 's', 't', 'u', 'v', 'w',
'x', 'y', 'z', "'",
]
# fmt: on
model_defaults = {'enc_hidden': 1024, 'pred_hidden': 64}
encoder = {
'cls': 'nemo.collections.asr.modules.ConvASREncoder',
'params': {
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': model_defaults['enc_hidden'],
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': True,
'se': True,
'se_context_size': -1,
}
],
},
}
decoder = {
'_target_': 'nemo.collections.asr.modules.RNNTDecoder',
'prednet': {'pred_hidden': model_defaults['pred_hidden'], 'pred_rnn_layers': 1},
}
joint = {
'_target_': 'nemo.collections.asr.modules.RNNTJoint',
'jointnet': {'joint_hidden': 32, 'activation': 'relu'},
}
decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}}
modelConfig = DictConfig(
{
'labels': ListConfig(labels),
'preprocessor': DictConfig(preprocessor),
'model_defaults': DictConfig(model_defaults),
'encoder': DictConfig(encoder),
'decoder': DictConfig(decoder),
'joint': DictConfig(joint),
'decoding': DictConfig(decoding),
}
)
model_instance = EncDecRNNTModel(cfg=modelConfig)
return model_instance
class TestEncDecRNNTModel:
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_constructor(self, asr_model):
asr_model.train()
# TODO: make proper config and assert correct number of weights
# Check to/from config_dict:
confdict = asr_model.to_config_dict()
instance2 = EncDecRNNTModel.from_config_dict(confdict)
assert isinstance(instance2, EncDecRNNTModel)
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_forward(self, asr_model):
asr_model = asr_model.eval()
asr_model.preprocessor.featurizer.dither = 0.0
asr_model.preprocessor.featurizer.pad_to = 0
asr_model.compute_eval_loss = False
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
with torch.no_grad():
# batch size 1
logprobs_instance = []
for i in range(input_signal.size(0)):
logprobs_ins, _ = asr_model.forward(
input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
)
logprobs_instance.append(logprobs_ins)
logprobs_instance = torch.cat(logprobs_instance, 0)
# batch size 4
logprobs_batch, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length)
assert logprobs_instance.shape == logprobs_batch.shape
diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch))
assert diff <= 1e-6
diff = torch.max(torch.abs(logprobs_instance - logprobs_batch))
assert diff <= 1e-6
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_vocab_change(self, asr_model):
old_vocab = copy.deepcopy(asr_model.joint.vocabulary)
nw1 = asr_model.num_weights
asr_model.change_vocabulary(new_vocabulary=old_vocab)
# No change
assert nw1 == asr_model.num_weights
new_vocab = copy.deepcopy(old_vocab)
new_vocab.append('!')
new_vocab.append('$')
new_vocab.append('@')
asr_model.change_vocabulary(new_vocabulary=new_vocab)
# fully connected + bias
# rnn embedding + joint + bias
pred_embedding = 3 * (asr_model.decoder.pred_hidden)
joint_joint = 3 * (asr_model.joint.joint_hidden + 1)
assert asr_model.num_weights == (nw1 + (pred_embedding + joint_joint))
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_change_conv_asr_se_context_window(self, asr_model):
old_cfg = copy.deepcopy(asr_model.cfg)
asr_model.change_conv_asr_se_context_window(context_window=32) # 32 * 0.01s context
new_config = asr_model.cfg
assert old_cfg.encoder.jasper[0].se_context_size == -1
assert new_config.encoder.jasper[0].se_context_size == 32
for name, m in asr_model.encoder.named_modules():
if type(m).__class__.__name__ == 'SqueezeExcite':
assert m.context_window == 32
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_change_conv_asr_se_context_window_no_config_update(self, asr_model):
old_cfg = copy.deepcopy(asr_model.cfg)
asr_model.change_conv_asr_se_context_window(context_window=32, update_config=False) # 32 * 0.01s context
new_config = asr_model.cfg
assert old_cfg.encoder.jasper[0].se_context_size == -1
assert new_config.encoder.jasper[0].se_context_size == -1 # no change
for name, m in asr_model.encoder.named_modules():
if type(m).__class__.__name__ == 'SqueezeExcite':
assert m.context_window == 32
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_decoding_change(self, asr_model):
assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
new_strategy = DictConfig({})
new_strategy.strategy = 'greedy'
new_strategy.greedy = DictConfig({'max_symbols': 10})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer)
new_strategy = DictConfig({})
new_strategy.strategy = 'beam'
new_strategy.beam = DictConfig({'beam_size': 1})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
assert asr_model.decoding.decoding.search_type == "default"
new_strategy = DictConfig({})
new_strategy.strategy = 'beam'
new_strategy.beam = DictConfig({'beam_size': 2})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
assert asr_model.decoding.decoding.search_type == "default"
new_strategy = DictConfig({})
new_strategy.strategy = 'tsd'
new_strategy.beam = DictConfig({'beam_size': 2})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
assert asr_model.decoding.decoding.search_type == "tsd"
new_strategy = DictConfig({})
new_strategy.strategy = 'alsd'
new_strategy.beam = DictConfig({'beam_size': 2})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
assert asr_model.decoding.decoding.search_type == "alsd"
@pytest.mark.unit
def test_GreedyRNNTInferConfig(self):
IGNORE_ARGS = ['decoder_model', 'joint_model', 'blank_index']
result = assert_dataclass_signature_match(
greedy_decode.GreedyRNNTInfer, greedy_decode.GreedyRNNTInferConfig, ignore_args=IGNORE_ARGS
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
@pytest.mark.unit
def test_GreedyBatchedRNNTInferConfig(self):
IGNORE_ARGS = ['decoder_model', 'joint_model', 'blank_index']
result = assert_dataclass_signature_match(
greedy_decode.GreedyBatchedRNNTInfer, greedy_decode.GreedyBatchedRNNTInferConfig, ignore_args=IGNORE_ARGS
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
@pytest.mark.unit
def test_BeamRNNTInferConfig(self):
IGNORE_ARGS = ['decoder_model', 'joint_model', 'blank_index']
result = assert_dataclass_signature_match(
beam_decode.BeamRNNTInfer, beam_decode.BeamRNNTInferConfig, ignore_args=IGNORE_ARGS
)
signatures_match, cls_subset, dataclass_subset = result
assert signatures_match
assert cls_subset is None
assert dataclass_subset is None
tts/NeMo/tests/collections/asr/test_asr_rnnt_encoder_model_bpe.py
+246
View File
@@ -0,0 +1,246 @@
# Copyright (c) 2020, 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 os
import shutil
import tempfile
import pytest
import torch
from omegaconf import DictConfig
from nemo.collections.asr.models.rnnt_bpe_models import EncDecRNNTBPEModel
from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode
from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
from nemo.collections.common import tokenizers
from nemo.core.utils import numba_utils
from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
@pytest.fixture()
def asr_model(test_data_dir):
preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
model_defaults = {'enc_hidden': 1024, 'pred_hidden': 64}
encoder = {
'cls': 'nemo.collections.asr.modules.ConvASREncoder',
'params': {
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': model_defaults['enc_hidden'],
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': True,
'se': True,
'se_context_size': -1,
}
],
},
}
decoder = {
'_target_': 'nemo.collections.asr.modules.RNNTDecoder',
'prednet': {'pred_hidden': model_defaults['pred_hidden'], 'pred_rnn_layers': 1,},
}
joint = {
'_target_': 'nemo.collections.asr.modules.RNNTJoint',
'jointnet': {'joint_hidden': 32, 'activation': 'relu',},
}
decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}}
tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'}
modelConfig = DictConfig(
{
'preprocessor': DictConfig(preprocessor),
'model_defaults': DictConfig(model_defaults),
'encoder': DictConfig(encoder),
'decoder': DictConfig(decoder),
'joint': DictConfig(joint),
'tokenizer': DictConfig(tokenizer),
'decoding': DictConfig(decoding),
}
)
model_instance = EncDecRNNTBPEModel(cfg=modelConfig)
return model_instance
class TestEncDecRNNTBPEModel:
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_constructor(self, asr_model):
asr_model.train()
# TODO: make proper config and assert correct number of weights
# Check to/from config_dict:
confdict = asr_model.to_config_dict()
instance2 = EncDecRNNTBPEModel.from_config_dict(confdict)
assert isinstance(instance2, EncDecRNNTBPEModel)
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_forward(self, asr_model):
asr_model = asr_model.eval()
asr_model.preprocessor.featurizer.dither = 0.0
asr_model.preprocessor.featurizer.pad_to = 0
asr_model.compute_eval_loss = False
input_signal = torch.randn(size=(4, 512))
length = torch.randint(low=161, high=500, size=[4])
with torch.no_grad():
# batch size 1
logprobs_instance = []
for i in range(input_signal.size(0)):
logprobs_ins, _ = asr_model.forward(
input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1]
)
logprobs_instance.append(logprobs_ins)
logits_instance = torch.cat(logprobs_instance, 0)
# batch size 4
logprobs_batch, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length)
assert logits_instance.shape == logprobs_batch.shape
diff = torch.mean(torch.abs(logits_instance - logprobs_batch))
assert diff <= 1e-6
diff = torch.max(torch.abs(logits_instance - logprobs_batch))
assert diff <= 1e-6
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_save_restore_artifact(self, asr_model):
asr_model.train()
with tempfile.TemporaryDirectory() as tmp_dir:
path = os.path.join(tmp_dir, 'rnnt_bpe.nemo')
asr_model.save_to(path)
new_model = EncDecRNNTBPEModel.restore_from(path)
assert isinstance(new_model, type(asr_model))
assert new_model.vocab_path.endswith('_vocab.txt')
assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_save_restore_artifact_spe(self, asr_model, test_data_dir):
asr_model.train()
with tempfile.TemporaryDirectory() as tmpdir:
tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128")
asr_model.change_vocabulary(new_tokenizer_dir=tokenizer_dir, new_tokenizer_type='bpe')
save_path = os.path.join(tmpdir, 'ctc_bpe.nemo')
asr_model.train()
asr_model.save_to(save_path)
new_model = EncDecRNNTBPEModel.restore_from(save_path)
assert isinstance(new_model, type(asr_model))
assert isinstance(new_model.tokenizer, tokenizers.SentencePieceTokenizer)
assert new_model.model_path.endswith('_tokenizer.model')
assert new_model.vocab_path.endswith('_vocab.txt')
assert new_model.spe_vocab_path.endswith('_tokenizer.vocab')
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_vocab_change(self, test_data_dir, asr_model):
with tempfile.TemporaryDirectory() as tmpdir:
old_tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt')
new_tokenizer_dir = os.path.join(tmpdir, 'tokenizer')
os.makedirs(new_tokenizer_dir, exist_ok=True)
shutil.copy2(old_tokenizer_dir, new_tokenizer_dir)
nw1 = asr_model.num_weights
asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
# No change
assert nw1 == asr_model.num_weights
with open(os.path.join(new_tokenizer_dir, 'vocab.txt'), 'a+') as f:
f.write("!\n")
f.write('$\n')
f.write('@\n')
asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe')
# rnn embedding + joint + bias
pred_embedding = 3 * (asr_model.decoder.pred_hidden)
joint_joint = 3 * (asr_model.joint.joint_hidden + 1)
assert asr_model.num_weights == (nw1 + (pred_embedding + joint_joint))
@pytest.mark.skipif(
not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.',
)
@pytest.mark.unit
def test_decoding_change(self, asr_model):
assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer)
new_strategy = DictConfig({})
new_strategy.strategy = 'greedy'
new_strategy.greedy = DictConfig({'max_symbols': 10})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer)
new_strategy = DictConfig({})
new_strategy.strategy = 'beam'
new_strategy.beam = DictConfig({'beam_size': 1})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
assert asr_model.decoding.decoding.search_type == "default"
new_strategy = DictConfig({})
new_strategy.strategy = 'beam'
new_strategy.beam = DictConfig({'beam_size': 2})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
assert asr_model.decoding.decoding.search_type == "default"
new_strategy = DictConfig({})
new_strategy.strategy = 'tsd'
new_strategy.beam = DictConfig({'beam_size': 2})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
assert asr_model.decoding.decoding.search_type == "tsd"
new_strategy = DictConfig({})
new_strategy.strategy = 'alsd'
new_strategy.beam = DictConfig({'beam_size': 2})
asr_model.change_decoding_strategy(decoding_cfg=new_strategy)
assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer)
assert asr_model.decoding.decoding.search_type == "alsd"
tts/NeMo/tests/collections/asr/test_jasper_block.py
+295
View File
@@ -0,0 +1,295 @@
# Copyright (c) 2021, 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.parts.submodules import jasper
class TestJasperBlock:
def jasper_base_config(self, **kwargs):
base = dict(
inplanes=16,
planes=8,
kernel_size=[11],
repeat=1,
stride=[1],
dilation=[1],
activation="relu",
conv_mask=True,
separable=False,
se=False,
)
base.update(kwargs)
return base
def check_module_exists(self, module, cls):
global _MODULE_EXISTS
_MODULE_EXISTS = 0
def _traverse(m):
if isinstance(m, cls):
global _MODULE_EXISTS
_MODULE_EXISTS += 1
module.apply(_traverse)
assert _MODULE_EXISTS > 0
@pytest.mark.unit
def test_basic_block(self):
config = self.jasper_base_config(residual=False)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_residual_block(self):
config = self.jasper_base_config(residual=True)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_basic_block_repeat(self):
config = self.jasper_base_config(residual=False, repeat=3)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
assert len(block.mconv) == 3 * 3 + 1 # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv)
@pytest.mark.unit
def test_basic_block_repeat_stride(self):
config = self.jasper_base_config(residual=False, repeat=3, stride=[2])
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 17]) # 131 // (stride ^ repeats)
assert ylen[0] == 17 # 131 // (stride ^ repeats)
assert len(block.mconv) == 3 * 3 + 1 # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv)
@pytest.mark.unit
def test_basic_block_repeat_stride_last(self):
config = self.jasper_base_config(residual=False, repeat=3, stride=[2], stride_last=True)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 66]) # 131 // stride
assert ylen[0] == 66 # 131 // stride
assert len(block.mconv) == 3 * 3 + 1 # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv)
@pytest.mark.unit
def test_basic_block_repeat_separable(self):
config = self.jasper_base_config(residual=False, repeat=3, separable=True)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
assert len(block.mconv) == 3 * 4 + 1 # (3 repeats x {1 dconv + 1 pconv + 1 norm + 1 dropout} + final conv)
@pytest.mark.unit
def test_basic_block_stride(self):
config = self.jasper_base_config(stride=[2], residual=False)
act = jasper.jasper_activations.get(config.pop('activation'))()
print(config)
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 66])
assert ylen[0] == 66
@pytest.mark.unit
def test_residual_block_stride(self):
config = self.jasper_base_config(stride=[2], residual=True, residual_mode='stride_add')
act = jasper.jasper_activations.get(config.pop('activation'))()
print(config)
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 66])
assert ylen[0] == 66
@pytest.mark.unit
def test_residual_block_activations(self):
for activation in jasper.jasper_activations.keys():
config = self.jasper_base_config(activation=activation)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
self.check_module_exists(block, act.__class__)
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_residual_block_normalizations(self):
NORMALIZATIONS = ["batch", "layer", "group"]
for normalization in NORMALIZATIONS:
config = self.jasper_base_config(normalization=normalization)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_residual_block_se(self):
config = self.jasper_base_config(se=True, se_reduction_ratio=8)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
self.check_module_exists(block, jasper.SqueezeExcite)
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
@pytest.mark.unit
def test_residual_block_asymmetric_pad_future_contexts(self):
# test future contexts at various values
# 0 = no future context
# 2 = limited future context
# 5 = symmetric context
# 8 = excess future context (more future context than present or past context)
future_contexts = [0, 2, 5, 8]
for future_context in future_contexts:
print(future_context)
config = self.jasper_base_config(future_context=future_context)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
self.check_module_exists(block, torch.nn.ConstantPad1d)
self.check_module_exists(block, jasper.MaskedConv1d)
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
assert block.mconv[0].pad_layer is not None
assert block.mconv[0]._padding == (config['kernel_size'][0] - 1 - future_context, future_context)
@pytest.mark.unit
def test_residual_block_asymmetric_pad_future_context_fallback(self):
# test future contexts at various values
# 15 = K < FC; fall back to symmetric context
future_context = 15
print(future_context)
config = self.jasper_base_config(future_context=future_context)
act = jasper.jasper_activations.get(config.pop('activation'))()
block = jasper.JasperBlock(**config, activation=act)
x = torch.randn(1, 16, 131)
xlen = torch.tensor([131])
y, ylen = block(([x], xlen))
self.check_module_exists(block, jasper.MaskedConv1d)
assert isinstance(block, jasper.JasperBlock)
assert y[0].shape == torch.Size([1, config['planes'], 131])
assert ylen[0] == 131
assert block.mconv[0].pad_layer is None
assert block.mconv[0]._padding == config['kernel_size'][0] // 2
@pytest.mark.unit
def test_padding_size_conv1d(self):
input_channels = 1
output_channels = 1
kernel_sizes = [3, 7, 11]
dilation_sizes = [2, 3, 4]
stride = 1
inp = torch.rand(2, 1, 40)
for kernel_size in kernel_sizes:
for dilation_size in dilation_sizes:
padding = jasper.get_same_padding(kernel_size, stride, dilation_size)
conv = torch.nn.Conv1d(
input_channels, output_channels, kernel_size=kernel_size, dilation=dilation_size, padding=padding
)
out = conv(inp)
assert out.shape == inp.shape
tts/NeMo/tests/collections/asr/test_label_datasets.py
+106
View File
@@ -0,0 +1,106 @@
# Copyright (c) 2020, 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 os
import pytest
import torch
from nemo.collections.asr.data.audio_to_label import TarredAudioToClassificationLabelDataset
from nemo.collections.asr.data.feature_to_label import FeatureToSeqSpeakerLabelDataset
from nemo.collections.asr.parts.preprocessing.feature_loader import ExternalFeatureLoader
from nemo.collections.asr.parts.preprocessing.features import WaveformFeaturizer
class TestASRDatasets:
labels = ["fash", "fbbh", "fclc"]
unique_labels_in_seq = ['0', '1', '2', '3', "zero", "one", "two", "three"]
@pytest.mark.unit
def test_tarred_dataset(self, test_data_dir):
manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json'))
# Test braceexpand loading
tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
featurizer = WaveformFeaturizer(sample_rate=16000, int_values=False, augmentor=None)
ds_braceexpand = TarredAudioToClassificationLabelDataset(
audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer
)
assert len(ds_braceexpand) == 32
count = 0
for _ in ds_braceexpand:
count += 1
assert count == 32
# Test loading via list
tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)]
ds_list_load = TarredAudioToClassificationLabelDataset(
audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer
)
count = 0
for _ in ds_list_load:
count += 1
assert count == 32
@pytest.mark.unit
def test_tarred_dataset_duplicate_name(self, test_data_dir):
manifest_path = os.path.abspath(
os.path.join(test_data_dir, 'asr/tarred_an4/tarred_duplicate_audio_manifest.json')
)
# Test braceexpand loading
tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar'))
featurizer = WaveformFeaturizer(sample_rate=16000, int_values=False, augmentor=None)
ds_braceexpand = TarredAudioToClassificationLabelDataset(
audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer
)
assert len(ds_braceexpand) == 6
count = 0
for _ in ds_braceexpand:
count += 1
assert count == 6
# Test loading via list
tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)]
ds_list_load = TarredAudioToClassificationLabelDataset(
audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer
)
count = 0
for _ in ds_list_load:
count += 1
assert count == 6
@pytest.mark.unit
def test_feat_seqlabel_dataset(self, test_data_dir):
manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/feat/emb.json'))
feature_loader = ExternalFeatureLoader(file_path=manifest_path, sample_rate=16000, augmentor=None)
ds_braceexpand = FeatureToSeqSpeakerLabelDataset(
manifest_filepath=manifest_path, labels=self.unique_labels_in_seq, feature_loader=feature_loader
)
# fmt: off
correct_label = torch.tensor(
[0.0, 1.0, 2.0, 2.0, 1.0, 2.0, 1.0, 2.0, 2.0, 1.0, 2.0, 2.0, 3.0, 1.0, 2.0, 2.0, 2.0, 0.0, 2.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 2.0, 2.0, 2.0, 1.0, 2.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 2.0, 1.0, 2.0, 1.0,]
)
# fmt: on
correct_label_length = torch.tensor(50)
assert ds_braceexpand[0][0].shape == (50, 32)
assert torch.equal(ds_braceexpand[0][2], correct_label)
assert torch.equal(ds_braceexpand[0][3], correct_label_length)
count = 0
for _ in ds_braceexpand:
count += 1
assert count == 2
tts/NeMo/tests/collections/asr/test_speaker_label_models.py
+63
View File
@@ -0,0 +1,63 @@
# Copyright (c) 2020, 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.
from unittest import TestCase
import pytest
from omegaconf import DictConfig
from nemo.collections.asr.models import EncDecSpeakerLabelModel
class EncDecSpeechLabelModelTest(TestCase):
@pytest.mark.unit
def test_constructor(self):
preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
encoder = {
'cls': 'nemo.collections.asr.modules.ConvASREncoder',
'params': {
'feat_in': 64,
'activation': 'relu',
'conv_mask': True,
'jasper': [
{
'filters': 512,
'repeat': 1,
'kernel': [1],
'stride': [1],
'dilation': [1],
'dropout': 0.0,
'residual': False,
'separable': False,
}
],
},
}
decoder = {
'cls': 'nemo.collections.asr.modules.SpeakerDecoder',
'params': {'feat_in': 512, 'num_classes': 2, 'pool_mode': 'xvector', 'emb_sizes': [1024]},
}
modelConfig = DictConfig(
{'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
)
speaker_model = EncDecSpeakerLabelModel(cfg=modelConfig)
speaker_model.train()
# TODO: make proper config and assert correct number of weights
# Check to/from config_dict:
confdict = speaker_model.to_config_dict()
instance2 = EncDecSpeakerLabelModel.from_config_dict(confdict)
self.assertTrue(isinstance(instance2, EncDecSpeakerLabelModel))