Add softmax and refactor activation functions

examples/arch/generic/layer/main.c

@@ -16,7 +16,7 @@ int main() {
  const int batch_size = 2;
 
  // Initialize a layer with the given input and output sizes, ReLU activation function, and dot product function
- if (!nn_layer_init(&layer, input_size, output_size, nn_activation_func_relu, nn_dot_product, &error)) {
+ if (!nn_layer_init(&layer, input_size, output_size, &error)) {
  fprintf(stderr, "error: %s\n", error.message);
  return 1;
  }
@@ -33,6 +33,19 @@ int main() {
  return 1;
  }
 
+ // Set the dot product function of the layer
+ if (!nn_layer_set_dot_product_func(&layer, nn_dot_product, &error)) {
+ fprintf(stderr, "error: %s\n", error.message);
+ return 1;
+ }
+
+ // Set the activation function of the layer
+ NNActivationFunction act_func = {.scalar = nn_activation_func_relu};
+ if (!nn_layer_set_activation_func(&layer, act_func, &error)) {
+ fprintf(stderr, "error: %s\n", error.message);
+ return 1;
+ }
+
  // Generate random inputs
  float inputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_INPUT_SIZE];
  for (size_t i = 0; i < batch_size; ++i) {

include/nn_activation.h

@@ -1,8 +1,21 @@
-#ifndef NN_ACTIVATION_FUNCTIONS_H
-#define NN_ACTIVATION_FUNCTIONS_H
+#ifndef NN_ACTIVATION_FUNCTION_H
+#define NN_ACTIVATION_FUNCTION_H
+
+#include "nn_error.h"
+#include <stdbool.h>
+#include <stddef.h>
+
+#ifndef NN_SOFTMAX_MAX_SIZE
+#define NN_SOFTMAX_MAX_SIZE 64
+#endif
 
 // NNActivationFunction represents an activation function.
-typedef float (*NNActivationFunction)(float);
+typedef float (*NNActivationFunctionScalar)(float);
+typedef bool (*NNActivationFunctionVector)(const float input[NN_SOFTMAX_MAX_SIZE], float output[NN_SOFTMAX_MAX_SIZE], size_t input_size, NNError *error);
+typedef union {
+ NNActivationFunctionScalar scalar;
+ NNActivationFunctionVector vector;
+} NNActivationFunction;
 
 // nn_activation_func_identity returns x.
 float nn_activation_func_identity(float x);
@@ -13,4 +26,7 @@ float nn_activation_func_sigmoid(float x);
 // nn_activation_func_relu returns the ReLU of x.
 float nn_activation_func_relu(float x);
 
-#endif // NN_ACTIVATION_FUNCTIONS_H
+// nn_activation_func_softmax calculates the softmax of the input and stores the result in the output.
+bool nn_activation_func_softmax(const float input[NN_SOFTMAX_MAX_SIZE], float output[NN_SOFTMAX_MAX_SIZE], size_t input_size, NNError *error);
+
+#endif // NN_ACTIVATION_FUNCTION_H

include/nn_error.h

@@ -6,8 +6,11 @@
 // NNErrorCode defines the error codes.
 typedef enum {
  NN_ERROR_NONE = 0, // no error
+ NN_ERROR_NOT_IMPLEMENTED, // not implemented
  NN_ERROR_INVALID_INSTANCE, // invalid instance
  NN_ERROR_INVALID_SIZE, // invalid size
+ NN_ERROR_INVALID_VALUE, // invalid value
+ NN_ERROR_INVALID_TYPE, // invalid type
  NN_ERROR_NEON_NOT_AVAILABLE, // NEON instructions not available
  NN_ERROR_CMSIS_DSP_NOT_AVAILABLE, // CMSIS-DSP functions not available
 } NNErrorCode;

include/nn_layer.h

@@ -33,12 +33,12 @@ typedef struct {
  size_t output_size;
  float weights[NN_LAYER_MAX_OUTPUT_SIZE][NN_LAYER_MAX_INPUT_SIZE];
  float biases[NN_LAYER_MAX_BIASES];
- NNActivationFunction act_func;
  NNDotProductFunction dot_product_func;
+ NNActivationFunction act_func;
 } NNLayer;
 
 // nn_layer_init initializes a layer with the given arguments.
-bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNActivationFunction act_func, NNDotProductFunction dot_product_func, NNError *error);
+bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNError *error);
 
 // nn_layer_init_weights_gaussian initializes the weights of the layer with a Gaussian distribution.
 bool nn_layer_init_weights_gaussian(NNLayer *layer, float scale, NNError *error);
@@ -52,6 +52,12 @@ bool nn_layer_set_weights(NNLayer *layer, const float weights[NN_LAYER_MAX_OUTPU
 // nn_layer_set_biases sets the biases of the given layer.
 bool nn_layer_set_biases(NNLayer *layer, const float biases[NN_LAYER_MAX_BIASES], NNError *error);
 
+// nn_layer_set_dot_product_func sets the dot product function of the given layer.
+bool nn_layer_set_dot_product_func(NNLayer *layer, NNDotProductFunction dot_product_func, NNError *error);
+
+// nn_layer_set_activation_func sets the activation function of the given layer.
+bool nn_layer_set_activation_func(NNLayer *layer, NNActivationFunction act_func, NNError *error);
+
 // nn_layer_forward computes the given layer with the given inputs and stores the result in outputs.
 bool nn_layer_forward(const NNLayer *layer, const float inputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_INPUT_SIZE], float outputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_OUTPUT_SIZE], size_t batch_size, NNError *error);
 

include/nn_neuron.h

@@ -19,19 +19,25 @@ typedef struct {
  float weights[NN_NEURON_MAX_WEIGHTS];
  size_t input_size;
  float bias;
- NNActivationFunction act_func;
  NNDotProductFunction dot_product_func;
+ NNActivationFunctionScalar act_func;
 } NNNeuron;
 
 // nn_neuron_init initializes a neuron with the given arguments.
-bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], size_t input_size, float bias, NNActivationFunction act_func, NNDotProductFunction dot_product_func, NNError *error);
+bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], size_t input_size, float bias, NNError *error);
 
 // nn_neuron_set_weights sets the weights of the given neuron.
 bool nn_neuron_set_weights(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], NNError *error);
 
 // nn_neuron_set_bias sets the bias of the given neuron.
 bool nn_neuron_set_bias(NNNeuron *neuron, float bias, NNError *error);
 
+// nn_neuron_set_dot_product_func sets the dot product function of the given neuron.
+bool nn_neuron_set_dot_product_func(NNNeuron *neuron, NNDotProductFunction dot_product_func, NNError *error);
+
+// nn_neuron_set_activation_func sets the activation function of the given neuron.
+bool nn_neuron_set_activation_func(NNNeuron *neuron, NNActivationFunctionScalar act_func, NNError *error);
+
 // nn_neuron_compute computes the given neuron and returns the output.
 float nn_neuron_compute(const NNNeuron *neuron, const float inputs[NN_NEURON_MAX_WEIGHTS], NNError *error);
 

src/nn_activation.c

@@ -1,8 +1,9 @@
 #include "nn_activation.h"
+#include "nn_error.h"
 #include <math.h>
+#include <stdbool.h>
 
 // TODO: Add tests
-// TODO: Add softmax activation function.
 
 // nn_activation_func_identity returns x.
 float nn_activation_func_identity(float x) {
@@ -18,3 +19,45 @@ float nn_activation_func_sigmoid(float x) {
 float nn_activation_func_relu(float x) {
  return fmaxf(0, x);
 }
+
+// nn_activation_func_softmax calculates the softmax of the input and stores the result in the output.
+bool nn_activation_func_softmax(const float input[NN_SOFTMAX_MAX_SIZE], float output[NN_SOFTMAX_MAX_SIZE], size_t input_size, NNError *error) {
+ nn_error_set(error, NN_ERROR_NONE, NULL);
+ if (input == NULL) {
+ nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "input is NULL");
+ return false;
+ } else if (output == NULL) {
+ nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "output is NULL");
+ return false;
+ } else if (input_size == 0 || input_size > NN_SOFTMAX_MAX_SIZE) {
+ nn_error_set(error, NN_ERROR_INVALID_SIZE, "invalid input size");
+ return false;
+ }
+
+ // Find the maximum input value
+ float max_input = input[0];
+ for (size_t i = 1; i < input_size; ++i) {
+ if (input[i] > max_input) {
+ max_input = input[i];
+ }
+ }
+
+ // Compute exp(input[i] - max_input) to prevent overflow
+ float sum = 0.0f;
+ for (size_t i = 0; i < input_size; ++i) {
+ output[i] = expf(input[i] - max_input);
+ sum += output[i];
+ }
+
+ if (sum == 0.0f) {
+ nn_error_set(error, NN_ERROR_INVALID_VALUE, "sum is zero");
+ return false;
+ }
+
+ // Normalize to form a probability distribution
+ for (size_t i = 0; i < input_size; ++i) {
+ output[i] /= sum;
+ }
+
+ return true;
+}

src/nn_layer.c

@@ -11,7 +11,7 @@
 #endif
 
 // nn_layer_init initializes a layer with the given arguments.
-bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNActivationFunction act_func, NNDotProductFunction dot_product_func, NNError *error) {
+bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNError *error) {
  nn_error_set(error, NN_ERROR_NONE, NULL);
  if (layer == NULL) {
  nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "layer is NULL");
@@ -27,12 +27,6 @@ bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNActi
  }
  layer->input_size = input_size;
  layer->output_size = output_size;
- if (act_func) {
- layer->act_func = act_func;
- }
- if (dot_product_func) {
- layer->dot_product_func = dot_product_func;
- }
 
  return true;
 }
@@ -104,6 +98,30 @@ bool nn_layer_set_biases(NNLayer *layer, const float biases[NN_LAYER_MAX_BIASES]
  return true;
 }
 
+// nn_layer_set_dot_product_func sets the dot product function of the given layer.
+bool nn_layer_set_dot_product_func(NNLayer *layer, NNDotProductFunction dot_product_func, NNError *error) {
+ nn_error_set(error, NN_ERROR_NONE, NULL);
+ if (layer == NULL) {
+ nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "layer is NULL");
+ return false;
+ }
+ layer->dot_product_func = dot_product_func;
+
+ return true;
+}
+
+// nn_layer_set_activation_func sets the activation function of the given layer.
+bool nn_layer_set_activation_func(NNLayer *layer, NNActivationFunction act_func, NNError *error) {
+ nn_error_set(error, NN_ERROR_NONE, NULL);
+ if (layer == NULL) {
+ nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "layer is NULL");
+ return false;
+ }
+ layer->act_func = act_func;
+
+ return true;
+}
+
 // nn_layer_forward computes the given layer with the given inputs and stores the result in outputs.
 bool nn_layer_forward(const NNLayer *layer, const float inputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_INPUT_SIZE], float outputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_OUTPUT_SIZE], size_t batch_size, NNError *error) {
  nn_error_set(error, NN_ERROR_NONE, NULL);
@@ -123,8 +141,8 @@ bool nn_layer_forward(const NNLayer *layer, const float inputs[NN_LAYER_MAX_BATC
  if (layer->dot_product_func != NULL) {
  outputs[i][j] += layer->dot_product_func(inputs[i], layer->weights[j], layer->input_size);
  }
- if (layer->act_func != NULL) {
- outputs[i][j] = layer->act_func(outputs[i][j]);
+ if (layer->act_func.scalar != NULL) {
+ outputs[i][j] = layer->act_func.scalar(outputs[i][j]);
  }
  }
  }

src/nn_neuron.c

@@ -7,7 +7,7 @@
 #include <stdio.h>
 
 // nn_neuron_init initializes a neuron with the given arguments.
-bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], size_t input_size, float bias, NNActivationFunction act_func, NNDotProductFunction dot_product_func, NNError *error) {
+bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], size_t input_size, float bias, NNError *error) {
  nn_error_set(error, NN_ERROR_NONE, NULL);
  if (neuron == NULL) {
  nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "neuron is NULL");
@@ -21,12 +21,6 @@ bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS]
  neuron->weights[i] = weights[i];
  }
  neuron->bias = bias;
- if (act_func != NULL) {
- neuron->act_func = act_func;
- }
- if (dot_product_func != NULL) {
- neuron->dot_product_func = dot_product_func;
- }
  return true;
 }
 
@@ -54,6 +48,28 @@ bool nn_neuron_set_bias(NNNeuron *neuron, float bias, NNError *error) {
  return true;
 }
 
+// nn_neuron_set_dot_product_func sets the dot product function of the given neuron.
+bool nn_neuron_set_dot_product_func(NNNeuron *neuron, NNDotProductFunction dot_product_func, NNError *error) {
+ nn_error_set(error, NN_ERROR_NONE, NULL);
+ if (neuron == NULL) {
+ nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "neuron is NULL");
+ return false;
+ }
+ neuron->dot_product_func = dot_product_func;
+ return true;
+}
+
+// nn_neuron_set_activation_func sets the activation function of the given neuron.
+bool nn_neuron_set_activation_func(NNNeuron *neuron, NNActivationFunctionScalar act_func, NNError *error) {
+ nn_error_set(error, NN_ERROR_NONE, NULL);
+ if (neuron == NULL) {
+ nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "neuron is NULL");
+ return false;
+ }
+ neuron->act_func = act_func;
+ return true;
+}
+
 // nn_neuron_compute computes the given neuron and returns the output.
 float nn_neuron_compute(const NNNeuron *neuron, const float inputs[NN_NEURON_MAX_WEIGHTS], NNError *error) {
  nn_error_set(error, NN_ERROR_NONE, NULL);

tests/arch/arm/cmsis-dsp/neuron/main.c

@@ -29,7 +29,11 @@ void run_test_cases(TestCase *test_cases, int n_cases, char *info, NNDotProductF
  NNNeuron neuron;
  NNError error;
 
- nn_neuron_init(&neuron, tc.weights, tc.input_size, tc.bias, nn_activation_func_identity, dot_product_func, &error);
+ nn_neuron_init(&neuron, tc.weights, tc.input_size, tc.bias, &error);
+ assert(error.code == NN_ERROR_NONE);
+ nn_neuron_set_dot_product_func(&neuron, dot_product_func, &error);
+ assert(error.code == NN_ERROR_NONE);
+ nn_neuron_set_activation_func(&neuron, nn_activation_func_identity, &error);
  assert(error.code == NN_ERROR_NONE);
  const float output = nn_neuron_compute(&neuron, tc.inputs, &error);
  assert(error.code == NN_ERROR_NONE);

tests/arch/arm/cmsis-dsp/neuron_perf/main.c

@@ -32,7 +32,13 @@ int main(int argc, char *argv[]) {
  inputs[i] = (float)rand() / (float)RAND_MAX;
  }
 
- if (!nn_neuron_init(&neuron, weights, input_size, bias, nn_activation_func_identity, nn_dot_product_cmsis, &error)) {
+ if (!nn_neuron_init(&neuron, weights, input_size, bias, &error)) {
+ printf("error: %s\n", error.message);
+ return 1;
+ } else if (!nn_neuron_set_activation_func(&neuron, nn_activation_func_identity, &error)) {
+ printf("error: %s\n", error.message);
+ return 1;
+ } else if (!nn_neuron_set_dot_product_func(&neuron, nn_dot_product_cmsis, &error)) {
  printf("error: %s\n", error.message);
  return 1;
  }

tests/arch/arm/neon/neuron/main.c

@@ -29,7 +29,11 @@ void run_test_cases(TestCase *test_cases, int n_cases, char *info, NNDotProductF
  NNNeuron neuron;
  NNError error;
 
- nn_neuron_init(&neuron, tc.weights, tc.input_size, tc.bias, nn_activation_func_identity, dot_product_func, &error);
+ nn_neuron_init(&neuron, tc.weights, tc.input_size, tc.bias, &error);
+ assert(error.code == NN_ERROR_NONE);
+ nn_neuron_set_dot_product_func(&neuron, dot_product_func, &error);
+ assert(error.code == NN_ERROR_NONE);
+ nn_neuron_set_activation_func(&neuron, nn_activation_func_identity, &error);
  assert(error.code == NN_ERROR_NONE);
  const float output = nn_neuron_compute(&neuron, tc.inputs, &error);
  assert(error.code == NN_ERROR_NONE);

tests/arch/arm/neon/neuron_perf/main.c

@@ -33,7 +33,13 @@ int main(int argc, char *argv[]) {
  inputs[i] = (float)rand() / (float)RAND_MAX;
  }
 
- if (!nn_neuron_init(&neuron, weights, input_size, bias, nn_activation_func_identity, nn_dot_product_neon, &error)) {
+ if (!nn_neuron_init(&neuron, weights, input_size, bias, &error)) {
+ printf("error: %s\n", error.message);
+ return 1;
+ } else if (!nn_neuron_set_activation_func(&neuron, nn_activation_func_identity, &error)) {
+ printf("error: %s\n", error.message);
+ return 1;
+ } else if (!nn_neuron_set_dot_product_func(&neuron, nn_dot_product_neon, &error)) {
  printf("error: %s\n", error.message);
  return 1;
  }