[FEATURE] Use Depthwise convolutions

[MainRo]

Is your feature request related to a problem? Please describe.
CPU consumption versus modeling performances will always be a tricky balance for NAM. Some people find that the wavenet models consume too much resources.

Describe the solution you'd like
A potential low hanging fruit on wavenet and convnet models is to replace some of the convolutions with depthwise convolutions. These reduce significantly the number of parameters of each convolution and the number of operations.

Describe alternatives you've considered
None.

Additional context

I Made some early tests on two of my captures. One on a clean amp, and another one on high-gain settings.

The change in the training code is very small:

diff --git a/nam/models/wavenet.py b/nam/models/wavenet.py
index b48e41e..cd03b79 100644
--- a/nam/models/wavenet.py
+++ b/nam/models/wavenet.py
@@ -63,7 +63,7 @@ class _Layer(nn.Module):
         super().__init__()
         # Input mixer takes care of the bias
         mid_channels = 2 * channels if gated else channels
-        self._conv = Conv1d(channels, mid_channels, kernel_size, dilation=dilation)
+        self._conv = Conv1d(channels, mid_channels, kernel_size, dilation=dilation, groups=channels)
         # Custom init: favors direct input-output
         # self._conv.weight.data.zero_()
         self._input_mixer = Conv1d(condition_size, mid_channels, 1, bias=False)

On the base wavenet architeture, the number of parameters drops from 13.8k to 4.9k. After 1000 training epochs on my settings I have the following results:

• clean, base wavenet. ESR: 0.0007998017827048898
• clean, depthwise wavenet. ESR: 0.000936342345084995
• high-gain, base wavenet. ESR: 0.002280422719195485
• high-gain, depthwise wavenet. ESR: 0.003808232955634594

As you can see, the drop in performance is quite low. However, I could not test/listen to the resulting models because I did not update the core C++ operations. If there is interest in this change, I may contribute.

[sdatkinson]

Yep, this is a nice standard thing to do with convolutions.

I want to think out loud through this. Putting this in here and NeuralAmpModelerCore will enable models to be made that can't be played by NAM products without an update. Let's prioritize getting it in the core first and releasing. This will give downstream folks the ability to react and release updates (I'll figure out how long seems reasonable).

We can implement it here too, but I'm going to delay putting it in a release until the above has a chance to diffuse out.

I don't think I like groups=channels. Instead, providing a collated groups=[...] sounds nice and gives more flexibility. Yes, it's more complicated for the downstream implementation, but I think it's worth it to do it once right instead of e.g. adding a depthwise: bool=False etc and then removing it in a later version in favor of the "full" functionality, or having two ways to do one thing.

So your depthwise model would be config'd by modifying nam_full_configs/models/wavenet.json to this:

{
    "_notes": [
        "Dev note: Ensure that tests/test_bin/test_train/test_main.py's data is ",
        "representative of this!"
    ],
    "net": {
        "name": "WaveNet",
        "config": {
            "layers_configs": [
                {
                    "condition_size": 1,
                    "input_size": 1,
                    "channels": 16,
                    "head_size": 8,
                    "kernel_size": 3,
                    "dilations": [
                        1,
                        2,
                        4,
                        8,
                        16,
                        32,
                        64,
                        128,
                        256,
                        512
                    ],
                    "groups": [
                        16,
                        16,
                        16,
                        16,
                        16,
                        16,
                        16,
                        16,
                        16,
                        16
                    ],
                    "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
                    ],
                    "groups": [
                        8,
                        8,
                        8,
                        8,
                        8,
                        8,
                        8,
                        8,
                        8,
                        8
                    ],
                    "activation": "Tanh",
                    "gated": false,
                    "head_bias": true
                }
            ],
            "head_scale": 0.02
        }
    },
    "optimizer": {
        "lr": 0.004
    },
    "lr_scheduler": {
        "class": "ExponentialLR",
        "kwargs": {
            "gamma": 0.993
        }
    }
}

The reason that I didn't just do a single hyperparameter like "groups": 16 is because channels need to stay constant within a layer array, but groups don't. (Not sure if it's useful though).

I also realize that we could accept either an array or a scalar--groups: Union[int, List[int]]...not sure about that yet. It's simple for the eyes but offends my "one way to do it" sense.

---

If you'd like to contribute this (tests too, please!) then I would take it--but core first/"patience, please".