Head 1x1 convolution (#189)

* Implemented head1x1 operation, API to config and activate it still not implemented

* Merged with bottleneck updates

* Fixed some shape issues, added tests, updated all tests to pass head1x1_params when instantiating Wavenet and LayerArray

* Updated head1x1 test to make sure it works when number of output channels is different from number of bottleneck channels

* Uncommented and simplified gated tests now that everything else is working.

* Fixing issue post merge

* Updated CMakeLists.txt to solve issue with Eigen

* Updated init issue in wavenet.cpp, relaxed test boundary since test was passing on ARM but failing on x86

* Addressing comments

---------

Co-authored-by: João Felipe Santos <santosjf@pm.me>

Author: jfsantos

CMakeLists.txt

@@ -25,6 +25,7 @@ else()
 endif()
 
 set(NAM_DEPS_PATH "${CMAKE_CURRENT_SOURCE_DIR}/Dependencies")
+include_directories(SYSTEM "${NAM_DEPS_PATH}/eigen")
 
 add_subdirectory(tools)
 

NAM/wavenet.cpp

@@ -19,15 +19,29 @@ void nam::wavenet::_Layer::SetMaxBufferSize(const int maxBufferSize)
   // Pre-allocate output buffers
   const long channels = this->get_channels();
   this->_output_next_layer.resize(channels, maxBufferSize);
-  // _output_head stores the activated portion: bottleneck rows (the actual bottleneck value, not doubled)
-  this->_output_head.resize(this->_bottleneck, maxBufferSize);
+  // _output_head stores the activated portion: bottleneck rows when no head1x1, or head1x1 out_channels when head1x1 is active
+  if (_head1x1)
+  {
+    this->_output_head.resize(_head1x1->get_out_channels(), maxBufferSize);
+    this->_output_head.setZero(); // Ensure consistent initialization across platforms
+    _head1x1->SetMaxBufferSize(maxBufferSize);
+  }
+  else
+  {
+    this->_output_head.resize(this->_bottleneck, maxBufferSize);
+    this->_output_head.setZero(); // Ensure consistent initialization across platforms
+  }
 }
 
 void nam::wavenet::_Layer::set_weights_(std::vector<float>::iterator& weights)
 {
   this->_conv.set_weights_(weights);
   this->_input_mixin.set_weights_(weights);
   this->_1x1.set_weights_(weights);
+  if (this->_head1x1)
+  {
+    this->_head1x1->set_weights_(weights);
+  }
 }
 
 void nam::wavenet::_Layer::Process(const Eigen::MatrixXf& input, const Eigen::MatrixXf& condition, const int num_frames)
@@ -61,31 +75,39 @@ void nam::wavenet::_Layer::Process(const Eigen::MatrixXf& input, const Eigen::Ma
     _1x1.process_(_z.topRows(bottleneck), num_frames); // Might not be RT safe
   }
 
-  // Store output to head (skip connection: activated conv output)
-  if (!this->_gated)
-    this->_output_head.leftCols(num_frames).noalias() = this->_z.leftCols(num_frames);
-  else
-    this->_output_head.leftCols(num_frames).noalias() = this->_z.topRows(bottleneck).leftCols(num_frames);
+  if (this->_head1x1) {
+    if (!this->_gated)
+      this->_head1x1->process_(this->_z.leftCols(num_frames), num_frames);
+    else
+      this->_head1x1->process(this->_z.topRows(bottleneck).leftCols(num_frames), num_frames);
+    this->_output_head.leftCols(num_frames).noalias() = this->_head1x1->GetOutput().leftCols(num_frames);
+  } else {
+    // Store output to head (skip connection: activated conv output)
+    if (!this->_gated)
+      this->_output_head.leftCols(num_frames).noalias() = this->_z.leftCols(num_frames);
+    else
+      this->_output_head.leftCols(num_frames).noalias() = this->_z.topRows(bottleneck).leftCols(num_frames);
+  }
+
   // Store output to next layer (residual connection: input + _1x1 output)
   this->_output_next_layer.leftCols(num_frames).noalias() =
     input.leftCols(num_frames) + _1x1.GetOutput().leftCols(num_frames);
 }
 
-
 // LayerArray =================================================================
 
 nam::wavenet::_LayerArray::_LayerArray(const int input_size, const int condition_size, const int head_size,
                                        const int channels, const int bottleneck, const int kernel_size,
                                        const std::vector<int>& dilations, const std::string activation,
                                        const bool gated, const bool head_bias, const int groups_input,
-                                       const int groups_1x1)
+                                       const int groups_1x1, const Head1x1Params& head1x1_params)
 : _rechannel(input_size, channels, false)
 , _head_rechannel(bottleneck, head_size, head_bias)
 , _bottleneck(bottleneck)
 {
   for (size_t i = 0; i < dilations.size(); i++)
     this->_layers.push_back(_Layer(
-      condition_size, channels, bottleneck, kernel_size, dilations[i], activation, gated, groups_input, groups_1x1));
+      condition_size, channels, bottleneck, kernel_size, dilations[i], activation, gated, groups_input, groups_1x1, head1x1_params));
 }
 
 void nam::wavenet::_LayerArray::SetMaxBufferSize(const int maxBufferSize)
@@ -212,7 +234,8 @@ nam::wavenet::WaveNet::WaveNet(const int in_channels,
       layer_array_params[i].input_size, layer_array_params[i].condition_size, layer_array_params[i].head_size,
       layer_array_params[i].channels, layer_array_params[i].bottleneck, layer_array_params[i].kernel_size,
       layer_array_params[i].dilations, layer_array_params[i].activation, layer_array_params[i].gated,
-      layer_array_params[i].head_bias, layer_array_params[i].groups_input, layer_array_params[i].groups_1x1));
+      layer_array_params[i].head_bias, layer_array_params[i].groups_input, layer_array_params[i].groups_1x1,
+      layer_array_params[i].head1x1_params));
     if (i > 0)
       if (layer_array_params[i].channels != layer_array_params[i - 1].head_size)
       {
@@ -324,10 +347,17 @@ std::unique_ptr<nam::DSP> nam::wavenet::Factory(const nlohmann::json& config, st
     const int groups_1x1 = layer_config.value("groups_1x1", 1); // defaults to 1
     const int channels = layer_config["channels"];
     const int bottleneck = layer_config.value("bottleneck", channels); // defaults to channels if not present
+    
+    // Parse head1x1 parameters
+    bool head1x1_active = layer_config.value("head1x1_active", false);
+    int head1x1_out_channels = layer_config.value("head1x1_out_channels", channels);
+    int head1x1_groups = layer_config.value("head1x1_groups", 1);
+    nam::wavenet::Head1x1Params head1x1_params(head1x1_active, head1x1_out_channels, head1x1_groups);
+    
     layer_array_params.push_back(nam::wavenet::LayerArrayParams(
       layer_config["input_size"], layer_config["condition_size"], layer_config["head_size"], channels, bottleneck,
       layer_config["kernel_size"], layer_config["dilations"], layer_config["activation"], layer_config["gated"],
-      layer_config["head_bias"], groups, groups_1x1));
+      layer_config["head_bias"], groups, groups_1x1, head1x1_params));
   }
   const bool with_head = !config["head"].is_null();
   const float head_scale = config["head_scale"];

NAM/wavenet.h

@@ -2,6 +2,7 @@
 
 #include <string>
 #include <vector>
+#include <memory>
 
 #include "json.hpp"
 #include <Eigen/Dense>
@@ -13,17 +14,36 @@ namespace nam
 {
 namespace wavenet
 {
+// Parameters for head1x1 configuration
+struct Head1x1Params
+{
+  Head1x1Params(bool active_, int out_channels_, int groups_)
+    : active(active_), out_channels(out_channels_), groups(groups_) {}
+
+  const bool active;
+  const int out_channels;
+  const int groups;
+};
+
 class _Layer
 {
 public:
   _Layer(const int condition_size, const int channels, const int bottleneck, const int kernel_size, const int dilation,
-         const std::string activation, const bool gated, const int groups_input, const int groups_1x1)
-  : _conv(channels, gated ? 2 * bottleneck : bottleneck, kernel_size, true, dilation, groups_input)
+         const std::string activation, const bool gated, const int groups_input, const int groups_1x1,
+         const Head1x1Params& head1x1_params)
+  : _conv(channels, gated ? 2 * bottleneck : bottleneck, kernel_size, true, dilation)
   , _input_mixin(condition_size, gated ? 2 * bottleneck : bottleneck, false)
-  , _1x1(bottleneck, channels, true, groups_1x1)
+  , _1x1(bottleneck, channels, groups_1x1)
   , _activation(activations::Activation::get_activation(activation)) // needs to support activations with parameters
-  , _gated(gated)
-  , _bottleneck(bottleneck) {};
+  , _gated(gated) 
+  , _bottleneck(bottleneck)
+  {
+    if (head1x1_params.active)
+    {
+      _head1x1 = std::make_unique<Conv1x1>(bottleneck, head1x1_params.out_channels, true, head1x1_params.groups);
+    }
+  };
+ 
   // Resize all arrays to be able to process `maxBufferSize` frames.
   void SetMaxBufferSize(const int maxBufferSize);
   // Set the parameters of this module
@@ -63,6 +83,8 @@ class _Layer
   Conv1x1 _input_mixin;
   // The post-activation 1x1 convolution
   Conv1x1 _1x1;
+  // The post-activation 1x1 convolution outputting to the head, optional
+  std::unique_ptr<Conv1x1> _head1x1;
   // The internal state
   Eigen::MatrixXf _z;
   // Output to next layer (residual connection: input + _1x1 output)
@@ -81,7 +103,7 @@ class LayerArrayParams
   LayerArrayParams(const int input_size_, const int condition_size_, const int head_size_, const int channels_,
                    const int bottleneck_, const int kernel_size_, const std::vector<int>&& dilations_,
                    const std::string activation_, const bool gated_, const bool head_bias_, const int groups_input,
-                   const int groups_1x1_)
+                   const int groups_1x1_, const Head1x1Params& head1x1_params_)
   : input_size(input_size_)
   , condition_size(condition_size_)
   , head_size(head_size_)
@@ -94,6 +116,7 @@ class LayerArrayParams
   , head_bias(head_bias_)
   , groups_input(groups_input)
   , groups_1x1(groups_1x1_)
+  , head1x1_params(head1x1_params_)
   {
   }
 
@@ -109,6 +132,7 @@ class LayerArrayParams
   const bool head_bias;
   const int groups_input;
   const int groups_1x1;
+  const Head1x1Params head1x1_params;
 };
 
 // An array of layers with the same channels, kernel sizes, activations.
@@ -118,7 +142,7 @@ class _LayerArray
   _LayerArray(const int input_size, const int condition_size, const int head_size, const int channels,
               const int bottleneck, const int kernel_size, const std::vector<int>& dilations,
               const std::string activation, const bool gated, const bool head_bias, const int groups_input,
-              const int groups_1x1);
+              const int groups_1x1, const Head1x1Params& head1x1_params);
 
   void SetMaxBufferSize(const int maxBufferSize);
 

tools/run_tests.cpp

@@ -14,6 +14,7 @@
 #include "test/test_wavenet/test_layer_array.cpp"
 #include "test/test_wavenet/test_full.cpp"
 #include "test/test_wavenet/test_real_time_safe.cpp"
+#include "test/test_wavenet/test_head1x1.cpp"
 #include "test/test_gating_activations.cpp"
 #include "test/test_wavenet_gating_compatibility.cpp"
 #include "test/test_blending_detailed.cpp"
@@ -115,6 +116,11 @@ int main()
   test_wavenet::test_full::test_wavenet_zero_input();
   test_wavenet::test_full::test_wavenet_different_buffer_sizes();
   test_wavenet::test_full::test_wavenet_prewarm();
+  test_wavenet::test_head1x1::test_head1x1_inactive();
+  test_wavenet::test_head1x1::test_head1x1_active();
+  test_wavenet::test_head1x1::test_head1x1_gated();
+  test_wavenet::test_head1x1::test_head1x1_groups();
+  test_wavenet::test_head1x1::test_head1x1_different_out_channels();
   test_wavenet::test_allocation_tracking_pass();
   test_wavenet::test_allocation_tracking_fail();
   test_wavenet::test_conv1d_process_realtime_safe();

tools/test/test_wavenet/test_full.cpp

@@ -29,9 +29,11 @@ void test_wavenet_model()
   const bool with_head = false;
   const int groups = 1;
   const int groups_1x1 = 1;
+  const bool head1x1_active = false;
 
+  nam::wavenet::Head1x1Params head1x1_params(head1x1_active, channels, 1);
   nam::wavenet::LayerArrayParams params(input_size, condition_size, head_size, channels, bottleneck, kernel_size,
-                                        std::move(dilations), activation, gated, head_bias, groups, groups_1x1);
+                                        std::move(dilations), activation, gated, head_bias, groups, groups_1x1, head1x1_params);
   std::vector<nam::wavenet::LayerArrayParams> layer_array_params;
   layer_array_params.push_back(std::move(params));
 
@@ -91,14 +93,17 @@ void test_wavenet_multiple_arrays()
   std::vector<int> dilations1{1};
   const int bottleneck = channels;
   const int groups_1x1 = 1;
+  const bool head1x1_active = false;
+
+  nam::wavenet::Head1x1Params head1x1_params(head1x1_active, channels, 1);
   layer_array_params.push_back(nam::wavenet::LayerArrayParams(input_size, condition_size, head_size, channels,
-                                                              bottleneck, kernel_size, std::move(dilations1),
-                                                              activation, gated, head_bias, groups, groups_1x1));
+                                                              bottleneck, kernel_size, std::move(dilations1), activation,
+                                                              gated, head_bias, groups, groups_1x1, head1x1_params));
   // Second array (head_size of first must match channels of second)
   std::vector<int> dilations2{1};
   layer_array_params.push_back(nam::wavenet::LayerArrayParams(head_size, condition_size, head_size, channels,
-                                                              bottleneck, kernel_size, std::move(dilations2),
-                                                              activation, gated, head_bias, groups, groups_1x1));
+                                                              bottleneck, kernel_size, std::move(dilations2), activation,
+                                                              gated, head_bias, groups, groups_1x1, head1x1_params));
 
   std::vector<float> weights;
   // Array 0: rechannel, layer, head_rechannel
@@ -145,9 +150,11 @@ void test_wavenet_zero_input()
   const bool with_head = false;
   const int groups = 1;
   const int groups_1x1 = 1;
+  const bool head1x1_active = false;
+  nam::wavenet::Head1x1Params head1x1_params(head1x1_active, channels, 1);
 
   nam::wavenet::LayerArrayParams params(input_size, condition_size, head_size, channels, bottleneck, kernel_size,
-                                        std::move(dilations), activation, gated, head_bias, groups, groups_1x1);
+                                        std::move(dilations), activation, gated, head_bias, groups, groups_1x1, head1x1_params);
   std::vector<nam::wavenet::LayerArrayParams> layer_array_params;
   layer_array_params.push_back(std::move(params));
 
@@ -190,9 +197,11 @@ void test_wavenet_different_buffer_sizes()
   const bool with_head = false;
   const int groups = 1;
   const int groups_1x1 = 1;
+  const bool head1x1_active = false;
+  nam::wavenet::Head1x1Params head1x1_params(head1x1_active, channels, 1);
 
   nam::wavenet::LayerArrayParams params(input_size, condition_size, head_size, channels, bottleneck, kernel_size,
-                                        std::move(dilations), activation, gated, head_bias, groups, groups_1x1);
+                                        std::move(dilations), activation, gated, head_bias, groups, groups_1x1, head1x1_params);
   std::vector<nam::wavenet::LayerArrayParams> layer_array_params;
   layer_array_params.push_back(std::move(params));
 
@@ -238,9 +247,12 @@ void test_wavenet_prewarm()
   const bool with_head = false;
   const int groups = 1;
   const int groups_1x1 = 1;
+  const bool head1x1_active = false;
+
+  nam::wavenet::Head1x1Params head1x1_params(head1x1_active, channels, 1);
 
   nam::wavenet::LayerArrayParams params(input_size, condition_size, head_size, channels, bottleneck, kernel_size,
-                                        std::move(dilations), activation, gated, head_bias, groups, groups_1x1);
+                                        std::move(dilations), activation, gated, head_bias, groups, groups_1x1, head1x1_params);
   std::vector<nam::wavenet::LayerArrayParams> layer_array_params;
   layer_array_params.push_back(std::move(params));