import math import os from typing import Tuple import torch import torch.nn.functional as F from parameterized import parameterized from torchaudio.models.wav2vec2 import ( hubert_base, hubert_large, hubert_pretrain_base, hubert_pretrain_large, hubert_pretrain_xlarge, hubert_xlarge, wav2vec2_base, wav2vec2_large, wav2vec2_large_lv60k, ) from torchaudio_unittest.common_utils import skipIfNoCuda, skipIfNoQengine, torch_script, TorchaudioTestCase TORCH_VERSION: Tuple[int, ...] = tuple(int(x) for x in torch.__version__.split(".")[:2]) if TORCH_VERSION >= (1, 10): import torch.ao.quantization as tq else: import torch.quantization as tq def _name_func(testcase_func, i, param): return f"{testcase_func.__name__}_{i}_{param[0][0].__name__}" factory_funcs = parameterized.expand( [ (wav2vec2_base,), (wav2vec2_large,), (wav2vec2_large_lv60k,), (hubert_base,), (hubert_large,), (hubert_xlarge,), ], name_func=_name_func, ) factory_funcs_hubert_pretrain = parameterized.expand( [ (hubert_pretrain_base,), (hubert_pretrain_large,), (hubert_pretrain_xlarge,), ], name_func=_name_func, ) class TestWav2Vec2Model(TorchaudioTestCase): def _smoke_test(self, model, device, dtype): model = model.to(device=device, dtype=dtype) model = model.eval() batch_size, num_frames = 3, 1024 waveforms = torch.randn(batch_size, num_frames, device=device, dtype=dtype) lengths = torch.randint( low=0, high=num_frames, size=[ batch_size, ], device=device, ) model(waveforms, lengths) @parameterized.expand([(torch.float32,), (torch.float64,)]) def test_cpu_smoke_test(self, dtype): model = wav2vec2_base() self._smoke_test(model, torch.device("cpu"), dtype) model = wav2vec2_base(aux_num_out=32) self._smoke_test(model, torch.device("cpu"), dtype) @parameterized.expand([(torch.float32,), (torch.float64,)]) @skipIfNoCuda def test_cuda_smoke_test(self, dtype): model = wav2vec2_base() self._smoke_test(model, torch.device("cuda"), dtype) model = wav2vec2_base(aux_num_out=32) self._smoke_test(model, torch.device("cuda"), dtype) def _feature_extractor_test(self, model): batch_size, num_frames = 3, 1024 model.eval() num_layers = len(model.encoder.transformer.layers) waveforms = torch.randn(batch_size, num_frames) lengths = torch.randint( low=0, high=num_frames, size=[ batch_size, ], ) # Not providing num_layers returns all the intermediate features from # tranformer layers all_features, lengths_ = model.extract_features(waveforms, lengths, num_layers=None) assert len(all_features) == num_layers for features in all_features: assert features.ndim == 3 assert features.shape[0] == batch_size assert lengths_.shape == torch.Size([batch_size]) # Limiting the number of layers to `l`. for l in range(1, num_layers + 1): features, lengths_ = model.extract_features(waveforms, lengths, num_layers=l) assert len(features) == l for i in range(l): self.assertEqual(all_features[i], features[i]) assert lengths_.shape == torch.Size([batch_size]) @factory_funcs def test_extract_feature(self, factory_func): """`extract_features` method does not fail""" self._feature_extractor_test(factory_func(aux_num_out=32)) def _test_batch_consistency(self, model): model.eval() batch_size, max_frames = 5, 5 * 1024 waveforms = torch.randn(batch_size, max_frames) input_lengths = torch.tensor([i * 3200 for i in range(1, 6)]) # Batch process with lengths batch_logits, output_lengths = model(waveforms, input_lengths) for i in range(batch_size): # Par-sample process without feeding length single_logit, _ = model(waveforms[i : i + 1, : input_lengths[i]], None) batch_logit = batch_logits[i : i + 1, : output_lengths[i]] # Convert to probability so that it's easier to interpretate the diff single_prob = F.softmax(single_logit, dim=2) batch_prob = F.softmax(batch_logit, dim=2) # We allow max atol=0.005 -> 0.5% self.assertEqual(single_prob, batch_prob, atol=0.005, rtol=0) @factory_funcs def test_pretrain_batch_consistency(self, factory_func): """Results from single process and batched process should be reasonably close""" self._test_batch_consistency(factory_func()) @factory_funcs def test_finetune_batch_consistency(self, factory_func): """Results from single process and batched process should be reasonably close""" self._test_batch_consistency(factory_func(aux_num_out=32)) def _test_zero_length(self, model): model.eval() batch_size = 3 waveforms = torch.randn(batch_size, 1024) input_lengths = torch.zeros(batch_size) _, output_lengths = model(waveforms, input_lengths) self.assertEqual(torch.zeros_like(output_lengths), output_lengths) _, output_lengths = model.extract_features(waveforms, input_lengths) self.assertEqual(torch.zeros_like(output_lengths), output_lengths) @factory_funcs def test_pretrain_zero_length(self, factory_func): """Passing zero length should not fail""" self._test_zero_length(factory_func()) @factory_funcs def test_finetune_zero_length(self, factory_func): """Passing zero length should not fail""" self._test_zero_length(factory_func(aux_num_out=32)) def _test_torchscript(self, model): model.eval() batch_size, num_frames = 3, 1024 waveforms = torch.randn(batch_size, num_frames) lengths = torch.randint( low=0, high=num_frames, size=[ batch_size, ], ) ref_out, ref_len = model(waveforms, lengths) scripted = torch_script(model) hyp_out, hyp_len = scripted(waveforms, lengths) self.assertEqual(hyp_out, ref_out) self.assertEqual(hyp_len, ref_len) @factory_funcs def test_pretrain_torchscript(self, factory_func): """Wav2Vec2Model should be scriptable""" if factory_func is hubert_xlarge and os.environ.get("CI") == "true": self.skipTest( "hubert_xlarge is known to fail on CI. " "See https://github.com/pytorch/pytorch/issues/65776" ) self._test_torchscript(factory_func()) @factory_funcs def test_finetune_torchscript(self, factory_func): """Wav2Vec2Model should be scriptable""" if factory_func is hubert_xlarge and os.environ.get("CI") == "true": self.skipTest( "hubert_xlarge is known to fail on CI. " "See https://github.com/pytorch/pytorch/issues/65776" ) self._test_torchscript(factory_func(aux_num_out=32)) def _test_quantize_smoke_test(self, model): model.eval() batch_size, num_frames = 3, 1024 # Remove the weight normalization forward hook model.encoder.transformer.pos_conv_embed.__prepare_scriptable__() quantized = tq.quantize_dynamic(model, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8) # A lazy way to check that Modules are different assert str(quantized) != str(model), "Dynamic quantization did not modify the module." waveforms = torch.randn(batch_size, num_frames) lengths = torch.randint( low=0, high=num_frames, size=[ batch_size, ], ) _, _ = quantized(waveforms, lengths) @factory_funcs @skipIfNoQengine def test_quantize(self, factory_func): """Wav2Vec2Model should support basic quantization""" self._test_quantize_smoke_test(factory_func(aux_num_out=32)) def _test_quantize_torchscript(self, model): model.eval() batch_size, num_frames = 3, 1024 # Remove the weight normalization forward hook model.encoder.transformer.pos_conv_embed.__prepare_scriptable__() quantized = tq.quantize_dynamic(model, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8) # A lazy way to check that Modules are different assert str(quantized) != str(model), "Dynamic quantization did not modify the module." waveforms = torch.randn(batch_size, num_frames) lengths = torch.randint( low=0, high=num_frames, size=[ batch_size, ], ) ref_out, ref_len = quantized(waveforms, lengths) # Script scripted = torch_script(quantized) hyp_out, hyp_len = scripted(waveforms, lengths) self.assertEqual(hyp_out, ref_out) self.assertEqual(hyp_len, ref_len) @factory_funcs @skipIfNoQengine def test_quantize_torchscript(self, factory_func): """Quantized Wav2Vec2Model should be scriptable""" self._test_quantize_torchscript(factory_func(aux_num_out=32)) def _compute_label_frame(audio_frame: int) -> int: """Compute number of frames in the label tensor based on the number of frames in the audio tensor.""" kernel_size = 25 stride = 20 sample_rate = 16 # 16 per millisecond label_frame = math.floor((audio_frame - kernel_size * sample_rate) / (stride * sample_rate)) + 1 return label_frame class TestHuBERTPretrainModel(TorchaudioTestCase): def _smoke_test(self, model, device, dtype): model = model.to(device=device, dtype=dtype) model = model.eval() batch_size, num_frames = 3, 1024 waveforms = torch.randn(batch_size, num_frames, device=device, dtype=dtype) labels = torch.randint( low=0, high=100, size=[ batch_size, _compute_label_frame(num_frames), ], device=device, ) lengths = torch.randint( low=0, high=num_frames, size=[ batch_size, ], device=device, ) model(waveforms, labels, lengths) @parameterized.expand([(torch.float32,), (torch.float64,)]) def test_cpu_smoke_test(self, dtype): model = hubert_pretrain_base() self._smoke_test(model, torch.device("cpu"), dtype) @parameterized.expand([(torch.float32,), (torch.float64,)]) @skipIfNoCuda def test_cuda_smoke_test(self, dtype): model = hubert_pretrain_base() self._smoke_test(model, torch.device("cuda"), dtype) def _feature_extractor_test(self, model): batch_size, num_frames = 3, 1024 model = model.wav2vec2 model.eval() num_layers = len(model.encoder.transformer.layers) waveforms = torch.randn(batch_size, num_frames) lengths = torch.randint( low=0, high=num_frames, size=[ batch_size, ], ) # Not providing num_layers returns all the intermediate features from # tranformer layers all_features, lengths_ = model.extract_features(waveforms, lengths, num_layers=None) assert len(all_features) == num_layers for features in all_features: assert features.ndim == 3 assert features.shape[0] == batch_size assert lengths_.shape == torch.Size([batch_size]) # Limiting the number of layers to `l`. for l in range(1, num_layers + 1): features, lengths_ = model.extract_features(waveforms, lengths, num_layers=l) assert len(features) == l for i in range(l): self.assertEqual(all_features[i], features[i]) assert lengths_.shape == torch.Size([batch_size]) @factory_funcs_hubert_pretrain def test_extract_feature(self, factory_func): """`extract_features` method does not fail""" self._feature_extractor_test(factory_func()) def _test_quantize_smoke_test(self, model): model.eval() batch_size, num_frames = 3, 1024 # Remove the weight normalization forward hook model.wav2vec2.encoder.transformer.pos_conv_embed.__prepare_scriptable__() quantized = tq.quantize_dynamic(model, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8) # A lazy way to check that Modules are different assert str(quantized) != str(model), "Dynamic quantization did not modify the module." waveforms = torch.randn(batch_size, num_frames) labels = torch.randint( low=0, high=100, size=[ batch_size, _compute_label_frame(num_frames), ], ) lengths = torch.randint( low=0, high=num_frames, size=[ batch_size, ], ) _, _, _ = quantized(waveforms, labels, lengths) @factory_funcs_hubert_pretrain @skipIfNoQengine def test_quantize(self, factory_func): """HuBERTPretrainModel should support basic quantization""" self._test_quantize_smoke_test(factory_func())