Configurable activations

NAM/activations.cpp

@@ -1,36 +1,174 @@
 #include "activations.h"
 
-nam::activations::ActivationTanh _TANH = nam::activations::ActivationTanh();
-nam::activations::ActivationFastTanh _FAST_TANH = nam::activations::ActivationFastTanh();
-nam::activations::ActivationHardTanh _HARD_TANH = nam::activations::ActivationHardTanh();
-nam::activations::ActivationReLU _RELU = nam::activations::ActivationReLU();
-nam::activations::ActivationLeakyReLU _LEAKY_RELU =
-  nam::activations::ActivationLeakyReLU(0.01); // FIXME does not parameterize LeakyReLU
-nam::activations::ActivationPReLU _PRELU = nam::activations::ActivationPReLU(0.01); // Same as leaky ReLU by default
-nam::activations::ActivationSigmoid _SIGMOID = nam::activations::ActivationSigmoid();
-nam::activations::ActivationSwish _SWISH = nam::activations::ActivationSwish();
-nam::activations::ActivationHardSwish _HARD_SWISH = nam::activations::ActivationHardSwish();
-nam::activations::ActivationLeakyHardTanh _LEAKY_HARD_TANH = nam::activations::ActivationLeakyHardTanh();
+// Global singleton instances (statically allocated, never deleted)
+static nam::activations::ActivationTanh _TANH;
+static nam::activations::ActivationFastTanh _FAST_TANH;
+static nam::activations::ActivationHardTanh _HARD_TANH;
+static nam::activations::ActivationReLU _RELU;
+static nam::activations::ActivationLeakyReLU _LEAKY_RELU(0.01); // FIXME does not parameterize LeakyReLU
+static nam::activations::ActivationPReLU _PRELU(0.01); // Same as leaky ReLU by default
+static nam::activations::ActivationSigmoid _SIGMOID;
+static nam::activations::ActivationSwish _SWISH;
+static nam::activations::ActivationHardSwish _HARD_SWISH;
+static nam::activations::ActivationLeakyHardTanh _LEAKY_HARD_TANH;
 
 bool nam::activations::Activation::using_fast_tanh = false;
 
-std::unordered_map<std::string, nam::activations::Activation*> nam::activations::Activation::_activations = {
-  {"Tanh", &_TANH},  {"Hardtanh", &_HARD_TANH},   {"Fasttanh", &_FAST_TANH},
-  {"ReLU", &_RELU},  {"LeakyReLU", &_LEAKY_RELU}, {"Sigmoid", &_SIGMOID},
-  {"SiLU", &_SWISH}, {"Hardswish", &_HARD_SWISH}, {"LeakyHardtanh", &_LEAKY_HARD_TANH},
-  {"PReLU", &_PRELU}};
+// Helper to create a non-owning shared_ptr (no-op deleter) for singletons
+template<typename T>
+nam::activations::Activation::Ptr make_singleton_ptr(T& singleton)
+{
+  return nam::activations::Activation::Ptr(&singleton, [](nam::activations::Activation*){});
+}
+
+std::unordered_map<std::string, nam::activations::Activation::Ptr> nam::activations::Activation::_activations = {
+  {"Tanh", make_singleton_ptr(_TANH)},
+  {"Hardtanh", make_singleton_ptr(_HARD_TANH)},
+  {"Fasttanh", make_singleton_ptr(_FAST_TANH)},
+  {"ReLU", make_singleton_ptr(_RELU)},
+  {"LeakyReLU", make_singleton_ptr(_LEAKY_RELU)},
+  {"Sigmoid", make_singleton_ptr(_SIGMOID)},
+  {"SiLU", make_singleton_ptr(_SWISH)},
+  {"Hardswish", make_singleton_ptr(_HARD_SWISH)},
+  {"LeakyHardtanh", make_singleton_ptr(_LEAKY_HARD_TANH)},
+  {"PReLU", make_singleton_ptr(_PRELU)}
+};
 
-nam::activations::Activation* tanh_bak = nullptr;
-nam::activations::Activation* sigmoid_bak = nullptr;
+nam::activations::Activation::Ptr tanh_bak = nullptr;
+nam::activations::Activation::Ptr sigmoid_bak = nullptr;
 
-nam::activations::Activation* nam::activations::Activation::get_activation(const std::string name)
+nam::activations::Activation::Ptr nam::activations::Activation::get_activation(const std::string name)
 {
   if (_activations.find(name) == _activations.end())
     return nullptr;
 
   return _activations[name];
 }
 
+// ActivationConfig implementation
+nam::activations::ActivationConfig nam::activations::ActivationConfig::simple(ActivationType t)
+{
+  ActivationConfig config;
+  config.type = t;
+  return config;
+}
+
+nam::activations::ActivationConfig nam::activations::ActivationConfig::from_json(const nlohmann::json& j)
+{
+  ActivationConfig config;
+
+  // Map from string to ActivationType
+  static const std::unordered_map<std::string, ActivationType> type_map = {
+    {"Tanh", ActivationType::Tanh},
+    {"Hardtanh", ActivationType::Hardtanh},
+    {"Fasttanh", ActivationType::Fasttanh},
+    {"ReLU", ActivationType::ReLU},
+    {"LeakyReLU", ActivationType::LeakyReLU},
+    {"PReLU", ActivationType::PReLU},
+    {"Sigmoid", ActivationType::Sigmoid},
+    {"SiLU", ActivationType::SiLU},
+    {"Hardswish", ActivationType::Hardswish},
+    {"LeakyHardtanh", ActivationType::LeakyHardtanh},
+    {"LeakyHardTanh", ActivationType::LeakyHardtanh} // Support both casings
+  };
+
+  // If it's a string, simple lookup
+  if (j.is_string())
+  {
+    std::string name = j.get<std::string>();
+    auto it = type_map.find(name);
+    if (it == type_map.end())
+    {
+      throw std::runtime_error("Unknown activation type: " + name);
+    }
+    config.type = it->second;
+    return config;
+  }
+
+  // If it's an object, parse type and parameters
+  if (j.is_object())
+  {
+    std::string type_str = j["type"].get<std::string>();
+    auto it = type_map.find(type_str);
+    if (it == type_map.end())
+    {
+      throw std::runtime_error("Unknown activation type: " + type_str);
+    }
+    config.type = it->second;
+
+    // Parse optional parameters based on activation type
+    if (config.type == ActivationType::PReLU)
+    {
+      if (j.find("negative_slope") != j.end())
+      {
+        config.negative_slope = j["negative_slope"].get<float>();
+      }
+      else if (j.find("negative_slopes") != j.end())
+      {
+        config.negative_slopes = j["negative_slopes"].get<std::vector<float>>();
+      }
+    }
+    else if (config.type == ActivationType::LeakyReLU)
+    {
+      config.negative_slope = j.value("negative_slope", 0.01f);
+    }
+    else if (config.type == ActivationType::LeakyHardtanh)
+    {
+      config.min_val = j.value("min_val", -1.0f);
+      config.max_val = j.value("max_val", 1.0f);
+      config.min_slope = j.value("min_slope", 0.01f);
+      config.max_slope = j.value("max_slope", 0.01f);
+    }
+
+    return config;
+  }
+
+  throw std::runtime_error("Invalid activation config: expected string or object");
+}
+
+nam::activations::Activation::Ptr nam::activations::Activation::get_activation(const ActivationConfig& config)
+{
+  switch (config.type)
+  {
+    case ActivationType::Tanh:
+      return _activations["Tanh"];
+    case ActivationType::Hardtanh:
+      return _activations["Hardtanh"];
+    case ActivationType::Fasttanh:
+      return _activations["Fasttanh"];
+    case ActivationType::ReLU:
+      return _activations["ReLU"];
+    case ActivationType::Sigmoid:
+      return _activations["Sigmoid"];
+    case ActivationType::SiLU:
+      return _activations["SiLU"];
+    case ActivationType::Hardswish:
+      return _activations["Hardswish"];
+    case ActivationType::LeakyReLU:
+      if (config.negative_slope.has_value())
+      {
+        return std::make_shared<ActivationLeakyReLU>(config.negative_slope.value());
+      }
+      return _activations["LeakyReLU"];
+    case ActivationType::PReLU:
+      if (config.negative_slopes.has_value())
+      {
+        return std::make_shared<ActivationPReLU>(config.negative_slopes.value());
+      }
+      else if (config.negative_slope.has_value())
+      {
+        return std::make_shared<ActivationPReLU>(config.negative_slope.value());
+      }
+      return std::make_shared<ActivationPReLU>(0.01f);
+    case ActivationType::LeakyHardtanh:
+      return std::make_shared<ActivationLeakyHardTanh>(
+        config.min_val.value_or(-1.0f), config.max_val.value_or(1.0f), config.min_slope.value_or(0.01f),
+        config.max_slope.value_or(0.01f));
+    default:
+      return nullptr;
+  }
+}
+
 void nam::activations::Activation::enable_fast_tanh()
 {
   nam::activations::Activation::using_fast_tanh = true;
@@ -69,8 +207,7 @@ void nam::activations::Activation::enable_lut(std::string function_name, float m
   {
     throw std::runtime_error("Tried to enable LUT for a function other than Tanh or Sigmoid");
   }
-  FastLUTActivation lut_activation(min, max, n_points, fn);
-  _activations[function_name] = &lut_activation;
+  _activations[function_name] = std::make_shared<FastLUTActivation>(min, max, n_points, fn);
 }
 
 void nam::activations::Activation::disable_lut(std::string function_name)

NAM/activations.h

@@ -1,16 +1,58 @@
 #pragma once
 
 #include <cassert>
-#include <string>
 #include <cmath> // expf
+#include <functional>
+#include <memory>
+#include <optional>
+#include <string>
 #include <unordered_map>
+#include <vector>
+
 #include <Eigen/Dense>
-#include <functional>
+
+#include "json.hpp"
 
 namespace nam
 {
 namespace activations
 {
+
+// Forward declaration
+class Activation;
+
+// Strongly-typed activation type enum
+enum class ActivationType
+{
+  Tanh,
+  Hardtanh,
+  Fasttanh,
+  ReLU,
+  LeakyReLU,
+  PReLU,
+  Sigmoid,
+  SiLU, // aka Swish
+  Hardswish,
+  LeakyHardtanh
+};
+
+// Strongly-typed activation configuration
+struct ActivationConfig
+{
+  ActivationType type;
+
+  // Optional parameters (used by specific activation types)
+  std::optional<float> negative_slope;               // LeakyReLU, PReLU (single)
+  std::optional<std::vector<float>> negative_slopes; // PReLU (per-channel)
+  std::optional<float> min_val;                      // LeakyHardtanh
+  std::optional<float> max_val;                      // LeakyHardtanh
+  std::optional<float> min_slope;                    // LeakyHardtanh
+  std::optional<float> max_slope;                    // LeakyHardtanh
+
+  // Convenience constructors
+  static ActivationConfig simple(ActivationType t);
+  static ActivationConfig from_json(const nlohmann::json& j);
+};
 inline float relu(float x)
 {
   return x > 0.0f ? x : 0.0f;
@@ -91,6 +133,9 @@ inline float hardswish(float x)
 class Activation
 {
 public:
+  // Type alias for shared pointer to Activation
+  using Ptr = std::shared_ptr<Activation>;
+
   Activation() = default;
   virtual ~Activation() = default;
   virtual void apply(Eigen::MatrixXf& matrix) { apply(matrix.data(), matrix.rows() * matrix.cols()); }
@@ -101,15 +146,17 @@ class Activation
   }
   virtual void apply(float* data, long size) {}
 
-  static Activation* get_activation(const std::string name);
+  static Ptr get_activation(const std::string name);
+  static Ptr get_activation(const ActivationConfig& config);
+  static Ptr get_activation(const nlohmann::json& activation_config);
   static void enable_fast_tanh();
   static void disable_fast_tanh();
   static bool using_fast_tanh;
   static void enable_lut(std::string function_name, float min, float max, std::size_t n_points);
   static void disable_lut(std::string function_name);
 
 protected:
-  static std::unordered_map<std::string, Activation*> _activations;
+  static std::unordered_map<std::string, Ptr> _activations;
 };
 
 // identity function activation
@@ -226,20 +273,21 @@ class ActivationPReLU : public Activation
   void apply(Eigen::MatrixXf& matrix) override
   {
     // Matrix is organized as (channels, time_steps)
-    int n_channels = negative_slopes.size();
-    int actual_channels = matrix.rows();
-
-    // NOTE: check not done during runtime on release builds
-    // model loader should make sure dimensions match
-    assert(actual_channels == n_channels);
-
+    unsigned long actual_channels = static_cast<unsigned long>(matrix.rows());
+
+    // Prepare the slopes for the current matrix size
+    std::vector<float> slopes_for_channels = negative_slopes;
+
+    // Fail loudly if input has more channels than activation
+    assert(actual_channels == negative_slopes.size());
+
     // Apply each negative slope to its corresponding channel
-    for (int channel = 0; channel < std::min(n_channels, actual_channels); channel++)
+    for (unsigned long channel = 0; channel < actual_channels; channel++)
     {
       // Apply the negative slope to all time steps in this channel
-      for (int time_step = 0; time_step < matrix.rows(); time_step++)
+      for (int time_step = 0; time_step < matrix.cols(); time_step++)
       {
-        matrix(channel, time_step) = leaky_relu(matrix(channel, time_step), negative_slopes[channel]);
+        matrix(channel, time_step) = leaky_relu(matrix(channel, time_step), slopes_for_channels[channel]);
       }
     }
   }

NAM/convnet.cpp

@@ -48,15 +48,16 @@ void nam::convnet::BatchNorm::process_(Eigen::MatrixXf& x, const long i_start, c
 }
 
 void nam::convnet::ConvNetBlock::set_weights_(const int in_channels, const int out_channels, const int _dilation,
-                                              const bool batchnorm, const std::string activation, const int groups,
+                                              const bool batchnorm,
+                                              const activations::ActivationConfig& activation_config, const int groups,
                                               std::vector<float>::iterator& weights)
 {
   this->_batchnorm = batchnorm;
   // HACK 2 kernel
   this->conv.set_size_and_weights_(in_channels, out_channels, 2, _dilation, !batchnorm, groups, weights);
   if (this->_batchnorm)
     this->batchnorm = BatchNorm(out_channels, weights);
-  this->activation = activations::Activation::get_activation(activation);
+  this->activation = activations::Activation::get_activation(activation_config);
 }
 
 void nam::convnet::ConvNetBlock::SetMaxBufferSize(const int maxBufferSize)
@@ -173,8 +174,9 @@ void nam::convnet::_Head::process_(const Eigen::MatrixXf& input, Eigen::MatrixXf
 }
 
 nam::convnet::ConvNet::ConvNet(const int in_channels, const int out_channels, const int channels,
-                               const std::vector<int>& dilations, const bool batchnorm, const std::string activation,
-                               std::vector<float>& weights, const double expected_sample_rate, const int groups)
+                               const std::vector<int>& dilations, const bool batchnorm,
+                               const activations::ActivationConfig& activation_config, std::vector<float>& weights,
+                               const double expected_sample_rate, const int groups)
 : Buffer(in_channels, out_channels, *std::max_element(dilations.begin(), dilations.end()), expected_sample_rate)
 {
   this->_verify_weights(channels, dilations, batchnorm, weights.size());
@@ -183,7 +185,7 @@ nam::convnet::ConvNet::ConvNet(const int in_channels, const int out_channels, co
   // First block takes in_channels input, subsequent blocks take channels input
   for (size_t i = 0; i < dilations.size(); i++)
     this->_blocks[i].set_weights_(
-      i == 0 ? in_channels : channels, channels, dilations[i], batchnorm, activation, groups, it);
+      i == 0 ? in_channels : channels, channels, dilations[i], batchnorm, activation_config, groups, it);
   // Only need _block_vals for the head (one entry)
   // Conv1D layers manage their own buffers now
   this->_block_vals.resize(1);
@@ -327,13 +329,15 @@ std::unique_ptr<nam::DSP> nam::convnet::Factory(const nlohmann::json& config, st
   const int channels = config["channels"];
   const std::vector<int> dilations = config["dilations"];
   const bool batchnorm = config["batchnorm"];
-  const std::string activation = config["activation"];
+  // Parse JSON into typed ActivationConfig at model loading boundary
+  const activations::ActivationConfig activation_config =
+    activations::ActivationConfig::from_json(config["activation"]);
   const int groups = config.value("groups", 1); // defaults to 1
   // Default to 1 channel in/out for backward compatibility
   const int in_channels = config.value("in_channels", 1);
   const int out_channels = config.value("out_channels", 1);
   return std::make_unique<nam::convnet::ConvNet>(
-    in_channels, out_channels, channels, dilations, batchnorm, activation, weights, expectedSampleRate, groups);
+    in_channels, out_channels, channels, dilations, batchnorm, activation_config, weights, expectedSampleRate, groups);
 }
 
 namespace