diff --git a/nam/models/linear.py b/nam/models/linear.py index 622bf2cd..25f43023 100644 --- a/nam/models/linear.py +++ b/nam/models/linear.py @@ -32,6 +32,7 @@ def import_weights(self, weights): self._net.weight.data = ( _torch.Tensor([w for w in weights[: self._net.weight.numel()]]) .reshape(self._net.weight.shape) + .flip(-1) .to(self._net.weight.device) ) if self._bias: @@ -55,7 +56,7 @@ def _export_config(self): } def _export_weights(self) -> _np.ndarray: - params_list = [self._net.weight.flatten()] + params_list = [self._net.weight.flip(-1).flatten()] if self._bias: params_list.append(self._net.bias.flatten()) params = _torch.cat(params_list).detach().cpu().numpy() diff --git a/tests/test_nam/test_models/test_linear.py b/tests/test_nam/test_models/test_linear.py index f65d032a..d8afd5d6 100644 --- a/tests/test_nam/test_models/test_linear.py +++ b/tests/test_nam/test_models/test_linear.py @@ -3,6 +3,7 @@ # Author: Steven Atkinson (steven@atkinson.mn) import pytest as _pytest +import torch as _torch from nam.models import linear as _linear @@ -16,3 +17,51 @@ def setup_class(cls): args = () kwargs = {"receptive_field": 2, "sample_rate": 44100} super().setup_class(C, args, kwargs) + + def test_export_weights_are_chronological(self): + model = _linear.Linear(receptive_field=3) + model._net.weight.data.copy_(_torch.tensor([[[0.5, -0.25, 0.125]]])) + + exported_weights = model._export_weights() + impulse_response = model(_torch.tensor([1.0, 0.0, 0.0])) + + _torch.testing.assert_close( + _torch.from_numpy(exported_weights), impulse_response + ) + _torch.testing.assert_close( + _torch.from_numpy(exported_weights), + _torch.tensor([0.125, -0.25, 0.5]), + ) + + def test_import_weights_are_chronological(self): + weights = _torch.tensor([0.5, -0.25, 0.125]) + model = _linear.Linear(receptive_field=3) + + model.import_weights(weights) + + _torch.testing.assert_close( + model._net.weight.data, _torch.tensor([[[0.125, -0.25, 0.5]]]) + ) + _torch.testing.assert_close( + model(_torch.tensor([1.0, 0.0, 0.0])), weights + ) + + def test_import_export_weights_round_trip_with_bias(self): + model = _linear.Linear(receptive_field=3, bias=True) + model._net.weight.data.copy_(_torch.tensor([[[0.5, -0.25, 0.125]]])) + model._net.bias.data.copy_(_torch.tensor([0.75])) + exported_weights = model._export_weights() + model2 = _linear.Linear(receptive_field=3, bias=True) + + model2.import_weights(exported_weights) + + _torch.testing.assert_close( + _torch.from_numpy(exported_weights), + _torch.tensor([0.125, -0.25, 0.5, 0.75]), + ) + _torch.testing.assert_close(model2._net.weight, model._net.weight) + _torch.testing.assert_close(model2._net.bias, model._net.bias) + _torch.testing.assert_close( + _torch.from_numpy(model2._export_weights()), + _torch.from_numpy(exported_weights), + ) diff --git a/tests/test_nam/test_models/test_sequential.py b/tests/test_nam/test_models/test_sequential.py index 19e206c4..6d1e1e1f 100644 --- a/tests/test_nam/test_models/test_sequential.py +++ b/tests/test_nam/test_models/test_sequential.py @@ -232,6 +232,18 @@ def test_batch_processing(self): assert y.shape[0] == 3 # Batch dimension preserved assert y.shape[1] == x.shape[1] - seq_model.receptive_field + 1 + def test_export_weights_uses_chronological_linear_taps(self): + linear1 = _linear.Linear(receptive_field=3) + linear1._net.weight.data.copy_(_torch.tensor([[[1.0, 2.0, 3.0]]])) + linear2 = _linear.Linear(receptive_field=1) + linear2._net.weight.data.copy_(_torch.tensor([[[4.0]]])) + model = _sequential.Sequential(models=[linear1, linear2]) + + _torch.testing.assert_close( + _torch.from_numpy(model._export_weights()), + _torch.tensor([3.0, 2.0, 1.0, 4.0]), + ) + if __name__ == "__main__": _pytest.main()