def test_pmsqe_pit(n_src, sample_rate):
# Define supported STFT
if sample_rate == 16000:
stft = Encoder(STFTFB(kernel_size=512, n_filters=512, stride=256))
else:
stft = Encoder(STFTFB(kernel_size=256, n_filters=256, stride=128))
# Usage by itself
ref, est = torch.randn(2, n_src, 16000), torch.randn(2, n_src, 16000)
ref_spec = transforms.mag(stft(ref))
est_spec = transforms.mag(stft(est))
loss_func = PITLossWrapper(SingleSrcPMSQE(sample_rate=sample_rate), pit_from="pw_pt")
# Assert forward ok.
loss_func(est_spec, ref_spec)
def test_pmsqe(sample_rate):
# Define supported STFT
if sample_rate == 16000:
stft = Encoder(STFTFB(kernel_size=512, n_filters=512, stride=256))
else:
stft = Encoder(STFTFB(kernel_size=256, n_filters=256, stride=128))
# Usage by itself
ref, est = torch.randn(2, 1, 16000), torch.randn(2, 1, 16000)
ref_spec = transforms.mag(stft(ref))
est_spec = transforms.mag(stft(est))
loss_func = SingleSrcPMSQE(sample_rate=sample_rate)
loss_value = loss_func(est_spec, ref_spec)
# Assert output has shape (batch,)
assert loss_value.shape[0] == ref.shape[0]
# Assert support for transposed inputs.
tr_loss_value = loss_func(est_spec.transpose(1, 2), ref_spec.transpose(1, 2))
assert_allclose(loss_value, tr_loss_value)