neural-amp-modeler

Neural network emulator for guitar amplifiers
Log | Files | Refs | README | LICENSE

commit dde0baeb3b51cea4ba6fb5941d899f3448c6ed5f
parent 4e97b46cc24417a550d92a82bce2e1d012f10ed5
Author: Steven Atkinson <steven@atkinson.mn>
Date:   Sun, 30 Apr 2023 11:52:58 -0700

Fix model device handling (#225)

* Models to CPU

* Update GPU kwargs

* End-to-end tests of CLI trainer
Diffstat:
Mbin/export/main.py | 1+
Mbin/train/inputs/data/single_pair.json | 4++++
Mbin/train/inputs/data/two_pairs.json | 4++++
Mbin/train/inputs/learning/default.json | 7++++++-
Mbin/train/inputs/learning/demo.json | 7++++++-
Mbin/train/inputs/models/catlstm.json | 5++++-
Mbin/train/inputs/models/convnet.json | 4++++
Mbin/train/inputs/models/demonet.json | 4++++
Mbin/train/inputs/models/lstm.json | 5++++-
Mbin/train/inputs/models/wavenet.json | 4++++
Mbin/train/main.py | 2++
Atests/test_bin/__init__.py | 0
Atests/test_bin/test_train/__init__.py | 0
Atests/test_bin/test_train/test_main.py | 203+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
14 files changed, 246 insertions(+), 4 deletions(-)

diff --git a/bin/export/main.py b/bin/export/main.py @@ -38,6 +38,7 @@ def main(args): k: Param.init_from_config(v) for k, v in json.load(fp).items() } export_args = (outdir, param_config) + net.cpu() net.eval() outdir.mkdir(parents=True, exist_ok=True) net.export(*export_args, include_snapshot=args.snapshot) diff --git a/bin/train/inputs/data/single_pair.json b/bin/train/inputs/data/single_pair.json @@ -1,4 +1,8 @@ { + "_notes": [ + "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ", + "representative of this!" + ], "train": { "start": null, "stop": -432000, diff --git a/bin/train/inputs/data/two_pairs.json b/bin/train/inputs/data/two_pairs.json @@ -1,4 +1,8 @@ { + "_notes": [ + "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ", + "representative of this!" + ], "train": { "x_path": "C:\\path\\to\\train\\source.wav", "y_path": "C:\\path\\to\\train\\target.wav", diff --git a/bin/train/inputs/learning/default.json b/bin/train/inputs/learning/default.json @@ -1,4 +1,8 @@ { + "_notes": [ + "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ", + "representative of this!" + ], "train_dataloader": { "batch_size": 16, "shuffle": true, @@ -8,7 +12,8 @@ }, "val_dataloader": {}, "trainer": { - "gpus": 1, + "accelerator": "gpu", + "devices": 1, "max_epochs": 100 }, "trainer_fit_kwargs": {} diff --git a/bin/train/inputs/learning/demo.json b/bin/train/inputs/learning/demo.json @@ -1,4 +1,8 @@ { + "_notes": [ + "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ", + "representative of this!" + ], "train_dataloader": { "batch_size": 16, "shuffle": true, @@ -8,7 +12,8 @@ }, "val_dataloader": {}, "trainer": { - "gpus": 0, + "accelerator": "gpu", + "devices": 1, "max_epochs": 10 }, "trainer_fit_kwargs": {} diff --git a/bin/train/inputs/models/catlstm.json b/bin/train/inputs/models/catlstm.json @@ -7,7 +7,10 @@ " non-parametric version, even if you're modeling a fair number of knobs.", " * You'll probably have a much larger dataset, so validating every so often ", " in steps instead of epochs helps. Make sure to also set val_check_interval", - " under the trainer dict in your learning config JSON." + " under the trainer dict in your learning config JSON.", + "", + "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ", + "representative of this!" ], "net": { "name": "CatLSTM", diff --git a/bin/train/inputs/models/convnet.json b/bin/train/inputs/models/convnet.json @@ -1,4 +1,8 @@ { + "_notes": [ + "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ", + "representative of this!" + ], "net": { "name": "ConvNet", "config": { diff --git a/bin/train/inputs/models/demonet.json b/bin/train/inputs/models/demonet.json @@ -1,4 +1,8 @@ { + "_notes": [ + "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ", + "representative of this!" + ], "net": { "name": "WaveNet", "config": { diff --git a/bin/train/inputs/models/lstm.json b/bin/train/inputs/models/lstm.json @@ -9,7 +9,10 @@ " 1e-4 after 1000 epochs. I've found LSTMs to work with a pretty aggressive", " learning rate that would be out of the question for other architectures.", " * Number of units between 8 and 96, layers from 1 to 5 all seem to be ok", - " depending on the dataset, though bigger models might not make real-time." + " depending on the dataset, though bigger models might not make real-time.", + "", + "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ", + "representative of this!" ], "net": { "name": "LSTM", diff --git a/bin/train/inputs/models/wavenet.json b/bin/train/inputs/models/wavenet.json @@ -1,4 +1,8 @@ { + "_notes": [ + "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ", + "representative of this!" + ], "net": { "name": "WaveNet", "config": { diff --git a/bin/train/main.py b/bin/train/main.py @@ -2,6 +2,7 @@ # Created Date: Saturday February 5th 2022 # Author: Steven Atkinson (steven@atkinson.mn) + # Hack to recover graceful shutdowns in Windows. # This has to happen ASAP # See: @@ -201,6 +202,7 @@ def main_inner( trainer.checkpoint_callback.best_model_path, **Model.parse_config(model_config), ) + model.cpu() model.eval() if make_plots: plot( diff --git a/tests/test_bin/__init__.py b/tests/test_bin/__init__.py diff --git a/tests/test_bin/test_train/__init__.py b/tests/test_bin/test_train/__init__.py diff --git a/tests/test_bin/test_train/test_main.py b/tests/test_bin/test_train/test_main.py @@ -0,0 +1,203 @@ +# File: test_main.py +# Created Date: Sunday April 30th 2023 +# Author: Steven Atkinson (steven@atkinson.mn) + +import json +from enum import Enum +from pathlib import Path +from subprocess import check_call +from tempfile import TemporaryDirectory +from typing import Dict, Tuple, Union + +import numpy as np +import pytest +import torch + +from nam.data import np_to_wav + +_BIN_TRAIN_MAIN_PY_PATH = Path(__file__).absolute().parent.parent.parent.parent / Path( + "bin", "train", "main.py" +) + + +class _Device(Enum): + CPU = "cpu" + GPU = "gpu" + MPS = "mps" + + +class Test(object): + @classmethod + def setup_class(cls): + cls._num_samples = 128 + cls._num_samples_validation = 15 + cls._ny = 2 + cls._batch_size = 2 + + def test_cpu(self): + self._t_main(_Device.CPU) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU test") + def test_gpu(self): + self._t_main(_Device.GPU) + + @pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS test") + def test_mps(self): + self._t_main(_Device.MPS) + + @classmethod + def _data_config_path(cls, root_path: Path) -> Path: + return Path(cls._input_path(root_path), "data_config.json") + + def _get_configs( + self, root_path: Path, device: _Device + ) -> Tuple[Dict, Dict, Dict]: # TODO pydantic models + data_config = { + "train": { + "start": None, + "stop": -self._num_samples_validation, + "ny": self._ny, + }, + "validation": { + "start": -self._num_samples_validation, + "stop": None, + "ny": None, + }, + "common": { + "x_path": str(self._x_path(root_path)), + "y_path": str(self._y_path(root_path)), + "delay": 0, + }, + } + stage_channels = (3, 2) + model_config = { + "net": { + "name": "WaveNet", + "config": { + "layers_configs": [ + { + "condition_size": 1, + "input_size": 1, + "channels": stage_channels[0], + "head_size": stage_channels[1], + "kernel_size": 3, + "dilations": [1], + "activation": "Tanh", + "gated": False, + "head_bias": False, + }, + { + "condition_size": 1, + "input_size": stage_channels[0], + "channels": stage_channels[1], + "head_size": 1, + "kernel_size": 3, + "dilations": [2], + "activation": "Tanh", + "gated": False, + "head_bias": False, + }, + ], + "head_scale": 0.02, + }, + }, + "optimizer": {"lr": 0.004}, + "lr_scheduler": {"class": "ExponentialLR", "kwargs": {"gamma": 0.993}}, + } + + def extra_trainer_kwargs(device) -> Dict[str, Union[int, str]]: + return { + _Device.GPU: {"accelerator": "gpu", "devices": 1}, + _Device.MPS: {"accelerator": "mps", "devices": 1}, + }.get(device, {}) + + learning_config = { + "train_dataloader": { + "batch_size": 3, + "shuffle": True, + "pin_memory": True, + "drop_last": True, + "num_workers": 0, + }, + "val_dataloader": {}, + "trainer": {"max_epochs": 2, **extra_trainer_kwargs(device)}, + "trainer_fit_kwargs": {}, + } + + return data_config, model_config, learning_config + + def _get_data(self) -> Tuple[np.ndarray, np.ndarray]: + """ + :return: (N,), (N,) + """ + x = np.random.rand(self._num_samples) - 0.5 + y = 1.1 * x + return x, y + + @classmethod + def _input_path(cls, root_path: Path, ensure: bool = False) -> Path: + p = Path(root_path, "inputs") + if ensure: + p.mkdir() + return p + + @classmethod + def _learning_config_path(cls, root_path: Path) -> Path: + return Path(cls._input_path(root_path), "learning_config.json") + + @classmethod + def _model_config_path(cls, root_path: Path) -> Path: + return Path(cls._input_path(root_path), "model_config.json") + + @classmethod + def _output_path(cls, root_path: Path, ensure: bool = False) -> Path: + p = Path(root_path, "outputs") + if ensure: + p.mkdir() + return p + + def _setup_files(self, root_path: Path, device: _Device): + x, y = self._get_data() + np_to_wav(x, self._x_path(root_path)) + np_to_wav(y, self._y_path(root_path)) + data_config, model_config, learning_config = self._get_configs( + root_path, device + ) + with open(self._data_config_path(root_path), "w") as fp: + json.dump(data_config, fp) + with open(self._model_config_path(root_path), "w") as fp: + json.dump(model_config, fp) + with open(self._learning_config_path(root_path), "w") as fp: + json.dump(learning_config, fp) + + def _t_main(self, device: _Device): + """ + End-to-end test of bin/train/main.py + """ + with TemporaryDirectory() as tempdir: + tempdir = Path(tempdir) + self._input_path(tempdir, ensure=True) + self._setup_files(tempdir, device) + check_call( + [ + "python", + str(_BIN_TRAIN_MAIN_PY_PATH), + str(self._data_config_path(tempdir)), + str(self._model_config_path(tempdir)), + str(self._learning_config_path(tempdir)), + str(self._output_path(tempdir, ensure=True)), + "--no-show", + ] + ) + + @classmethod + def _x_path(cls, root_path: Path) -> Path: + return Path(cls._input_path(root_path), "input.wav") + + @classmethod + def _y_path(cls, root_path: Path) -> Path: + return Path(cls._input_path(root_path), "output.wav") + + +if __name__ == "__main__": + pytest.main()