neural-amp-modeler

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

commit 6ecb5e86970c22c3a69325c8fc01b9932707d32c
parent e0c73b26277936e3da8f42bb28114a181a8b4bd9
Author: Steven Atkinson <steven@atkinson.mn>
Date:   Tue, 31 Jan 2023 19:24:07 -0800

Add different pre-made architectures to easy mode (#81)


Diffstat:
Mbin/train/easy_colab.ipynb | 5++++-
Mnam/train/colab.py | 137+++++++++++++++++++++++++++++++++++++++++++++++++++++++++----------------------
2 files changed, 103 insertions(+), 39 deletions(-)

diff --git a/bin/train/easy_colab.ipynb b/bin/train/easy_colab.ipynb @@ -98,7 +98,10 @@ "outputs": [], "source": [ "%tensorboard --logdir /content/lightning_logs\n", - "run(epochs=100)\n", + "run(\n", + " epochs=100,\n", + " architecture=\"standard\" # standard, lite, feather\n", + ")\n", "# Psst! Curious how it's going?\n", "# You can look under lightning_logs/version_0/checkpoints to see it saving its progress.\n", "# Look for the number after \"ESR\" to go down. It will start around 1.0. 0.1 is ok, and \n", diff --git a/nam/train/colab.py b/nam/train/colab.py @@ -6,6 +6,7 @@ Hide the mess in Colab to make things look pretty for users. """ +from enum import Enum from pathlib import Path from time import time from typing import Optional, Tuple @@ -20,6 +21,12 @@ from nam.data import REQUIRED_RATE, Split, init_dataset, wav_to_np from nam.models import Model +class _Architecture(Enum): + STANDARD = "standard" + LITE = "lite" + FEATHER = "feather" + + class _Version: def __init__(self, major: int, minor: int, patch: int): self.major = major @@ -133,13 +140,97 @@ def _calibrate_delay( return delay +def _get_wavenet_config(architecture): + return { + _Architecture.STANDARD: { + "layers_configs": [ + { + "input_size": 1, + "condition_size": 1, + "channels": 16, + "head_size": 8, + "kernel_size": 3, + "dilations": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512], + "activation": "Tanh", + "gated": False, + "head_bias": False, + }, + { + "condition_size": 1, + "input_size": 16, + "channels": 8, + "head_size": 1, + "kernel_size": 3, + "dilations": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512], + "activation": "Tanh", + "gated": False, + "head_bias": True, + }, + ], + "head_scale": 0.02, + }, + _Architecture.LITE: { + "layers_configs": [ + { + "input_size": 1, + "condition_size": 1, + "channels": 12, + "head_size": 6, + "kernel_size": 3, + "dilations": [1, 2, 4, 8, 16, 32, 64], + "activation": "Tanh", + "gated": False, + "head_bias": False, + }, + { + "condition_size": 1, + "input_size": 12, + "channels": 6, + "head_size": 1, + "kernel_size": 3, + "dilations": [128, 256, 512, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512], + "activation": "Tanh", + "gated": False, + "head_bias": True, + }, + ], + "head_scale": 0.02, + }, + _Architecture.FEATHER: { + "layers_configs": [ + { + "input_size": 1, + "condition_size": 1, + "channels": 8, + "head_size": 4, + "kernel_size": 3, + "dilations": [1, 2, 4, 8, 16, 32, 64], + "activation": "Tanh", + "gated": False, + "head_bias": False, + }, + { + "condition_size": 1, + "input_size": 8, + "channels": 4, + "head_size": 1, + "kernel_size": 3, + "dilations": [128, 256, 512, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512], + "activation": "Tanh", + "gated": False, + "head_bias": True, + }, + ], + "head_scale": 0.02, + }, + }[architecture] + + def _get_configs( input_basename: str, delay: int, epochs: int, - stage_1_channels: int, - stage_2_channels: int, - head_scale: float, + architecture: _Architecture, lr: float, lr_decay: float, ): @@ -159,33 +250,7 @@ def _get_configs( "name": "WaveNet", # This should do decently. If you really want a nice model, try turning up # "channels" in the first block and "input_size" in the second from 12 to 16. - "config": { - "layers_configs": [ - { - "input_size": 1, - "condition_size": 1, - "head_size": stage_2_channels, - "channels": stage_1_channels, - "kernel_size": 3, - "dilations": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512], - "activation": "Tanh", - "gated": False, - "head_bias": False, - }, - { - "input_size": stage_1_channels, - "condition_size": 1, - "head_size": 1, - "channels": stage_2_channels, - "kernel_size": 3, - "dilations": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512], - "activation": "Tanh", - "gated": False, - "head_bias": True, - }, - ], - "head_scale": head_scale, - }, + "config": _get_wavenet_config(architecture), }, "loss": {"val_loss": "esr"}, "optimizer": {"lr": lr}, @@ -258,11 +323,9 @@ def _get_valid_export_directory(): def run( - epochs=100, - delay=None, - stage_1_channels=16, - stage_2_channels=8, - head_scale: float = 0.02, + epochs: int = 100, + delay: Optional[int] = None, + architecture: str = "standard", lr=0.004, lr_decay=0.007, seed=0, @@ -284,9 +347,7 @@ def run( input_basename, delay, epochs, - stage_1_channels, - stage_2_channels, - head_scale, + _Architecture(architecture), lr, lr_decay, )