1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
|
import torch
from parameterized import parameterized
from torchaudio.models import squim_objective_base, squim_subjective_base
from torchaudio_unittest.common_utils import skipIfNoCuda, torch_script, TorchaudioTestCase
class TestSquimObjective(TorchaudioTestCase):
def _smoke_test_objective(self, model, device, dtype):
model = model.to(device=device, dtype=dtype)
model = model.eval()
batch_size, num_frames = 3, 16000
waveforms = torch.randn(batch_size, num_frames, device=device, dtype=dtype)
model(waveforms)
@parameterized.expand([(torch.float32,), (torch.float64,)])
def test_cpu_smoke_test(self, dtype):
model = squim_objective_base()
self._smoke_test_objective(model, torch.device("cpu"), dtype)
@parameterized.expand([(torch.float32,), (torch.float64,)])
@skipIfNoCuda
def test_cuda_smoke_test(self, dtype):
model = squim_objective_base()
self._smoke_test_objective(model, torch.device("cuda"), dtype)
def test_batch_consistency(self):
model = squim_objective_base()
model.eval()
batch_size, num_frames = 3, 16000
waveforms = torch.randn(batch_size, num_frames)
ref_scores = model(waveforms)
hyp_scores = [torch.zeros(batch_size), torch.zeros(batch_size), torch.zeros(batch_size)]
for i in range(batch_size):
scores = model(waveforms[i : i + 1])
for j in range(3):
hyp_scores[j][i] = scores[j]
self.assertEqual(len(hyp_scores), len(ref_scores))
for i in range(len(ref_scores)):
self.assertEqual(hyp_scores[i], ref_scores[i])
def test_torchscript_consistency(self):
model = squim_objective_base()
model.eval()
batch_size, num_frames = 3, 16000
waveforms = torch.randn(batch_size, num_frames)
ref_scores = model(waveforms)
scripted = torch_script(model)
hyp_scores = scripted(waveforms)
self.assertEqual(len(hyp_scores), len(ref_scores))
for i in range(len(ref_scores)):
self.assertEqual(hyp_scores[i], ref_scores[i])
class TestSquimSubjective(TorchaudioTestCase):
def _smoke_test_subjective(self, model, device, dtype):
model = model.to(device=device, dtype=dtype)
model = model.eval()
batch_size, num_frames = 3, 16000
waveforms = torch.randn(batch_size, num_frames, device=device, dtype=dtype)
reference = torch.randn(batch_size, num_frames, device=device, dtype=dtype)
model(waveforms, reference)
@parameterized.expand([(torch.float32,), (torch.float64,)])
def test_cpu_smoke_test(self, dtype):
model = squim_subjective_base()
self._smoke_test_subjective(model, torch.device("cpu"), dtype)
@parameterized.expand([(torch.float32,), (torch.float64,)])
@skipIfNoCuda
def test_cuda_smoke_test(self, dtype):
model = squim_subjective_base()
self._smoke_test_subjective(model, torch.device("cuda"), dtype)
def test_batch_consistency(self):
model = squim_subjective_base()
model.eval()
batch_size, num_frames = 3, 16000
waveforms = torch.randn(batch_size, num_frames)
reference = torch.randn(batch_size, num_frames)
ref_scores = model(waveforms, reference)
hyp_scores = []
for i in range(batch_size):
scores = model(waveforms[i : i + 1], reference[i : i + 1])
hyp_scores.append(scores)
hyp_scores = torch.tensor(hyp_scores)
self.assertEqual(hyp_scores, ref_scores)
def test_torchscript_consistency(self):
model = squim_subjective_base()
model.eval()
batch_size, num_frames = 3, 16000
waveforms = torch.randn(batch_size, num_frames)
reference = torch.randn(batch_size, num_frames)
ref_scores = model(waveforms, reference)
scripted = torch_script(model)
hyp_scores = scripted(waveforms, reference)
self.assertEqual(hyp_scores, ref_scores)
|
