mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-20 08:04:35 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -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__])
|
||||
@@ -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
|
||||
@@ -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__])
|
||||
@@ -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__])
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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}"
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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))
|
||||
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
@@ -0,0 +1,69 @@
|
||||
# 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.
|
||||
|
||||
# Copyright 2018-2020 William Falcon
|
||||
#
|
||||
# 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 dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
from .pl_utils import BATCH_SIZE, NUM_BATCHES, NUM_CLASSES
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LossInput:
|
||||
"""
|
||||
The input for ``nemo.collections.common.metrics.GlobalAverageLossMetric`` metric tests.
|
||||
|
||||
Args:
|
||||
loss_sum_or_avg: a one dimensional float tensor which contains losses for averaging. Each element is either a
|
||||
sum or mean of several losses depending on the parameter ``take_avg_loss`` of the
|
||||
``nemo.collections.common.metrics.GlobalAverageLossMetric`` class.
|
||||
num_measurements: a one dimensional integer tensor which contains number of measurements which sums or average
|
||||
values are in ``loss_sum_or_avg``.
|
||||
"""
|
||||
|
||||
loss_sum_or_avg: torch.Tensor
|
||||
num_measurements: torch.Tensor
|
||||
|
||||
|
||||
NO_ZERO_NUM_MEASUREMENTS = LossInput(
|
||||
loss_sum_or_avg=torch.rand(NUM_BATCHES) * 2.0 - 1.0, num_measurements=torch.randint(1, 100, (NUM_BATCHES,)),
|
||||
)
|
||||
|
||||
SOME_NUM_MEASUREMENTS_ARE_ZERO = LossInput(
|
||||
loss_sum_or_avg=torch.rand(NUM_BATCHES) * 2.0 - 1.0,
|
||||
num_measurements=torch.cat(
|
||||
(
|
||||
torch.randint(1, 100, (NUM_BATCHES // 2,), dtype=torch.int32),
|
||||
torch.zeros(NUM_BATCHES - NUM_BATCHES // 2, dtype=torch.int32),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
ALL_NUM_MEASUREMENTS_ARE_ZERO = LossInput(
|
||||
loss_sum_or_avg=torch.rand(NUM_BATCHES) * 2.0 - 1.0, num_measurements=torch.zeros(NUM_BATCHES, dtype=torch.int32),
|
||||
)
|
||||
@@ -0,0 +1,49 @@
|
||||
# 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.
|
||||
|
||||
# Copyright 2018-2020 William Falcon
|
||||
#
|
||||
# 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 collections import namedtuple
|
||||
|
||||
import torch
|
||||
|
||||
from .pl_utils import BATCH_SIZE, NUM_BATCHES, NUM_CLASSES
|
||||
|
||||
Input = namedtuple('Input', ["probs", "logits"])
|
||||
|
||||
|
||||
ONLY_PROBS = Input(probs=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES), logits=None)
|
||||
|
||||
ONLY_LOGITS1 = Input(probs=None, logits=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES))
|
||||
|
||||
ONLY_LOGITS100 = Input(probs=None, logits=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES) * 200 - 100)
|
||||
|
||||
PROBS_AND_LOGITS = Input(
|
||||
probs=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES),
|
||||
logits=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES) * 200 - 100,
|
||||
)
|
||||
|
||||
NO_PROBS_NO_LOGITS = Input(probs=None, logits=None)
|
||||
@@ -0,0 +1,557 @@
|
||||
# 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.
|
||||
|
||||
# Copyright 2018-2020 William Falcon
|
||||
#
|
||||
# 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 pickle
|
||||
import sys
|
||||
from functools import partial
|
||||
from typing import Callable, Optional
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from scipy.stats import entropy
|
||||
from torch.distributions.utils import logits_to_probs
|
||||
from torch.multiprocessing import Pool, set_start_method
|
||||
from torchmetrics import Metric
|
||||
|
||||
from nemo.collections.common.metrics import GlobalAverageLossMetric, Perplexity
|
||||
|
||||
NUM_PROCESSES = 2
|
||||
NUM_BATCHES = 10
|
||||
BATCH_SIZE = 16
|
||||
NUM_CLASSES = 5
|
||||
EXTRA_DIM = 3
|
||||
THRESHOLD = 0.5
|
||||
|
||||
|
||||
def setup_ddp(rank, world_size):
|
||||
""" Setup ddp enviroment """
|
||||
os.environ["MASTER_ADDR"] = 'localhost'
|
||||
os.environ['MASTER_PORT'] = '8088'
|
||||
|
||||
if torch.distributed.is_available() and sys.platform not in ['win32', 'cygwin']:
|
||||
torch.distributed.init_process_group("gloo", rank=rank, world_size=world_size)
|
||||
|
||||
|
||||
def _class_test(
|
||||
rank: int,
|
||||
worldsize: int,
|
||||
preds: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
metric_class: Metric,
|
||||
sk_metric: Callable,
|
||||
dist_sync_on_step: bool,
|
||||
metric_args: dict = {},
|
||||
check_dist_sync_on_step: bool = True,
|
||||
check_batch: bool = True,
|
||||
atol: float = 1e-8,
|
||||
):
|
||||
""" Utility function doing the actual comparison between lightning class metric
|
||||
and reference metric.
|
||||
Args:
|
||||
rank: rank of current process
|
||||
worldsize: number of processes
|
||||
preds: torch tensor with predictions
|
||||
target: torch tensor with targets
|
||||
metric_class: lightning metric class that should be tested
|
||||
sk_metric: callable function that is used for comparison
|
||||
dist_sync_on_step: bool, if true will synchronize metric state across
|
||||
processes at each ``forward()``
|
||||
metric_args: dict with additional arguments used for class initialization
|
||||
check_dist_sync_on_step: bool, if true will check if the metric is also correctly
|
||||
calculated per batch per device (and not just at the end)
|
||||
check_batch: bool, if true will check if the metric is also correctly
|
||||
calculated across devices for each batch (and not just at the end)
|
||||
"""
|
||||
# Instanciate lightning metric
|
||||
metric = metric_class(compute_on_step=True, dist_sync_on_step=dist_sync_on_step, **metric_args)
|
||||
|
||||
# verify metrics work after being loaded from pickled state
|
||||
pickled_metric = pickle.dumps(metric)
|
||||
metric = pickle.loads(pickled_metric)
|
||||
|
||||
for i in range(rank, NUM_BATCHES, worldsize):
|
||||
batch_result = metric(preds[i], target[i])
|
||||
|
||||
if metric.dist_sync_on_step:
|
||||
if rank == 0:
|
||||
ddp_preds = torch.stack([preds[i + r] for r in range(worldsize)])
|
||||
ddp_target = torch.stack([target[i + r] for r in range(worldsize)])
|
||||
sk_batch_result = sk_metric(ddp_preds, ddp_target)
|
||||
# assert for dist_sync_on_step
|
||||
if check_dist_sync_on_step:
|
||||
assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol)
|
||||
else:
|
||||
sk_batch_result = sk_metric(preds[i], target[i])
|
||||
# assert for batch
|
||||
if check_batch:
|
||||
assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol)
|
||||
|
||||
# check on all batches on all ranks
|
||||
result = metric.compute()
|
||||
assert isinstance(result, torch.Tensor)
|
||||
|
||||
total_preds = torch.stack([preds[i] for i in range(NUM_BATCHES)])
|
||||
total_target = torch.stack([target[i] for i in range(NUM_BATCHES)])
|
||||
sk_result = sk_metric(total_preds, total_target)
|
||||
|
||||
# assert after aggregation
|
||||
assert np.allclose(result.numpy(), sk_result, atol=atol)
|
||||
|
||||
|
||||
def _functional_test(
|
||||
preds: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
metric_functional: Callable,
|
||||
sk_metric: Callable,
|
||||
metric_args: dict = {},
|
||||
atol: float = 1e-8,
|
||||
):
|
||||
""" Utility function doing the actual comparison between lightning functional metric
|
||||
and reference metric.
|
||||
Args:
|
||||
preds: torch tensor with predictions
|
||||
target: torch tensor with targets
|
||||
metric_functional: lightning metric functional that should be tested
|
||||
sk_metric: callable function that is used for comparison
|
||||
metric_args: dict with additional arguments used for class initialization
|
||||
"""
|
||||
metric = partial(metric_functional, **metric_args)
|
||||
|
||||
for i in range(NUM_BATCHES):
|
||||
lightning_result = metric(preds[i], target[i])
|
||||
sk_result = sk_metric(preds[i], target[i])
|
||||
|
||||
# assert its the same
|
||||
assert np.allclose(lightning_result.numpy(), sk_result, atol=atol)
|
||||
|
||||
|
||||
class MetricTester:
|
||||
""" Class used for efficiently run alot of parametrized tests in ddp mode.
|
||||
Makes sure that ddp is only setup once and that pool of processes are
|
||||
used for all tests.
|
||||
All tests should subclass from this and implement a new method called
|
||||
`test_metric_name`
|
||||
where the method `self.run_metric_test` is called inside.
|
||||
"""
|
||||
|
||||
atol = 1e-8
|
||||
|
||||
def setup_class(self):
|
||||
""" Setup the metric class. This will spawn the pool of workers that are
|
||||
used for metric testing and setup_ddp
|
||||
"""
|
||||
try:
|
||||
set_start_method('spawn')
|
||||
except RuntimeError:
|
||||
pass
|
||||
self.poolSize = NUM_PROCESSES
|
||||
self.pool = Pool(processes=self.poolSize)
|
||||
self.pool.starmap(setup_ddp, [(rank, self.poolSize) for rank in range(self.poolSize)])
|
||||
|
||||
def teardown_class(self):
|
||||
""" Close pool of workers """
|
||||
self.pool.close()
|
||||
self.pool.join()
|
||||
|
||||
def run_functional_metric_test(
|
||||
self,
|
||||
preds: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
metric_functional: Callable,
|
||||
sk_metric: Callable,
|
||||
metric_args: dict = {},
|
||||
):
|
||||
""" Main method that should be used for testing functions. Call this inside
|
||||
testing method
|
||||
Args:
|
||||
preds: torch tensor with predictions
|
||||
target: torch tensor with targets
|
||||
metric_functional: lightning metric class that should be tested
|
||||
sk_metric: callable function that is used for comparison
|
||||
metric_args: dict with additional arguments used for class initialization
|
||||
"""
|
||||
_functional_test(
|
||||
preds=preds,
|
||||
target=target,
|
||||
metric_functional=metric_functional,
|
||||
sk_metric=sk_metric,
|
||||
metric_args=metric_args,
|
||||
atol=self.atol,
|
||||
)
|
||||
|
||||
def run_class_metric_test(
|
||||
self,
|
||||
ddp: bool,
|
||||
preds: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
metric_class: Metric,
|
||||
sk_metric: Callable,
|
||||
dist_sync_on_step: bool,
|
||||
metric_args: dict = {},
|
||||
check_dist_sync_on_step: bool = True,
|
||||
check_batch: bool = True,
|
||||
):
|
||||
""" Main method that should be used for testing class. Call this inside testing
|
||||
methods.
|
||||
Args:
|
||||
ddp: bool, if running in ddp mode or not
|
||||
preds: torch tensor with predictions
|
||||
target: torch tensor with targets
|
||||
metric_class: lightning metric class that should be tested
|
||||
sk_metric: callable function that is used for comparison
|
||||
dist_sync_on_step: bool, if true will synchronize metric state across
|
||||
processes at each ``forward()``
|
||||
metric_args: dict with additional arguments used for class initialization
|
||||
check_dist_sync_on_step: bool, if true will check if the metric is also correctly
|
||||
calculated per batch per device (and not just at the end)
|
||||
check_batch: bool, if true will check if the metric is also correctly
|
||||
calculated across devices for each batch (and not just at the end)
|
||||
"""
|
||||
if ddp:
|
||||
if sys.platform == "win32":
|
||||
pytest.skip("DDP not supported on windows")
|
||||
|
||||
self.pool.starmap(
|
||||
partial(
|
||||
_class_test,
|
||||
preds=preds,
|
||||
target=target,
|
||||
metric_class=metric_class,
|
||||
sk_metric=sk_metric,
|
||||
dist_sync_on_step=dist_sync_on_step,
|
||||
metric_args=metric_args,
|
||||
check_dist_sync_on_step=check_dist_sync_on_step,
|
||||
check_batch=check_batch,
|
||||
atol=self.atol,
|
||||
),
|
||||
[(rank, self.poolSize) for rank in range(self.poolSize)],
|
||||
)
|
||||
else:
|
||||
_class_test(
|
||||
0,
|
||||
1,
|
||||
preds=preds,
|
||||
target=target,
|
||||
metric_class=metric_class,
|
||||
sk_metric=sk_metric,
|
||||
dist_sync_on_step=dist_sync_on_step,
|
||||
metric_args=metric_args,
|
||||
check_dist_sync_on_step=check_dist_sync_on_step,
|
||||
check_batch=check_batch,
|
||||
atol=self.atol,
|
||||
)
|
||||
|
||||
|
||||
def reference_perplexity_func(probs):
|
||||
ent = entropy(probs, axis=-1)
|
||||
ppl = np.exp(ent)
|
||||
return ppl.mean()
|
||||
|
||||
|
||||
def _perplexity_class_test(
|
||||
rank: int,
|
||||
worldsize: int,
|
||||
probs: Optional[torch.Tensor],
|
||||
logits: Optional[torch.Tensor],
|
||||
dist_sync_on_step: bool,
|
||||
metric_args: dict = {},
|
||||
check_dist_sync_on_step: bool = True,
|
||||
check_batch: bool = True,
|
||||
atol: float = 1e-8,
|
||||
):
|
||||
""" Utility function doing the actual comparison between lightning class metric
|
||||
and reference metric.
|
||||
Args:
|
||||
rank: rank of current process
|
||||
worldsize: number of processes
|
||||
probs: torch tensor with probabilities
|
||||
logits: torch tensor with logits. The function checks ``probs`` and ``logits are mutually exclusive for
|
||||
``Perplexity`` metric.
|
||||
dist_sync_on_step: bool, if true will synchronize metric state across
|
||||
processes at each ``forward()``
|
||||
metric_args: dict with additional arguments used for class initialization
|
||||
check_dist_sync_on_step: bool, if true will check if the metric is also correctly
|
||||
calculated per batch per device (and not just at the end)
|
||||
check_batch: bool, if true will check if the metric is also correctly
|
||||
calculated across devices for each batch (and not just at the end)
|
||||
"""
|
||||
# Instanciate lightning metric
|
||||
perplexity = Perplexity(compute_on_step=True, dist_sync_on_step=dist_sync_on_step, **metric_args)
|
||||
if (probs is None) == (logits is None):
|
||||
with pytest.raises(ValueError):
|
||||
perplexity(probs, logits)
|
||||
return
|
||||
|
||||
# verify perplexity works after being loaded from pickled state
|
||||
pickled_metric = pickle.dumps(perplexity)
|
||||
perplexity = pickle.loads(pickled_metric)
|
||||
|
||||
for i in range(rank, NUM_BATCHES, worldsize):
|
||||
batch_result = perplexity(None if probs is None else probs[i], None if logits is None else logits[i])
|
||||
|
||||
if perplexity.dist_sync_on_step:
|
||||
if rank == 0:
|
||||
if probs is not None:
|
||||
ddp_probs = torch.stack([probs[i + r] for r in range(worldsize)])
|
||||
else:
|
||||
ddp_logits = torch.stack([logits[i + r] for r in range(worldsize)])
|
||||
ddp_probs = logits_to_probs(ddp_logits, is_binary=False)
|
||||
sk_batch_result = reference_perplexity_func(ddp_probs)
|
||||
# assert for dist_sync_on_step
|
||||
if check_dist_sync_on_step:
|
||||
assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol)
|
||||
else:
|
||||
if probs is None:
|
||||
p = logits_to_probs(logits[i], is_binary=False)
|
||||
else:
|
||||
p = probs[i]
|
||||
sk_batch_result = reference_perplexity_func(p)
|
||||
# assert for batch
|
||||
if check_batch:
|
||||
assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol)
|
||||
|
||||
assert (probs is None) != (logits is None)
|
||||
# check on all batches on all ranks
|
||||
result = perplexity.compute()
|
||||
assert isinstance(result, torch.Tensor)
|
||||
|
||||
if probs is None:
|
||||
probs = logits_to_probs(logits, is_binary=False)
|
||||
sk_result = reference_perplexity_func(probs)
|
||||
|
||||
# assert after aggregation
|
||||
assert np.allclose(result.numpy(), sk_result, atol=atol)
|
||||
|
||||
|
||||
class PerplexityTester(MetricTester):
|
||||
def run_class_perplexity_test(
|
||||
self,
|
||||
ddp: bool,
|
||||
probs: Optional[torch.Tensor],
|
||||
logits: Optional[torch.Tensor],
|
||||
dist_sync_on_step: bool,
|
||||
metric_args: dict = {},
|
||||
check_dist_sync_on_step: bool = True,
|
||||
check_batch: bool = True,
|
||||
):
|
||||
""" Main method that should be used for testing class. Call this inside testing
|
||||
methods.
|
||||
Args:
|
||||
ddp: bool, if running in ddp mode or not
|
||||
probs: torch tensor with probabilities.
|
||||
logits: torch tensor with logits. This test checks that probs and logits are mutually exclusive for
|
||||
``Perplexity`` metric.
|
||||
dist_sync_on_step: bool, if true will synchronize metric state across
|
||||
processes at each ``forward()``
|
||||
metric_args: dict with additional arguments used for class initialization
|
||||
check_dist_sync_on_step: bool, if true will check if the metric is also correctly
|
||||
calculated per batch per device (and not just at the end)
|
||||
check_batch: bool, if true will check if the metric is also correctly
|
||||
calculated across devices for each batch (and not just at the end)
|
||||
"""
|
||||
if ddp:
|
||||
if sys.platform == "win32":
|
||||
pytest.skip("DDP not supported on windows")
|
||||
|
||||
self.pool.starmap(
|
||||
partial(
|
||||
_perplexity_class_test,
|
||||
probs=probs,
|
||||
logits=logits,
|
||||
dist_sync_on_step=dist_sync_on_step,
|
||||
metric_args=metric_args,
|
||||
check_dist_sync_on_step=check_dist_sync_on_step,
|
||||
check_batch=check_batch,
|
||||
atol=self.atol,
|
||||
),
|
||||
[(rank, self.poolSize) for rank in range(self.poolSize)],
|
||||
)
|
||||
else:
|
||||
_perplexity_class_test(
|
||||
0,
|
||||
1,
|
||||
probs=probs,
|
||||
logits=logits,
|
||||
dist_sync_on_step=dist_sync_on_step,
|
||||
metric_args=metric_args,
|
||||
check_dist_sync_on_step=check_dist_sync_on_step,
|
||||
check_batch=check_batch,
|
||||
atol=self.atol,
|
||||
)
|
||||
|
||||
|
||||
def reference_loss_func(loss_sum_or_avg: torch.Tensor, num_measurements: torch.Tensor, take_avg_loss: bool):
|
||||
"""
|
||||
Returns average loss for data from``loss_sum_or_avg``. This function sums all losses from ``loss_sum_or_avg`` and
|
||||
divides the sum by the sum of ``num_measurements`` elements.
|
||||
|
||||
If ``take_avg_loss`` is ``True`` then ``loss_sum_or_avg[i]`` elements are mean values of ``num_measurements[i]``
|
||||
losses. In that case before computing sum of losses each element of ``loss_sum_or_avg`` is multiplied by
|
||||
corresponding element of ``num_measurements``.
|
||||
|
||||
If ``num_measurements`` sum is zero then the function returns NaN tensor.
|
||||
|
||||
The function is used for testing ``nemo.collections.common.metrics.GlobalAverageLossMetric`` class.
|
||||
|
||||
Args:
|
||||
loss_sum_or_avg: a one dimensional float ``torch.Tensor``. Sums or mean values of loss.
|
||||
num_measurements: a one dimensional integer ``torch.Tensor``. Number of values on which sums of means in
|
||||
``loss_sum_or_avg`` are calculated.
|
||||
take_avg_loss: if ``True`` then ``loss_sum_or_avg`` contains mean losses else ``loss_sum_or_avg`` contains
|
||||
sums of losses.
|
||||
"""
|
||||
if take_avg_loss:
|
||||
for i in range(loss_sum_or_avg.shape[0]):
|
||||
loss_sum_or_avg[i] *= num_measurements[i]
|
||||
nm_sum = num_measurements.sum()
|
||||
if nm_sum.eq(0):
|
||||
return torch.tensor(float('nan'))
|
||||
return loss_sum_or_avg.sum() / nm_sum
|
||||
|
||||
|
||||
def _loss_class_test(
|
||||
rank: int,
|
||||
worldsize: int,
|
||||
loss_sum_or_avg: Optional[torch.Tensor],
|
||||
num_measurements: Optional[torch.Tensor],
|
||||
dist_sync_on_step: bool,
|
||||
take_avg_loss: bool,
|
||||
check_dist_sync_on_step: bool = True,
|
||||
check_batch: bool = True,
|
||||
atol: float = 1e-8,
|
||||
):
|
||||
""" Utility function doing the actual comparison between lightning class metric
|
||||
and reference metric.
|
||||
Args:
|
||||
rank: rank of current process
|
||||
worldsize: number of processes
|
||||
loss_sum_or_avg: a one dimensional float torch tensor with loss sums or means.
|
||||
num_measurements: a one dimensional integer torch tensor with number of values on which sums or means from
|
||||
``loss_sum_or_avg`` were computed.
|
||||
dist_sync_on_step: bool, if true will synchronize metric state across
|
||||
processes at each ``forward()``
|
||||
take_avg_loss: dict with additional arguments used for class initialization
|
||||
check_dist_sync_on_step: bool, if true will check if the metric is also correctly
|
||||
calculated per batch per device (and not just at the end)
|
||||
check_batch: bool, if true will check if the metric is also correctly
|
||||
calculated across devices for each batch (and not just at the end)
|
||||
"""
|
||||
# Instantiate lightning metric
|
||||
loss_metric = GlobalAverageLossMetric(
|
||||
compute_on_step=True, dist_sync_on_step=dist_sync_on_step, take_avg_loss=take_avg_loss
|
||||
)
|
||||
|
||||
# verify perplexity works after being loaded from pickled state
|
||||
pickled_metric = pickle.dumps(loss_metric)
|
||||
loss_metric = pickle.loads(pickled_metric)
|
||||
|
||||
for i in range(rank, NUM_BATCHES, worldsize):
|
||||
batch_result = loss_metric(loss_sum_or_avg[i], num_measurements[i])
|
||||
if loss_metric.dist_sync_on_step:
|
||||
if rank == 0:
|
||||
ddp_loss_sum_or_avg = torch.stack([loss_sum_or_avg[i + r] for r in range(worldsize)])
|
||||
ddp_num_measurements = torch.stack([num_measurements[i + r] for r in range(worldsize)])
|
||||
sk_batch_result = reference_loss_func(ddp_loss_sum_or_avg, ddp_num_measurements, take_avg_loss)
|
||||
# assert for dist_sync_on_step
|
||||
if check_dist_sync_on_step:
|
||||
if sk_batch_result.isnan():
|
||||
assert batch_result.isnan()
|
||||
else:
|
||||
assert (
|
||||
np.allclose(batch_result.numpy(), sk_batch_result, atol=atol),
|
||||
f"batch_result = {batch_result.numpy()}, sk_batch_result = {sk_batch_result}",
|
||||
)
|
||||
else:
|
||||
ls = loss_sum_or_avg[i : i + 1]
|
||||
nm = num_measurements[i : i + 1]
|
||||
sk_batch_result = reference_loss_func(ls, nm, take_avg_loss)
|
||||
# assert for batch
|
||||
if check_batch:
|
||||
if sk_batch_result.isnan():
|
||||
assert batch_result.isnan()
|
||||
else:
|
||||
assert (
|
||||
np.allclose(batch_result.numpy(), sk_batch_result, atol=atol),
|
||||
f"batch_result = {batch_result.numpy()}, sk_batch_result = {sk_batch_result}",
|
||||
)
|
||||
# check on all batches on all ranks
|
||||
result = loss_metric.compute()
|
||||
assert isinstance(result, torch.Tensor)
|
||||
sk_result = reference_loss_func(loss_sum_or_avg, num_measurements, take_avg_loss)
|
||||
|
||||
# assert after aggregation
|
||||
if sk_result.isnan():
|
||||
assert result.isnan()
|
||||
else:
|
||||
assert (
|
||||
np.allclose(result.numpy(), sk_result, atol=atol),
|
||||
f"result = {result.numpy()}, sk_result = {sk_result}",
|
||||
)
|
||||
|
||||
|
||||
class LossTester(MetricTester):
|
||||
def run_class_loss_test(
|
||||
self,
|
||||
ddp: bool,
|
||||
loss_sum_or_avg: torch.Tensor,
|
||||
num_measurements: torch.Tensor,
|
||||
dist_sync_on_step: bool,
|
||||
take_avg_loss: bool,
|
||||
check_dist_sync_on_step: bool = True,
|
||||
check_batch: bool = True,
|
||||
):
|
||||
if ddp:
|
||||
if sys.platform == "win32":
|
||||
pytest.skip("DDP not supported on windows")
|
||||
self.pool.starmap(
|
||||
partial(
|
||||
_loss_class_test,
|
||||
loss_sum_or_avg=loss_sum_or_avg,
|
||||
num_measurements=num_measurements,
|
||||
dist_sync_on_step=dist_sync_on_step,
|
||||
take_avg_loss=take_avg_loss,
|
||||
check_dist_sync_on_step=check_dist_sync_on_step,
|
||||
check_batch=check_batch,
|
||||
atol=self.atol,
|
||||
),
|
||||
[(rank, self.poolSize) for rank in range(self.poolSize)],
|
||||
)
|
||||
else:
|
||||
_loss_class_test(
|
||||
0,
|
||||
1,
|
||||
loss_sum_or_avg=loss_sum_or_avg,
|
||||
num_measurements=num_measurements,
|
||||
dist_sync_on_step=dist_sync_on_step,
|
||||
take_avg_loss=take_avg_loss,
|
||||
check_dist_sync_on_step=check_dist_sync_on_step,
|
||||
check_batch=check_batch,
|
||||
atol=self.atol,
|
||||
)
|
||||
@@ -0,0 +1,150 @@
|
||||
# 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 .loss_inputs import ALL_NUM_MEASUREMENTS_ARE_ZERO, NO_ZERO_NUM_MEASUREMENTS, SOME_NUM_MEASUREMENTS_ARE_ZERO
|
||||
from .perplexity_inputs import NO_PROBS_NO_LOGITS, ONLY_LOGITS1, ONLY_LOGITS100, ONLY_PROBS, PROBS_AND_LOGITS
|
||||
from .pl_utils import LossTester, PerplexityTester
|
||||
from nemo.collections.common.metrics.classification_accuracy import TopKClassificationAccuracy
|
||||
|
||||
|
||||
class TestCommonMetrics:
|
||||
top_k_logits = torch.tensor([[0.1, 0.3, 0.2, 0.0], [0.9, 0.6, 0.2, 0.3], [0.2, 0.1, 0.4, 0.3]],) # 1 # 0 # 2
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_top_1_accuracy(self):
|
||||
labels = torch.tensor([0, 0, 2], dtype=torch.long)
|
||||
|
||||
accuracy = TopKClassificationAccuracy(top_k=None)
|
||||
acc = accuracy(logits=self.top_k_logits, labels=labels)
|
||||
|
||||
assert accuracy.correct_counts_k.shape == torch.Size([1])
|
||||
assert accuracy.total_counts_k.shape == torch.Size([1])
|
||||
assert abs(acc[0] - 0.667) < 1e-3
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_top_1_2_accuracy(self):
|
||||
labels = torch.tensor([0, 1, 0], dtype=torch.long)
|
||||
|
||||
accuracy = TopKClassificationAccuracy(top_k=[1, 2])
|
||||
top1_acc, top2_acc = accuracy(logits=self.top_k_logits, labels=labels)
|
||||
|
||||
assert accuracy.correct_counts_k.shape == torch.Size([2])
|
||||
assert accuracy.total_counts_k.shape == torch.Size([2])
|
||||
|
||||
assert abs(top1_acc - 0.0) < 1e-3
|
||||
assert abs(top2_acc - 0.333) < 1e-3
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_top_1_accuracy_distributed(self):
|
||||
# Simulate test on 2 process DDP execution
|
||||
labels = torch.tensor([[0, 0, 2], [2, 0, 0]], dtype=torch.long)
|
||||
|
||||
accuracy = TopKClassificationAccuracy(top_k=None)
|
||||
proc1_acc = accuracy(logits=self.top_k_logits, labels=labels[0])
|
||||
correct1, total1 = accuracy.correct_counts_k, accuracy.total_counts_k
|
||||
|
||||
accuracy.reset()
|
||||
proc2_acc = accuracy(logits=torch.flip(self.top_k_logits, dims=[1]), labels=labels[1]) # reverse logits
|
||||
correct2, total2 = accuracy.correct_counts_k, accuracy.total_counts_k
|
||||
|
||||
correct = torch.stack([correct1, correct2])
|
||||
total = torch.stack([total1, total2])
|
||||
|
||||
assert correct.shape == torch.Size([2, 1])
|
||||
assert total.shape == torch.Size([2, 1])
|
||||
|
||||
assert abs(proc1_acc[0] - 0.667) < 1e-3 # 2/3
|
||||
assert abs(proc2_acc[0] - 0.333) < 1e-3 # 1/3
|
||||
|
||||
accuracy.reset()
|
||||
accuracy.correct_counts_k = torch.tensor([correct.sum()])
|
||||
accuracy.total_counts_k = torch.tensor([total.sum()])
|
||||
acc_topk = accuracy.compute()
|
||||
acc_top1 = acc_topk[0]
|
||||
|
||||
assert abs(acc_top1 - 0.5) < 1e-3 # 3/6
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_top_1_accuracy_distributed_uneven_batch(self):
|
||||
# Simulate test on 2 process DDP execution
|
||||
accuracy = TopKClassificationAccuracy(top_k=None)
|
||||
|
||||
proc1_acc = accuracy(logits=self.top_k_logits, labels=torch.tensor([0, 0, 2]))
|
||||
correct1, total1 = accuracy.correct_counts_k, accuracy.total_counts_k
|
||||
|
||||
proc2_acc = accuracy(
|
||||
logits=torch.flip(self.top_k_logits, dims=[1])[:2, :], # reverse logits, select first 2 samples
|
||||
labels=torch.tensor([2, 0]),
|
||||
) # reduce number of labels
|
||||
correct2, total2 = accuracy.correct_counts_k, accuracy.total_counts_k
|
||||
|
||||
correct = torch.stack([correct1, correct2])
|
||||
total = torch.stack([total1, total2])
|
||||
|
||||
assert correct.shape == torch.Size([2, 1])
|
||||
assert total.shape == torch.Size([2, 1])
|
||||
|
||||
assert abs(proc1_acc[0] - 0.667) < 1e-3 # 2/3
|
||||
assert abs(proc2_acc[0] - 0.500) < 1e-3 # 1/2
|
||||
|
||||
accuracy.correct_counts_k = torch.tensor([correct.sum()])
|
||||
accuracy.total_counts_k = torch.tensor([total.sum()])
|
||||
acc_topk = accuracy.compute()
|
||||
acc_top1 = acc_topk[0]
|
||||
|
||||
assert abs(acc_top1 - 0.6) < 1e-3 # 3/5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ddp", [True, False])
|
||||
@pytest.mark.parametrize("dist_sync_on_step", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
"probs, logits",
|
||||
[
|
||||
(ONLY_PROBS.probs, ONLY_PROBS.logits),
|
||||
(ONLY_LOGITS1.probs, ONLY_LOGITS1.logits),
|
||||
(ONLY_LOGITS100.probs, ONLY_LOGITS100.logits),
|
||||
(PROBS_AND_LOGITS.probs, PROBS_AND_LOGITS.logits),
|
||||
(NO_PROBS_NO_LOGITS.probs, NO_PROBS_NO_LOGITS.logits),
|
||||
],
|
||||
)
|
||||
class TestPerplexity(PerplexityTester):
|
||||
def test_perplexity(self, ddp, dist_sync_on_step, probs, logits):
|
||||
self.run_class_perplexity_test(
|
||||
ddp=ddp, probs=probs, logits=logits, dist_sync_on_step=dist_sync_on_step,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ddp", [True, False])
|
||||
@pytest.mark.parametrize("dist_sync_on_step", [True, False])
|
||||
@pytest.mark.parametrize("take_avg_loss", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
"loss_sum_or_avg, num_measurements",
|
||||
[
|
||||
(NO_ZERO_NUM_MEASUREMENTS.loss_sum_or_avg, NO_ZERO_NUM_MEASUREMENTS.num_measurements),
|
||||
(SOME_NUM_MEASUREMENTS_ARE_ZERO.loss_sum_or_avg, SOME_NUM_MEASUREMENTS_ARE_ZERO.num_measurements),
|
||||
(ALL_NUM_MEASUREMENTS_ARE_ZERO.loss_sum_or_avg, ALL_NUM_MEASUREMENTS_ARE_ZERO.num_measurements),
|
||||
],
|
||||
)
|
||||
class TestLoss(LossTester):
|
||||
def test_loss(self, ddp, dist_sync_on_step, loss_sum_or_avg, num_measurements, take_avg_loss):
|
||||
self.run_class_loss_test(
|
||||
ddp=ddp,
|
||||
loss_sum_or_avg=loss_sum_or_avg,
|
||||
num_measurements=num_measurements,
|
||||
dist_sync_on_step=dist_sync_on_step,
|
||||
take_avg_loss=take_avg_loss,
|
||||
)
|
||||
@@ -0,0 +1,267 @@
|
||||
# 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
|
||||
|
||||
from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer
|
||||
from nemo.collections.common.tokenizers.youtokentome_tokenizer import YouTokenToMeTokenizer
|
||||
|
||||
MODEL_SPECIAL_TOKENS = {
|
||||
'unk_token': '[UNK]',
|
||||
'sep_token': '[SEP]',
|
||||
'pad_token': '[PAD]',
|
||||
'bos_token': '[CLS]',
|
||||
'mask_token': '[MASK]',
|
||||
'eos_token': '[SEP]',
|
||||
'cls_token': '[CLS]',
|
||||
}
|
||||
|
||||
|
||||
class TestSentencePieceTokenizerLegacy:
|
||||
model_name = "/m_common.model"
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_add_special_tokens(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True)
|
||||
special_tokens = MODEL_SPECIAL_TOKENS
|
||||
tokenizer.add_special_tokens(special_tokens)
|
||||
assert tokenizer.vocab_size == tokenizer.original_vocab_size + len(set(special_tokens.values()))
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_text_to_tokens(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True)
|
||||
special_tokens = MODEL_SPECIAL_TOKENS
|
||||
tokenizer.add_special_tokens(special_tokens)
|
||||
|
||||
text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]"
|
||||
tokens = tokenizer.text_to_tokens(text)
|
||||
|
||||
assert len(tokens) == len(text.split())
|
||||
assert tokens.count("[CLS]") == 1
|
||||
assert tokens.count("[MASK]") == 1
|
||||
assert tokens.count("[SEP]") == 2
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_tokens_to_text(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True)
|
||||
|
||||
text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]"
|
||||
tokens = tokenizer.text_to_tokens(text)
|
||||
result = tokenizer.tokens_to_text(tokens)
|
||||
|
||||
assert text == result
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_text_to_ids(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True)
|
||||
special_tokens = MODEL_SPECIAL_TOKENS
|
||||
tokenizer.add_special_tokens(special_tokens)
|
||||
|
||||
text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]"
|
||||
ids = tokenizer.text_to_ids(text)
|
||||
|
||||
assert len(ids) == len(text.split())
|
||||
assert ids.count(tokenizer.token_to_id("[CLS]")) == 1
|
||||
assert ids.count(tokenizer.token_to_id("[MASK]")) == 1
|
||||
assert ids.count(tokenizer.token_to_id("[SEP]")) == 2
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_ids_to_text(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True)
|
||||
special_tokens = MODEL_SPECIAL_TOKENS
|
||||
tokenizer.add_special_tokens(special_tokens)
|
||||
|
||||
text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]"
|
||||
ids = tokenizer.text_to_ids(text)
|
||||
result = tokenizer.ids_to_text(ids)
|
||||
|
||||
assert text == result
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_tokens_to_ids(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True)
|
||||
special_tokens = MODEL_SPECIAL_TOKENS
|
||||
tokenizer.add_special_tokens(special_tokens)
|
||||
|
||||
text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]"
|
||||
tokens = tokenizer.text_to_tokens(text)
|
||||
ids = tokenizer.tokens_to_ids(tokens)
|
||||
|
||||
assert len(ids) == len(tokens)
|
||||
assert ids.count(tokenizer.token_to_id("[CLS]")) == 1
|
||||
assert ids.count(tokenizer.token_to_id("[MASK]")) == 1
|
||||
assert ids.count(tokenizer.token_to_id("[SEP]")) == 2
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_ids_to_tokens(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True)
|
||||
special_tokens = MODEL_SPECIAL_TOKENS
|
||||
tokenizer.add_special_tokens(special_tokens)
|
||||
|
||||
text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]"
|
||||
tokens = tokenizer.text_to_tokens(text)
|
||||
ids = tokenizer.tokens_to_ids(tokens)
|
||||
result = tokenizer.ids_to_tokens(ids)
|
||||
|
||||
assert len(result) == len(tokens)
|
||||
|
||||
for i in range(len(result)):
|
||||
assert result[i] == tokens[i]
|
||||
|
||||
|
||||
class TestSentencePieceTokenizer:
|
||||
model_name = "/m_new.model"
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_text_to_tokens(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
# <cls> is user_defined_symbol in the test tokenizer model
|
||||
# <unk>, <sep>, <s>, and </s> are control symbols
|
||||
text = "<cls> a b c <sep> e f g h i </s>"
|
||||
tokens = tokenizer.text_to_tokens(text)
|
||||
|
||||
assert tokens.count("<cls>") == 1
|
||||
assert tokens.count("<sep>") == 0
|
||||
assert tokens.count("</s>") == 0
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_tokens_to_text(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
# <cls> is user_defined_symbol in the test tokenizer model
|
||||
text = "<cls> a b c e f g h i"
|
||||
tokens = tokenizer.text_to_tokens(text)
|
||||
result = tokenizer.tokens_to_text(tokens)
|
||||
|
||||
assert text == result
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_text_to_ids(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
# <cls> is user_defined_symbol in the test tokenizer model
|
||||
# <unk>, <sep>, <s>, and </s> are control symbols
|
||||
text = "<cls> a b c <sep> e f g h i </s>"
|
||||
tokens = tokenizer.text_to_ids(text)
|
||||
|
||||
assert tokens.count(tokenizer.token_to_id("<cls>")) == 1
|
||||
assert tokens.count(tokenizer.token_to_id("<sep>")) == 0
|
||||
assert tokens.count(tokenizer.token_to_id("</s>")) == 0
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_ids_to_text(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
text = "<cls> a b c <sep> e f g h i </s>"
|
||||
ids = tokenizer.text_to_ids(text)
|
||||
result = tokenizer.ids_to_text(ids)
|
||||
|
||||
assert text == result
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_tokens_to_ids(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
tokens = ["<cls>", "a", "b", "c", "<sep>", "e", "f", "<sep>", "g", "h", "i", "</s>"]
|
||||
ids = tokenizer.tokens_to_ids(tokens)
|
||||
|
||||
assert len(ids) == len(tokens)
|
||||
assert ids.count(tokenizer.token_to_id("<cls>")) == 1
|
||||
assert ids.count(tokenizer.token_to_id("</s>")) == 1
|
||||
assert ids.count(tokenizer.token_to_id("<sep>")) == 2
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_ids_to_tokens(self, test_data_dir):
|
||||
tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
tokens = ["<cls>", "a", "b", "c", "<sep>", "e", "f", "<sep>", "g", "h", "i", "</s>"]
|
||||
ids = tokenizer.tokens_to_ids(tokens)
|
||||
result = tokenizer.ids_to_tokens(ids)
|
||||
|
||||
assert len(result) == len(tokens)
|
||||
|
||||
for i in range(len(result)):
|
||||
assert result[i] == tokens[i]
|
||||
|
||||
|
||||
class TestYouTokenToMeTokenizer:
|
||||
model_name = "/yttm.4096.en-de.model"
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_text_to_tokens(self, test_data_dir):
|
||||
tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
text = "<BOS> a b c e <UNK> f g h i <EOS>"
|
||||
tokens = tokenizer.text_to_tokens(text)
|
||||
|
||||
assert tokens.count("<BOS>") == 0
|
||||
assert tokens.count("<UNK>") == 0
|
||||
assert tokens.count("<EOS>") == 0
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_tokens_to_text(self, test_data_dir):
|
||||
tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
text = "a b c e f g h i"
|
||||
tokens = tokenizer.text_to_tokens(text)
|
||||
result = tokenizer.tokens_to_text(tokens)
|
||||
|
||||
assert text == result
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_text_to_ids(self, test_data_dir):
|
||||
tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
text = "<BOS> a b c <UNK> e f g h i <EOS>"
|
||||
tokens = tokenizer.text_to_ids(text)
|
||||
|
||||
assert tokens.count(tokenizer.bos_id) == 0
|
||||
assert tokens.count(tokenizer.unk_id) == 0
|
||||
assert tokens.count(tokenizer.eos_id) == 0
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_ids_to_text(self, test_data_dir):
|
||||
tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
text = "a b c e f g h i"
|
||||
ids = tokenizer.text_to_ids(text)
|
||||
result = tokenizer.ids_to_text(ids)
|
||||
|
||||
assert text == result
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_tokens_to_ids(self, test_data_dir):
|
||||
tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
tokens = ["<BOS>", "a", "b", "c", "<UNK>", "e", "f", "<UNK>", "g", "h", "i", "<EOS>"]
|
||||
ids = tokenizer.tokens_to_ids(tokens)
|
||||
|
||||
assert len(ids) == len(tokens)
|
||||
assert ids.count(tokenizer.bos_id) == 1
|
||||
assert ids.count(tokenizer.eos_id) == 1
|
||||
assert ids.count(tokenizer.unk_id) == 2
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_ids_to_tokens(self, test_data_dir):
|
||||
tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name)
|
||||
|
||||
tokens = ["<BOS>", "a", "b", "c", "<UNK>", "e", "f", "<UNK>", "g", "h", "i", "<EOS>"]
|
||||
ids = tokenizer.tokens_to_ids(tokens)
|
||||
result = tokenizer.ids_to_tokens(ids)
|
||||
|
||||
assert len(result) == len(tokens)
|
||||
|
||||
for i in range(len(result)):
|
||||
assert result[i] == tokens[i]
|
||||
@@ -0,0 +1,48 @@
|
||||
# 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 sklearn.metrics import precision_recall_fscore_support
|
||||
|
||||
from nemo.collections.nlp.metrics.classification_report import ClassificationReport
|
||||
|
||||
|
||||
class ClassificationReportTests:
|
||||
num_classes = 3
|
||||
label_ids = {'a': 0, 'b': 1, 'c': 2}
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_classification_report(self):
|
||||
|
||||
preds = torch.Tensor([0, 1, 1, 1, 2, 2, 0])
|
||||
labels = torch.Tensor([1, 0, 0, 1, 2, 1, 0])
|
||||
|
||||
def __convert_to_tensor(sklearn_metric):
|
||||
return torch.Tensor([round(sklearn_metric * 100)])[0]
|
||||
|
||||
for mode in ['macro', 'micro', 'weighted']:
|
||||
|
||||
classification_report_nemo = ClassificationReport(
|
||||
num_classes=self.num_classes, label_ids=self.label_ids, mode=mode
|
||||
)
|
||||
# pytest.set_trace()
|
||||
precision, recall, f1, _ = classification_report_nemo(preds, labels)
|
||||
tp, fp, fn = classification_report_nemo.tp, classification_report_nemo.fp, classification_report_nemo.fn
|
||||
pr_sklearn, recall_sklearn, f1_sklearn, _ = precision_recall_fscore_support(labels, preds, average=mode)
|
||||
|
||||
self.assertEqual(torch.round(precision), __convert_to_tensor(pr_sklearn), f'wrong precision for {mode}')
|
||||
self.assertEqual(torch.round(recall), __convert_to_tensor(recall_sklearn), f'wrong recall for {mode}')
|
||||
self.assertEqual(torch.round(f1), __convert_to_tensor(f1_sklearn), f'wrong f1 for {mode}')
|
||||
@@ -0,0 +1,83 @@
|
||||
# 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 shutil
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
import wget
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from nemo.collections.nlp.models import EntityLinkingModel
|
||||
|
||||
|
||||
def get_cfg():
|
||||
|
||||
language_model = OmegaConf.create(
|
||||
{"pretrained_model_name": "bert-base-uncased", "config_file": None, "config": None, "lm_checkpoint": None}
|
||||
)
|
||||
|
||||
tokenizer = OmegaConf.create(
|
||||
{"tokenizer_name": "bert-base-uncased", "vocab_file": None, "tokenizer_model": None, "do_lower_case": True}
|
||||
)
|
||||
|
||||
model = OmegaConf.create(
|
||||
{
|
||||
"nemo_path": "sap_entity_linking.nemo",
|
||||
"max_seq_length": 128,
|
||||
"language_model": language_model,
|
||||
"tokenizer": tokenizer,
|
||||
"train_ds": None,
|
||||
"validation_ds": None,
|
||||
}
|
||||
)
|
||||
|
||||
cfg = OmegaConf.create({"model": model})
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
class TestEntityLinkingModel:
|
||||
@pytest.mark.unit
|
||||
def test_creation_saving_restoring(self):
|
||||
# Create a new temporary directory
|
||||
with tempfile.TemporaryDirectory() as restore_dir:
|
||||
with tempfile.TemporaryDirectory() as save_dir:
|
||||
model = EntityLinkingModel(cfg=get_cfg().model)
|
||||
assert isinstance(model, EntityLinkingModel)
|
||||
|
||||
save_dir_path = save_dir
|
||||
|
||||
# Where model will be saved
|
||||
model_save_path = os.path.join(save_dir, f"{model.__class__.__name__}.nemo")
|
||||
model.save_to(save_path=model_save_path)
|
||||
|
||||
# Where model will be restored from
|
||||
model_restore_path = os.path.join(restore_dir, f"{model.__class__.__name__}.nemo")
|
||||
shutil.copy(model_save_path, model_restore_path)
|
||||
|
||||
# at this point save_dir should not exist
|
||||
assert save_dir_path is not None and not os.path.exists(save_dir_path)
|
||||
assert not os.path.exists(model_save_path)
|
||||
assert os.path.exists(model_restore_path)
|
||||
|
||||
# attempt to restore
|
||||
model_copy = model.__class__.restore_from(restore_path=model_restore_path)
|
||||
assert model.num_weights == model_copy.num_weights
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = TestEntityLinkingModel()
|
||||
t.test_creation_saving_restoring()
|
||||
@@ -0,0 +1,86 @@
|
||||
# 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 pytest
|
||||
import torch
|
||||
|
||||
import nemo.collections.nlp as nemo_nlp
|
||||
from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
|
||||
from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
|
||||
|
||||
|
||||
def do_export(model, name: str):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Generate filename in the temporary directory.
|
||||
tmp_file_name = os.path.join(tmpdir, name + '.onnx')
|
||||
# Test export.
|
||||
if torch.cuda.is_available():
|
||||
model = model.cuda()
|
||||
model.export(tmp_file_name)
|
||||
|
||||
|
||||
class TestHuggingFace:
|
||||
@pytest.mark.unit
|
||||
def test_list_pretrained_models(self):
|
||||
pretrained_lm_models = nemo_nlp.modules.get_pretrained_lm_models_list()
|
||||
assert len(pretrained_lm_models) > 0
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_bert_model(self):
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name='bert-base-uncased')
|
||||
assert isinstance(model, nemo_nlp.modules.BertEncoder)
|
||||
do_export(model, "bert-base-uncased")
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_distilbert_model(self):
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name='distilbert-base-uncased')
|
||||
assert isinstance(model, nemo_nlp.modules.DistilBertEncoder)
|
||||
do_export(model, "distilbert-base-uncased")
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_roberta_model(self):
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name='roberta-base')
|
||||
assert isinstance(model, nemo_nlp.modules.RobertaEncoder)
|
||||
do_export(model, "roberta-base-uncased")
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_albert_model(self):
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name='albert-base-v1')
|
||||
assert isinstance(model, nemo_nlp.modules.AlbertEncoder)
|
||||
do_export(model, "albert-base-v1")
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_chinese_bert_wwm_model(self):
|
||||
model_name = 'hfl/chinese-bert-wwm'
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name)
|
||||
assert isinstance(model, nemo_nlp.modules.BertModule)
|
||||
tokenizer = get_tokenizer(tokenizer_name=model_name)
|
||||
assert isinstance(tokenizer, AutoTokenizer)
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
def test_get_pretrained_arabic_model(self):
|
||||
model_name = 'asafaya/bert-base-arabic'
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name)
|
||||
assert isinstance(model, nemo_nlp.modules.BertModule)
|
||||
tokenizer = get_tokenizer(tokenizer_name=model_name)
|
||||
assert isinstance(tokenizer, AutoTokenizer)
|
||||
@@ -0,0 +1,76 @@
|
||||
# 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.
|
||||
|
||||
try:
|
||||
import apex
|
||||
|
||||
apex_available = True
|
||||
except Exception:
|
||||
apex_available = False
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import onnx
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import nemo.collections.nlp as nemo_nlp
|
||||
from nemo.core.classes import typecheck
|
||||
|
||||
|
||||
def get_pretrained_bert_345m_uncased_model():
|
||||
model_name = "megatron-bert-345m-uncased"
|
||||
model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name)
|
||||
if torch.cuda.is_available():
|
||||
model = model.cuda()
|
||||
return model
|
||||
|
||||
|
||||
class TestMegatron:
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_list_pretrained_models(self):
|
||||
pretrained_lm_models = nemo_nlp.modules.get_pretrained_lm_models_list()
|
||||
assert len(pretrained_lm_models) > 0
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.skip("Only one Megatron model is allowed")
|
||||
def test_get_model(self):
|
||||
model = get_pretrained_bert_345m_uncased_model()
|
||||
assert isinstance(model, nemo_nlp.modules.MegatronBertEncoder)
|
||||
|
||||
typecheck.set_typecheck_enabled(enabled=False)
|
||||
inp = model.input_example()
|
||||
out = model.forward(*inp)
|
||||
typecheck.set_typecheck_enabled(enabled=True)
|
||||
|
||||
@pytest.mark.skipif(not os.path.exists('/home/TestData/nlp'), reason='Not a Jenkins machine')
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_onnx_export(self):
|
||||
model = get_pretrained_bert_345m_uncased_model()
|
||||
assert model
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Generate filename in the temporary directory.
|
||||
# Test export.
|
||||
model.export(os.path.join(".", "megatron.onnx"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = TestMegatron()
|
||||
t.test_onnx_export()
|
||||
@@ -0,0 +1,176 @@
|
||||
# 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
|
||||
import pytorch_lightning as pl
|
||||
import torch
|
||||
import wget
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from nemo.collections import nlp as nemo_nlp
|
||||
from nemo.collections.nlp.models import IntentSlotClassificationModel
|
||||
from nemo.collections.nlp.modules.common import (
|
||||
SequenceClassifier,
|
||||
SequenceRegression,
|
||||
SequenceTokenClassifier,
|
||||
TokenClassifier,
|
||||
)
|
||||
|
||||
|
||||
def classifier_export(obj):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, obj.__class__.__name__ + '.onnx')
|
||||
obj = obj.cuda()
|
||||
obj.export(output=filename)
|
||||
|
||||
|
||||
class TestExportableClassifiers:
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_token_classifier_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(TokenClassifier(hidden_size=256, num_layers=num_layers, num_classes=16))
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_bert_pretraining_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(TokenClassifier(hidden_size=256, num_layers=num_layers, num_classes=16))
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_sequence_token_classifier_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(
|
||||
SequenceTokenClassifier(hidden_size=256, num_slots=8, num_intents=8, num_layers=num_layers)
|
||||
)
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_sequence_classifier_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(SequenceClassifier(hidden_size=256, num_classes=16, num_layers=num_layers))
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_sequence_regression_export_to_onnx(self):
|
||||
for num_layers in [1, 2, 4]:
|
||||
classifier_export(SequenceRegression(hidden_size=256, num_layers=num_layers))
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_IntentSlotClassificationModel_export_to_onnx(self, dummy_data):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
wget.download(
|
||||
'https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/'
|
||||
'nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml',
|
||||
tmpdir,
|
||||
)
|
||||
config_file = os.path.join(tmpdir, 'intent_slot_classification_config.yaml')
|
||||
config = OmegaConf.load(config_file)
|
||||
config = OmegaConf.create(OmegaConf.to_container(config, resolve=True))
|
||||
config.model.data_dir = dummy_data
|
||||
config.trainer.gpus = 1
|
||||
config.trainer.precision = 32
|
||||
config.trainer.accelerator = None
|
||||
trainer = pl.Trainer(**config.trainer)
|
||||
model = IntentSlotClassificationModel(config.model, trainer=trainer)
|
||||
filename = os.path.join(tmpdir, 'isc.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 == 'input_ids'
|
||||
assert onnx_model.graph.input[1].name == 'attention_mask'
|
||||
assert onnx_model.graph.input[2].name == 'token_type_ids'
|
||||
assert onnx_model.graph.output[0].name == 'intent_logits'
|
||||
assert onnx_model.graph.output[1].name == 'slot_logits'
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.skip('Are we exporting to torchscript?')
|
||||
def test_IntentSlotClassificationModel_export_to_torchscript(self, dummy_data):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
wget.download(
|
||||
'https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/'
|
||||
'nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml',
|
||||
tmpdir,
|
||||
)
|
||||
config_file = os.path.join(tmpdir, 'intent_slot_classification_config.yaml')
|
||||
config = OmegaConf.load(config_file)
|
||||
config = OmegaConf.create(OmegaConf.to_container(config, resolve=True))
|
||||
config.model.data_dir = dummy_data
|
||||
config.trainer.gpus = 1
|
||||
config.trainer.precision = 32
|
||||
config.trainer.accelerator = None
|
||||
trainer = pl.Trainer(**config.trainer)
|
||||
model = IntentSlotClassificationModel(config.model, trainer=trainer)
|
||||
filename = os.path.join(tmpdir, 'isc.pt')
|
||||
model.export(output=filename, check_trace=True)
|
||||
torchscript_model = torch.jit.load(filename)
|
||||
inputs = list(torchscript_model.graph.inputs())
|
||||
assert inputs[1].debugName() == 'input_ids.1'
|
||||
assert inputs[2].debugName() == 'input.1'
|
||||
assert inputs[3].debugName() == 'attention_mask.1'
|
||||
modules = list(torchscript_model.modules())
|
||||
assert modules[2].original_name == 'BertEmbeddings'
|
||||
assert modules[46].original_name == 'BertSelfAttention'
|
||||
assert modules[217].original_name == 'SequenceTokenClassifier'
|
||||
assert modules[219].original_name == 'MultiLayerPerceptron'
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_TokenClassificationModel_export_to_onnx(self):
|
||||
model = nemo_nlp.models.TokenClassificationModel.from_pretrained(model_name="ner_en_bert")
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, 'ner.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 == 'input_ids'
|
||||
assert onnx_model.graph.output[0].name == 'logits'
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_PunctuationCapitalizationModel_export_to_onnx(self):
|
||||
model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(model_name="punctuation_en_distilbert")
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, 'puncap.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 == 'input_ids'
|
||||
# assert onnx_model.graph.input[2].name == 'token_type_ids'
|
||||
assert onnx_model.graph.output[0].name == 'punct_logits'
|
||||
assert onnx_model.graph.output[1].name == 'capit_logits'
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_QAModel_export_to_onnx(self):
|
||||
model = nemo_nlp.models.QAModel.from_pretrained(model_name="qa_squadv2.0_bertbase")
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, 'qa.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 == 'input_ids'
|
||||
assert onnx_model.graph.input[2].name == 'token_type_ids'
|
||||
assert onnx_model.graph.output[0].name == 'logits'
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def dummy_data(test_data_dir):
|
||||
return os.path.join(test_data_dir, 'nlp', 'dummy_data')
|
||||
@@ -0,0 +1,144 @@
|
||||
# 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, OmegaConf
|
||||
|
||||
from nemo.collections.nlp.models import MTEncDecModel
|
||||
from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import AAYNBaseConfig
|
||||
|
||||
|
||||
def export_test(model, suffix, try_script=False):
|
||||
with tempfile.TemporaryDirectory() as restore_folder:
|
||||
enc_filename = os.path.join(restore_folder, 'nmt_e' + suffix)
|
||||
dec_filename = os.path.join(restore_folder, 'nmt_d' + suffix)
|
||||
model.encoder.export(output=enc_filename, try_script=try_script)
|
||||
model.decoder.export(output=dec_filename, try_script=try_script)
|
||||
assert os.path.exists(enc_filename)
|
||||
assert os.path.exists(dec_filename)
|
||||
|
||||
|
||||
def get_cfg():
|
||||
cfg = AAYNBaseConfig()
|
||||
cfg.encoder_tokenizer.tokenizer_name = 'yttm'
|
||||
cfg.encoder_tokenizer.tokenizer_model = 'tests/.data/yttm.4096.en-de.model'
|
||||
cfg.decoder_tokenizer.tokenizer_name = 'yttm'
|
||||
cfg.decoder_tokenizer.tokenizer_model = 'tests/.data/yttm.4096.en-de.model'
|
||||
cfg.train_ds = None
|
||||
cfg.validation_ds = None
|
||||
cfg.test_ds = None
|
||||
return cfg
|
||||
|
||||
|
||||
class TestMTEncDecModel:
|
||||
@pytest.mark.unit
|
||||
def test_creation_saving_restoring(self):
|
||||
model = MTEncDecModel(cfg=get_cfg())
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
# Create a new temporary directory
|
||||
with tempfile.TemporaryDirectory() as restore_folder:
|
||||
with tempfile.TemporaryDirectory() as save_folder:
|
||||
save_folder_path = save_folder
|
||||
# Where model will be saved
|
||||
model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo")
|
||||
model.save_to(save_path=model_save_path)
|
||||
# Where model will be restored from
|
||||
model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo")
|
||||
shutil.copy(model_save_path, model_restore_path)
|
||||
# at this point save_folder should not exist
|
||||
assert save_folder_path is not None and not os.path.exists(save_folder_path)
|
||||
assert not os.path.exists(model_save_path)
|
||||
assert os.path.exists(model_restore_path)
|
||||
# attempt to restore
|
||||
model_copy = model.__class__.restore_from(restore_path=model_restore_path)
|
||||
assert model.num_weights == model_copy.num_weights
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_no_artifact_name_collision(self):
|
||||
model = MTEncDecModel(cfg=get_cfg())
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
with tempfile.TemporaryDirectory() as tmpdir1:
|
||||
model.save_to("nmt_model.nemo")
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
model._unpack_nemo_file(path2file="nmt_model.nemo", out_folder=tmpdir)
|
||||
conf = OmegaConf.load(os.path.join(tmpdir, "model_config.yaml"))
|
||||
# Make sure names now differ in saved config
|
||||
assert conf.encoder_tokenizer.tokenizer_model != conf.decoder_tokenizer.tokenizer_model
|
||||
# Make sure names in config start with "nemo:" prefix
|
||||
assert conf.encoder_tokenizer.tokenizer_model.startswith("nemo:")
|
||||
assert conf.decoder_tokenizer.tokenizer_model.startswith("nemo:")
|
||||
# Check if both tokenizers were included
|
||||
assert os.path.exists(os.path.join(tmpdir, conf.encoder_tokenizer.tokenizer_model[5:]))
|
||||
assert os.path.exists(os.path.join(tmpdir, conf.decoder_tokenizer.tokenizer_model[5:]))
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_train_eval_loss(self):
|
||||
cfg = get_cfg()
|
||||
cfg.label_smoothing = 0.5
|
||||
model = MTEncDecModel(cfg=cfg)
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
batch_size = 10
|
||||
time = 32
|
||||
vocab_size = 32000
|
||||
torch.manual_seed(42)
|
||||
tgt_ids = torch.LongTensor(batch_size, time).random_(1, model.decoder_tokenizer.vocab_size)
|
||||
logits = torch.FloatTensor(batch_size, time, vocab_size).random_(-1, 1)
|
||||
log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
|
||||
train_loss = model.loss_fn(log_probs=log_probs, labels=tgt_ids)
|
||||
eval_loss = model.eval_loss_fn(log_probs=log_probs, labels=tgt_ids)
|
||||
assert not torch.allclose(train_loss, eval_loss) # , (train_loss, eval_loss)
|
||||
|
||||
cfg.label_smoothing = 0
|
||||
model = MTEncDecModel(cfg=cfg)
|
||||
# Train loss == val loss when label smoothing = 0
|
||||
train_loss = model.loss_fn(log_probs=log_probs, labels=tgt_ids)
|
||||
eval_loss = model.eval_loss_fn(log_probs=log_probs, labels=tgt_ids)
|
||||
assert torch.allclose(train_loss, eval_loss)
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_cpu_export_onnx(self):
|
||||
model = MTEncDecModel(cfg=get_cfg())
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
export_test(model, ".ts")
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_cpu_export_ts(self):
|
||||
model = MTEncDecModel(cfg=get_cfg())
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
export_test(model, ".onnx")
|
||||
|
||||
@pytest.mark.skipif(not os.path.exists('/home/TestData/nlp'), reason='Not a Jenkins machine')
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_gpu_export_ts(self):
|
||||
model = MTEncDecModel(cfg=get_cfg()).cuda()
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
export_test(model, ".ts")
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_gpu_export_onnx(self):
|
||||
model = MTEncDecModel(cfg=get_cfg()).cuda()
|
||||
assert isinstance(model, MTEncDecModel)
|
||||
export_test(model, ".onnx")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = TestMTEncDecModel()
|
||||
# t.test_gpu_export_ts()
|
||||
t.test_train_eval_loss()
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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
|
||||
from pathlib import Path
|
||||
from shutil import rmtree
|
||||
from unittest import TestCase
|
||||
|
||||
import pytest
|
||||
import pytorch_lightning as pl
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
import nemo.collections.nlp.models as models
|
||||
|
||||
|
||||
def get_metrics(data_dir, model):
|
||||
trainer = pl.Trainer(gpus=[0])
|
||||
|
||||
model.set_trainer(trainer)
|
||||
model.update_data_dir(data_dir)
|
||||
|
||||
test_ds = OmegaConf.create(
|
||||
{
|
||||
'text_file': 'text_dev.txt',
|
||||
'labels_file': 'labels_dev.txt',
|
||||
'shuffle': False,
|
||||
'num_samples': -1,
|
||||
'batch_size': 8,
|
||||
}
|
||||
)
|
||||
|
||||
model._cfg.dataset.use_cache = False
|
||||
trainer.test(model)
|
||||
model.setup_test_data(test_data_config=test_ds)
|
||||
metrics = trainer.test(model)[0]
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
def data_exists(data_dir):
|
||||
return os.path.exists(data_dir)
|
||||
|
||||
|
||||
class TestPretrainedModelPerformance:
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.skipif(
|
||||
not data_exists('/home/TestData/nlp/token_classification_punctuation/fisher'), reason='Not a Jenkins machine'
|
||||
)
|
||||
def test_punct_capit_with_bert(self):
|
||||
data_dir = '/home/TestData/nlp/token_classification_punctuation/fisher'
|
||||
model = models.PunctuationCapitalizationModel.from_pretrained("punctuation_en_bert")
|
||||
metrics = get_metrics(data_dir, model)
|
||||
|
||||
assert abs(metrics['punct_precision'] - 52.3024) < 0.001
|
||||
assert abs(metrics['punct_recall'] - 58.9220) < 0.001
|
||||
assert abs(metrics['punct_f1'] - 53.2976) < 0.001
|
||||
assert abs(metrics['capit_precision'] - 87.0707) < 0.001
|
||||
assert abs(metrics['capit_recall'] - 87.0707) < 0.001
|
||||
assert abs(metrics['capit_f1'] - 87.0707) < 0.001
|
||||
assert int(model.punct_class_report.total_examples) == 128
|
||||
|
||||
preds_512 = model.add_punctuation_capitalization(['what can i do for you today'], max_seq_length=512)[0]
|
||||
assert preds_512 == 'What can I do for you today?'
|
||||
preds_5 = model.add_punctuation_capitalization(['what can i do for you today'], max_seq_length=5)[0]
|
||||
assert preds_5 == 'What can i'
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.skipif(
|
||||
not data_exists('/home/TestData/nlp/token_classification_punctuation/fisher'), reason='Not a Jenkins machine'
|
||||
)
|
||||
def test_punct_capit_with_distilbert(self):
|
||||
data_dir = '/home/TestData/nlp/token_classification_punctuation/fisher'
|
||||
model = models.PunctuationCapitalizationModel.from_pretrained("punctuation_en_distilbert")
|
||||
metrics = get_metrics(data_dir, model)
|
||||
|
||||
assert abs(metrics['punct_precision'] - 53.0826) < 0.001
|
||||
assert abs(metrics['punct_recall'] - 56.2905) < 0.001
|
||||
assert abs(metrics['punct_f1'] - 52.4225) < 0.001
|
||||
assert int(model.punct_class_report.total_examples) == 128
|
||||
|
||||
@pytest.mark.with_downloads()
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.skipif(
|
||||
not data_exists('/home/TestData/nlp/token_classification_punctuation/gmb'), reason='Not a Jenkins machine'
|
||||
)
|
||||
def test_ner_model(self):
|
||||
data_dir = '/home/TestData/nlp/token_classification_punctuation/gmb'
|
||||
model = models.TokenClassificationModel.from_pretrained("ner_en_bert")
|
||||
metrics = get_metrics(data_dir, model)
|
||||
|
||||
assert abs(metrics['precision'] - 96.0937) < 0.001
|
||||
assert abs(metrics['recall'] - 96.0146) < 0.001
|
||||
assert abs(metrics['f1'] - 95.6076) < 0.001
|
||||
assert int(model.classification_report.total_examples) == 202
|
||||
@@ -0,0 +1,43 @@
|
||||
# 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
|
||||
import torch
|
||||
|
||||
from nemo.collections.tts.data.datalayers import FastSpeech2Dataset
|
||||
|
||||
|
||||
class TestTTSDatasets:
|
||||
@pytest.mark.unit
|
||||
def test_fs2_dataset(self, test_data_dir):
|
||||
manifest_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/manifest.json')
|
||||
mappings_file = os.path.join(test_data_dir, 'tts/mini_ljspeech/mappings.json')
|
||||
ignore_file = os.path.join(test_data_dir, 'tts/mini_ljspeech/wavs_to_ignore.pkl')
|
||||
|
||||
# Test loading data (including supplementary data) with ignore file
|
||||
ds = FastSpeech2Dataset(
|
||||
manifest_filepath=manifest_path,
|
||||
mappings_filepath=mappings_file,
|
||||
sample_rate=22050,
|
||||
ignore_file=ignore_file,
|
||||
load_supplementary_values=True,
|
||||
)
|
||||
|
||||
assert len(ds) == 4
|
||||
count = 0
|
||||
for _ in ds:
|
||||
count += 1
|
||||
assert count == 4
|
||||
@@ -0,0 +1,62 @@
|
||||
# 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, OmegaConf
|
||||
|
||||
from nemo.collections.tts.models import FastPitchModel, HifiGanModel, TalkNetSpectModel, WaveGlowModel
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def fastpitch_model():
|
||||
test_root = os.path.dirname(os.path.abspath(__file__))
|
||||
conf = OmegaConf.load(os.path.join(test_root, '../../../examples/tts/conf/fastpitch.yaml'))
|
||||
conf.train_dataset = conf.validation_datasets = '.'
|
||||
conf.model.train_ds = conf.model.test_ds = conf.model.validation_ds = None
|
||||
model = FastPitchModel(cfg=conf.model)
|
||||
return model
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def hifigan_model():
|
||||
test_root = os.path.dirname(os.path.abspath(__file__))
|
||||
model = HifiGanModel.from_pretrained(model_name="tts_hifigan")
|
||||
return model
|
||||
|
||||
|
||||
class TestExportable:
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_FastPitchModel_export_to_onnx(self, fastpitch_model):
|
||||
model = fastpitch_model.cuda()
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, 'fp.onnx')
|
||||
model.export(output=filename, verbose=True, check_trace=True)
|
||||
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_HifiGanModel_export_to_onnx(self, hifigan_model):
|
||||
model = hifigan_model.cuda()
|
||||
assert hifigan_model.generator is not None
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
filename = os.path.join(tmpdir, 'hfg.pt')
|
||||
model.export(output=filename, verbose=True, check_trace=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = TestExportable()
|
||||
t.test_FastPitchModel_export_to_onnx(fastpitch_model())
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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
|
||||
from unittest import TestCase
|
||||
|
||||
import onnx
|
||||
import pytest
|
||||
import torch
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from nemo.collections.tts.models import WaveGlowModel
|
||||
from nemo.collections.tts.modules import WaveGlowModule
|
||||
from nemo.core.classes import typecheck
|
||||
|
||||
mcfg = DictConfig(
|
||||
{
|
||||
"_target_": "nemo.collections.tts.modules.waveglow.WaveGlowModule",
|
||||
"n_flows": 12,
|
||||
"n_group": 8,
|
||||
"n_mel_channels": 80,
|
||||
"n_early_every": 4,
|
||||
"n_early_size": 2,
|
||||
"n_wn_channels": 512,
|
||||
"n_wn_layers": 8,
|
||||
"wn_kernel_size": 3,
|
||||
}
|
||||
)
|
||||
|
||||
pcfg = DictConfig(
|
||||
{
|
||||
"_target_": "nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures",
|
||||
"dither": 0.0,
|
||||
"nfilt": 80,
|
||||
"stft_conv": True,
|
||||
}
|
||||
)
|
||||
|
||||
wcfg = DictConfig({"waveglow": mcfg, "sigma": 1.0, "preprocessor": pcfg,})
|
||||
|
||||
|
||||
def input_example(sz):
|
||||
mel = torch.randn(1, 1, 80, sz).cuda().half()
|
||||
z = torch.randn(1, 8, sz * 256 // 8, 1).cuda().half()
|
||||
return (
|
||||
mel,
|
||||
z,
|
||||
)
|
||||
|
||||
|
||||
def taco2wg(spec, z):
|
||||
spec = spec.permute(0, 3, 2, 1).contiguous()
|
||||
return spec.view(spec.size(0), spec.size(1), -1), z.view(z.size(0), z.size(1), -1)
|
||||
|
||||
|
||||
# Wrapper method to convert Jasper's Taco2 output to WG input and call inference
|
||||
def forward_wrapper(self, spec, z=None):
|
||||
spec, z = taco2wg(spec, z)
|
||||
audio = self.waveglow.norm_dist_to_audio(spec=spec, sigma=1.0, z=z)
|
||||
return audio
|
||||
|
||||
|
||||
class TestWaveGlow:
|
||||
@pytest.mark.run_only_on('GPU')
|
||||
@pytest.mark.unit
|
||||
def test_export_to_onnx(self):
|
||||
model = WaveGlowModel(wcfg)
|
||||
model = model.cuda().half()
|
||||
typecheck.set_typecheck_enabled(enabled=False)
|
||||
with tempfile.TemporaryDirectory() as tmpdir, model.nemo_infer():
|
||||
# Generate filename in the temporary directory.
|
||||
# TODO: Change `waveglow.ts` to `waveglow.onnx` for > 21.05
|
||||
tmp_file_name = os.path.join("waveglow.ts")
|
||||
|
||||
n_mels = 80
|
||||
# Test export.
|
||||
inp = input_example(n_mels)
|
||||
inp1 = taco2wg(*inp)
|
||||
inp2 = inp1
|
||||
res1 = model.waveglow(*inp1)
|
||||
res2 = model.waveglow(*inp2)
|
||||
assert torch.allclose(res1, res2, rtol=0.01, atol=0.1)
|
||||
WaveGlowModel.forward_for_export = forward_wrapper
|
||||
model.export(
|
||||
tmp_file_name,
|
||||
verbose=True,
|
||||
input_example=inp,
|
||||
output_example=res1,
|
||||
try_script=False,
|
||||
check_trace=False,
|
||||
do_constant_folding=True,
|
||||
dynamic_axes={"spec": [0], "z": [0], "audio": [0]},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = TestWaveGlow()
|
||||
t.test_export_to_onnx()
|
||||
Reference in New Issue
Block a user