def test_separation_model_extra_modules(one_item):
with tempfile.NamedTemporaryFile(suffix='.json', delete=True) as tmp:
dpcl_config['modules']['test'] = {'class': 'MyModule'}
dpcl_config['connections'].append(('test', ('mix_magnitude', {
'embedding': 'embedding',
'flip': False
})))
dpcl_config['output'].append('test')
with open(tmp.name, 'w') as f:
json.dump(dpcl_config, f)
configs = [dpcl_config, tmp.name, json.dumps(dpcl_config)]
nussl.ml.register_module(MyModule)
for config in configs:
model = SeparationModel(config)
output = model(one_item)
assert (output['embedding'].shape == (
one_item['mix_magnitude'].shape + (20, )))
assert torch.allclose(one_item['mix_magnitude'], output['test'])
model = SeparationModel(config)
copy_one_item = copy.deepcopy(one_item)
output = model(copy_one_item, flip=True)
assert torch.allclose(one_item['mix_magnitude'], -output['test'])
def test_separation_model_save_and_load():
model = SeparationModel(dpcl_config)
tfms = datasets.transforms.Compose([
datasets.transforms.PhaseSensitiveSpectrumApproximation(),
datasets.transforms.ToSeparationModel(),
datasets.transforms.Cache('tests/local/sep_model/cache')
])
class DummyData:
def __init__(self):
self.stft_params = None
self.sample_rate = None
self.num_channels = None
self.metadata = {'transforms': tfms}
class DummyState:
def __init__(self):
self.epoch = 0
self.epoch_length = 100
self.max_epochs = 100
self.output = None
self.metrics = {}
self.seed = None
self.epoch_history = {}
class DummyTrainer:
def __init__(self):
self.state = DummyState()
dummy_data = DummyData()
dummy_trainer = DummyTrainer()
with tempfile.NamedTemporaryFile(suffix='.pth', delete=True) as tmp:
loc = model.save(tmp.name,
train_data=dummy_data,
val_data=dummy_data,
trainer=dummy_trainer)
new_model, metadata = SeparationModel.load(tmp.name)
assert metadata['nussl_version'] == nussl.__version__
new_model_params = {}
old_model_params = {}
for name, param in new_model.named_parameters():
new_model_params[name] = param
for name, param in model.named_parameters():
old_model_params[name] = param
for key in new_model_params:
assert torch.allclose(new_model_params[key], old_model_params[key])
def test_separation_end_to_end(one_item):
for c in [end_to_end_real_config, end_to_end_complex_config]:
with tempfile.NamedTemporaryFile(suffix='.json', delete=True) as tmp:
with open(tmp.name, 'w') as f:
json.dump(c, f)
configs = [c, tmp.name, json.dumps(c)]
for config in configs:
model = SeparationModel(config)
output = model(one_item)
assert (
output['audio'].shape == one_item['source_audio'].shape)
def test_separation_dprnn(one_item):
# dprnn network
with tempfile.NamedTemporaryFile(suffix='.json', delete=True) as tmp:
with open(tmp.name, 'w') as f:
json.dump(dual_path_recurrent_config, f)
configs = [
dual_path_recurrent_config, tmp.name,
json.dumps(dual_path_recurrent_config)
]
for config in configs:
model = SeparationModel(config, verbose=True)
output = model(one_item)
assert (output['audio'].shape == one_item['source_audio'].shape)
def test_separation_model_dpcl(one_item):
n_features = one_item['mix_magnitude'].shape[2]
# dpcl network
with tempfile.NamedTemporaryFile(suffix='.json', delete=True) as tmp:
with open(tmp.name, 'w') as f:
json.dump(dpcl_config, f)
configs = [dpcl_config, tmp.name, json.dumps(dpcl_config)]
for config in configs:
model = SeparationModel(config)
output = model(one_item)
assert (output['embedding'].shape == (
one_item['mix_magnitude'].shape + (20, )))
def test_separation_model_save():
model = SeparationModel(dpcl_config)
with tempfile.NamedTemporaryFile(suffix='.pth', delete=True) as tmp:
loc = model.save(tmp.name)
checkpoint = torch.load(loc)
assert checkpoint['metadata']['nussl_version'] == nussl.__version__
new_model = SeparationModel(checkpoint['config'])
new_model.load_state_dict(checkpoint['state_dict'])
new_model_params = {}
old_model_params = {}
for name, param in new_model.named_parameters():
new_model_params[name] = param
for name, param in model.named_parameters():
old_model_params[name] = param
for key in new_model_params:
assert torch.allclose(new_model_params[key], old_model_params[key])
def test_separation_model_split(one_item):
n_features = one_item['mix_magnitude'].shape[2]
with tempfile.NamedTemporaryFile(suffix='.json', delete=True) as tmp:
with open(tmp.name, 'w') as f:
json.dump(split_config, f)
configs = [split_config, tmp.name, json.dumps(split_config)]
for config in configs:
model = SeparationModel(config)
output = model(one_item)
assert (output['estimates'].shape == (
one_item['mix_magnitude'].shape + (2, )))
assert (torch.allclose(output['estimates'].sum(dim=-1),
one_item['mix_magnitude']))
assert (output['split:0'].shape[2] == 100)
assert (output['split:1'].shape[2] == 157)
def test_separation_model_open_unmix_like(one_item):
n_features = one_item['mix_magnitude'].shape[2]
with tempfile.NamedTemporaryFile(suffix='.json', delete=True) as tmp:
with open(tmp.name, 'w') as f:
json.dump(open_unmix_like_config, f)
configs = [
open_unmix_like_config, tmp.name,
json.dumps(open_unmix_like_config)
]
for config in configs:
model = SeparationModel(config)
output = model(one_item)
assert (output['estimates'].shape == (
one_item['mix_magnitude'].shape + (2, )))
assert (output['mask'].shape == (one_item['mix_magnitude'].shape +
(2, )))
assert (output['embedding'].shape == (
one_item['mix_magnitude'].shape + (20, )))
def test_separation_model_mask_inference(one_item):
n_features = one_item['mix_magnitude'].shape[2]
with tempfile.NamedTemporaryFile(suffix='.json', delete=True) as tmp:
with open(tmp.name, 'w') as f:
json.dump(mi_config, f)
configs = [mi_config, tmp.name, json.dumps(mi_config)]
for config in configs:
model = SeparationModel(config)
bad_item = copy.deepcopy(one_item)
bad_item.pop('mix_magnitude')
pytest.raises(ValueError, model, bad_item)
output = model(one_item)
assert (output['estimates'].shape == (
one_item['mix_magnitude'].shape + (2, )))
assert (torch.allclose(output['estimates'].sum(dim=-1),
one_item['mix_magnitude']))
def test_separation_model_gmm_unfold(one_item):
n_features = one_item['mix_magnitude'].shape[2]
with tempfile.NamedTemporaryFile(suffix='.json', delete=True) as tmp:
with open(tmp.name, 'w') as f:
json.dump(gmm_unfold_config, f)
configs = [gmm_unfold_config, tmp.name, json.dumps(gmm_unfold_config)]
for config in configs:
model = SeparationModel(config)
one_item['init_means'] = torch.randn(
one_item['mix_magnitude'].shape[0], 2,
20).to(one_item['mix_magnitude'].device)
output = model(one_item)
assert (output['estimates'].shape == (
one_item['mix_magnitude'].shape + (2, )))
assert (torch.allclose(output['estimates'].sum(dim=-1),
one_item['mix_magnitude']))
assert (output['embedding'].shape == (
one_item['mix_magnitude'].shape + (20, )))
def test_separation_model_repr_and_verbose(one_item):
model = SeparationModel(end_to_end_real_config, verbose=True)
print(model)
model(one_item)
