path: root/src/lib/src/train.c

#include <neuralnet/train.h>

#include "neuralnet_impl.h"
#include <neuralnet/matrix.h>

#include <random/mt19937-64.h>
#include <random/normal.h>

#include <assert.h>
#include <math.h>
#include <stdlib.h>

#include <stdio.h>
#define LOGD printf

// If debug mode is requested, we will show progress every this many iterations.
static const int PROGRESS_THRESHOLD = 5; // %

/// Computes the total MSE from the output error matrix.
R ComputeMSE(const nnMatrix* errors) {
  R         sum_sq = 0;
  const int N      = errors->rows * errors->cols;
  const R*  value  = errors->values;
  for (int i = 0; i < N; ++i) {
    sum_sq += *value * *value;
    value++;
  }
  return sum_sq / (R)N;
}

/// Holds the bits required to compute a sigmoid gradient.
typedef struct nnSigmoidGradientElements {
  nnMatrix ones; // A vector of just ones, same size as the layer.
} nnSigmoidGradientElements;

/// Holds the various elements required to compute gradients. These depend on
/// what activation function are used, so they'll potentially be different for
/// each layer. A data type is defined for these because we allocate all the
/// required memory up front before entering the training loop.
typedef struct nnGradientElements {
  nnLayerType type;
  // Gradient vector, same size as the layer.
  // This will contain the gradient expression except for the output value of
  // the previous layer.
  nnMatrix gradient;
  union {
    nnSigmoidGradientElements sigmoid;
  };
} nnGradientElements;

// Initialize the network's weights randomly and set their biases to 0.
void nnInitNet(
    nnNeuralNetwork* net, uint64_t seed, const nnWeightInitStrategy strategy) {
  assert(net);

  mt19937_64 rng = mt19937_64_make();
  mt19937_64_init(&rng, seed);

  for (int l = 0; l < net->num_layers; ++l) {
    // Get the layer's weights and biases, if any.
    nnMatrix* weights = 0;
    nnMatrix* biases  = 0;
    switch (net->layers[l].type) {
    case nnLinear: {
      nnLinearImpl* linear = &net->layers[l].linear;

      weights = &linear->weights;
      biases  = &linear->biases;
      break;
    }
    // Activations.
    case nnRelu:
    case nnSigmoid:
      break;
    }
    if (!weights || !biases) {
      continue;
    }

    const R layer_size = (R)nnLayerInputSize(net, l);
    const R scale      = 1. / layer_size;
    const R stdev      = 1. / sqrt((R)layer_size);
    const R sigma      = stdev * stdev;

    R* value = weights->values;
    for (int k = 0; k < weights->rows * weights->cols; ++k) {
      switch (strategy) {
      case nnWeightInit01: {
        const R x01 = mt19937_64_gen_real3(&rng); // (0, +1) interval.
        *value++    = scale * x01;
        break;
      }
      case nnWeightInit11: {
        const R x11 = mt19937_64_gen_real4(&rng); // (-1, +1) interval.
        *value++    = scale * x11;
        break;
      }
      case nnWeightInitNormal: {
        // Using initialization with a normal distribution of standard
        // deviation 1 / sqrt(num_layer_weights) to prevent saturation when
        // multiplying inputs.
        const R u01 = mt19937_64_gen_real3(&rng); // (0, +1) interval.
        const R v01 = mt19937_64_gen_real3(&rng); // (0, +1) interval.
        R       z0, z1;
        normal2(u01, v01, &z0, &z1);
        z0       = normal_transform(z0, /*mu=*/0, sigma);
        z1       = normal_transform(z1, /*mu=*/0, sigma);
        *value++ = z0;
        if (k < weights->rows * weights->cols - 1) {
          *value++ = z1;
          ++k;
        }
        break;
      }
      default:
        assert(false);
      }
    }

    // Initialize biases.
    // 0 is used so that functions originally go through the origin.
    value = biases->values;
    for (int k = 0; k < biases->rows * biases->cols; ++k, ++value) {
      *value = 0;
    }
  }
}

// |inputs|  has one row vector per sample.
// |targets| has one row vector per sample.
//
// For now, each iteration trains with one sample (row) at a time.
void nnTrain(
    nnNeuralNetwork* net, const nnMatrix* inputs, const nnMatrix* targets,
    const nnTrainingParams* params) {
  assert(net);
  assert(inputs);
  assert(targets);
  assert(params);
  assert(nnNetOutputSize(net) == targets->cols);
  assert(net->num_layers > 0);

  // Allocate error vectors to hold the backpropagated error values.
  // For now, these are one row vector per layer, meaning that we will train
  // with one sample at a time.
  nnMatrix* errors = calloc(net->num_layers, sizeof(nnMatrix));

  // Allocate the weight delta matrices.
  nnMatrix* weight_deltas = calloc(net->num_layers, sizeof(nnMatrix));

  // Allocate the data structures required to compute gradients.
  // This depends on each layer's activation type.
  nnGradientElements* gradient_elems =
      calloc(net->num_layers, sizeof(nnGradientElements));

  assert(errors != 0);
  assert(weight_deltas != 0);
  assert(gradient_elems);

  for (int l = 0; l < net->num_layers; ++l) {
    const int          layer_input_size  = nnLayerInputSize(net, l);
    const int          layer_output_size = nnLayerOutputSize(net, l);
    const nnLayerImpl* layer             = &net->layers[l];

    errors[l]        = nnMatrixMake(1, layer_output_size);
    weight_deltas[l] = nnMatrixMake(layer_input_size, layer_output_size);

    // Allocate the gradient elements and vectors for weight delta calculation.
    nnGradientElements* elems = &gradient_elems[l];
    elems->type               = layer->type;
    switch (layer->type) {
    case nnLinear:
      break; // Gradient vector will be borrowed, no need to allocate.

    case nnSigmoid:
      elems->gradient = nnMatrixMake(1, layer_output_size);
      // Allocate the 1s vectors.
      elems->sigmoid.ones = nnMatrixMake(1, layer_output_size);
      nnMatrixInitConstant(&elems->sigmoid.ones, 1);
      break;

    case nnRelu:
      elems->gradient = nnMatrixMake(1, layer_output_size);
      break;
    }
  }

  // Construct the query object with a size of 1 since we are training with one
  // sample at a time.
  nnQueryObject* query = nnMakeQueryObject(net, 1);

  // Network outputs are given by the query object. Every network query updates
  // the outputs.
  const nnMatrix* const training_outputs = query->network_outputs;

  // If debug mode is requested, we will show progress every Nth iteration.
  const int progress_frame =
      (params->max_iterations < PROGRESS_THRESHOLD)
          ? 1
          : (params->max_iterations * PROGRESS_THRESHOLD / 100);

  // --- TRAIN

  nnInitNet(net, params->seed, params->weight_init);

  for (int iter = 0; iter < params->max_iterations; ++iter) {

    // For now, we train with one sample at a time.
    for (int sample = 0; sample < inputs->rows; ++sample) {
      // Slice the input and target matrices with the batch size.
      // We are not mutating the inputs, but we need the cast to borrow.
      const nnMatrix training_inputs =
          nnMatrixBorrowRows((nnMatrix*)inputs, sample, 1);
      const nnMatrix training_targets =
          nnMatrixBorrowRows((nnMatrix*)targets, sample, 1);

      // Forward pass.
      nnQuery(net, query, &training_inputs);

      // Compute the error derivative: o-t.
      //   Error: 1/2 (t-o)^2
      //   dE/do = -(t-o)
      //         = +(o-t)
      // Note that we compute o-t instead to remove that outer negative sign.
      // The 2 is dropped because we are only interested in the direction of the
      // gradient. The learning rate controls the magnitude.
      nnMatrixSub(
          training_outputs, &training_targets, &errors[net->num_layers - 1]);

      // Update weights and biases for each internal layer, back-propagating
      // errors along the way.
      for (int l = net->num_layers - 1; l >= 0; --l) {
        const nnMatrix* layer_input =
            (l == 0) ? &training_inputs : &query->layer_outputs[l - 1];
        const nnMatrix*     layer_output = &query->layer_outputs[l];
        nnGradientElements* elems        = &gradient_elems[l];
        nnMatrix*           gradient     = &elems->gradient;
        nnLayerImpl*        layer        = &net->layers[l];

        // Compute this layer's gradient.
        //
        // By 'gradient' we mean the subexpression common to all the gradients
        // for this layer.
        // For linear layers, this is the subexpression common to both the
        // weights and bias gradients.
        //
        // Linear:   G = id
        // Relu:     G = (output_k > 0 ? 1 : 0)
        // Sigmoid:  G = output_k * (1 - output_k)
        switch (layer->type) {
        case nnLinear: {
          break;
        }
        case nnRelu:
          nnMatrixGt(layer_output, 0, gradient);
          break;
        case nnSigmoid:
          nnMatrixSub(&elems->sigmoid.ones, layer_output, gradient);
          nnMatrixMulPairs(layer_output, gradient, gradient);
          break;
        }

        // Back-propagate the error.
        //
        // This combines this layer's gradient with the back-propagated error,
        // which is the combination of the gradients of subsequent layers down
        // to the output layer error.
        //
        // Note that this step uses the layer's original weights.
        if (l > 0) {
          switch (layer->type) {
          case nnLinear: {
            const nnMatrix* layer_weights = &layer->linear.weights;
            // E * W^T == E *^T W.
            // Using nnMatrixMulRows, we avoid having to transpose the weight
            // matrix.
            nnMatrixMulRows(&errors[l], layer_weights, &errors[l - 1]);
            break;
          }
          // For activations, the error back-propagates as is but multiplied by
          // the layer's gradient.
          case nnRelu:
          case nnSigmoid:
            nnMatrixMulPairs(&errors[l], gradient, &errors[l - 1]);
            break;
          }
        }

        // Update layer weights.
        if (layer->type == nnLinear) {
          nnLinearImpl* linear        = &layer->linear;
          nnMatrix*     layer_weights = &linear->weights;
          nnMatrix*     layer_biases  = &linear->biases;

          // Outer product to compute the weight deltas.
          nnMatrixMulOuter(layer_input, &errors[l], &weight_deltas[l]);

          // Update weights.
          nnMatrixScale(&weight_deltas[l], params->learning_rate);
          nnMatrixSub(layer_weights, &weight_deltas[l], layer_weights);

          // Update biases.
          // This is the same formula as for weights, except that the o_j term
          // is just 1.
          nnMatrixMulSub(
              layer_biases, &errors[l], params->learning_rate, layer_biases);
        }
      }

      // TODO: Add this under a verbose debugging mode.
      // if (params->debug) {
      //   LOGD("Iter: %d, Sample: %d, Error: %f\n", iter, sample,
      //   ComputeMSE(&errors[net->num_layers - 1])); LOGD("TGT: "); for (int i
      //   = 0; i < training_targets.cols; ++i) {
      //     printf("%.3f  ", training_targets.values[i]);
      //   }
      //   printf("\n");
      //   LOGD("OUT: ");
      //   for (int i = 0; i < training_outputs->cols; ++i) {
      //     printf("%.3f  ", training_outputs->values[i]);
      //   }
      //   printf("\n");
      // }
    }

    if (params->debug && ((iter % progress_frame) == 0)) {
      LOGD(
          "Iter: %d/%d, Error: %f\n", iter, params->max_iterations,
          ComputeMSE(&errors[net->num_layers - 1]));
    }
  }

  // Print the final error.
  if (params->debug) {
    LOGD(
        "Iter: %d/%d, Error: %f\n", params->max_iterations,
        params->max_iterations, ComputeMSE(&errors[net->num_layers - 1]));
  }

  // Clean up.
  for (int l = 0; l < net->num_layers; ++l) {
    nnMatrixDel(&errors[l]);
    nnMatrixDel(&weight_deltas[l]);

    nnGradientElements* elems = &gradient_elems[l];
    switch (elems->type) {
    case nnLinear:
      break; // Gradient vector is borrowed, no need to deallocate.

    case nnSigmoid:
      nnMatrixDel(&elems->gradient);
      nnMatrixDel(&elems->sigmoid.ones);
      break;

    case nnRelu:
      nnMatrixDel(&elems->gradient);
      break;
    }
  }
  free(errors);
  free(weight_deltas);
  free(gradient_elems);
}