Modularizing the RNN model

7 files changed, 88 insertions, 47 deletions

diff --git a/examples/rnnoise_demo.c b/examples/rnnoise_demo.c
index e1e239a..83d0709 100644
--- a/examples/rnnoise_demo.c
+++ b/examples/rnnoise_demo.c
@@ -1,4 +1,5 @@
-/* Copyright (c) 2017 Mozilla */
+/* Copyright (c) 2018 Gregor Richards
+ * Copyright (c) 2017 Mozilla */
 /*
    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions
@@ -35,7 +36,7 @@ int main(int argc, char **argv) {
   float x[FRAME_SIZE];
   FILE *f1, *fout;
   DenoiseState *st;
-  st = rnnoise_create();
+  st = rnnoise_create(NULL);
   if (argc!=3) {
     fprintf(stderr, "usage: %s <noisy speech> <output denoised>\n", argv[0]);
     return 1;
diff --git a/include/rnnoise.h b/include/rnnoise.h
index 562b34c..3e56b9f 100644
--- a/include/rnnoise.h
+++ b/include/rnnoise.h
@@ -1,4 +1,5 @@
-/* Copyright (c) 2017 Mozilla */
+/* Copyright (c) 2018 Gregor Richards
+ * Copyright (c) 2017 Mozilla */
 /*
    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions
@@ -40,12 +41,13 @@
 
 
 typedef struct DenoiseState DenoiseState;
+typedef struct RNNModel RNNModel;
 
 RNNOISE_EXPORT int rnnoise_get_size();
 
-RNNOISE_EXPORT int rnnoise_init(DenoiseState *st);
+RNNOISE_EXPORT int rnnoise_init(DenoiseState *st, RNNModel *model);
 
-RNNOISE_EXPORT DenoiseState *rnnoise_create();
+RNNOISE_EXPORT DenoiseState *rnnoise_create(RNNModel *model);
 
 RNNOISE_EXPORT void rnnoise_destroy(DenoiseState *st);
 
diff --git a/src/denoise.c b/src/denoise.c
index 128cd99..a350ffa 100644
--- a/src/denoise.c
+++ b/src/denoise.c
@@ -1,4 +1,5 @@
-/* Copyright (c) 2017 Mozilla */
+/* Copyright (c) 2018 Gregor Richards
+ * Copyright (c) 2017 Mozilla */
 /*
    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions
@@ -70,6 +71,11 @@
 #define TRAINING 0
 #endif
 
+
+/* The built-in model, used if no file is given as input */
+extern const struct RNNModel rnnoise_model_orig;
+
+
 static const opus_int16 eband5ms[] = {
 /*0  200 400 600 800  1k 1.2 1.4 1.6  2k 2.4 2.8 3.2  4k 4.8 5.6 6.8  8k 9.6 12k 15.6 20k*/
   0,  1,  2,  3,  4,  5,  6,  7,  8, 10, 12, 14, 16, 20, 24, 28, 34, 40, 48, 60, 78, 100
@@ -284,19 +290,29 @@ int rnnoise_get_size() {
   return sizeof(DenoiseState);
 }
 
-int rnnoise_init(DenoiseState *st) {
+int rnnoise_init(DenoiseState *st, RNNModel *model) {
   memset(st, 0, sizeof(*st));
+  if (model)
+    st->rnn.model = model;
+  else
+    st->rnn.model = &rnnoise_model_orig;
+  st->rnn.vad_gru_state = calloc(sizeof(float), st->rnn.model->vad_gru_size);
+  st->rnn.noise_gru_state = calloc(sizeof(float), st->rnn.model->noise_gru_size);
+  st->rnn.denoise_gru_state = calloc(sizeof(float), st->rnn.model->denoise_gru_size);
   return 0;
 }
 
-DenoiseState *rnnoise_create() {
+DenoiseState *rnnoise_create(RNNModel *model) {
   DenoiseState *st;
   st = malloc(rnnoise_get_size());
-  rnnoise_init(st);
+  rnnoise_init(st, model);
   return st;
 }
 
 void rnnoise_destroy(DenoiseState *st) {
+  free(st->rnn.vad_gru_state);
+  free(st->rnn.noise_gru_state);
+  free(st->rnn.denoise_gru_state);
   free(st);
 }
 
@@ -542,9 +558,9 @@ int main(int argc, char **argv) {
   DenoiseState *st;
   DenoiseState *noise_state;
   DenoiseState *noisy;
-  st = rnnoise_create();
-  noise_state = rnnoise_create();
-  noisy = rnnoise_create();
+  st = rnnoise_create(NULL);
+  noise_state = rnnoise_create(NULL);
+  noisy = rnnoise_create(NULL);
   if (argc!=4) {
     fprintf(stderr, "usage: %s <speech> <noise> <output denoised>\n", argv[0]);
     return 1;
diff --git a/src/rnn.c b/src/rnn.c
index 1daa7b6..c54958e 100644
--- a/src/rnn.c
+++ b/src/rnn.c
@@ -162,17 +162,17 @@ void compute_rnn(RNNState *rnn, float *gains, float *vad, const float *input) {
   float dense_out[MAX_NEURONS];
   float noise_input[MAX_NEURONS*3];
   float denoise_input[MAX_NEURONS*3];
-  compute_dense(&input_dense, dense_out, input);
-  compute_gru(&vad_gru, rnn->vad_gru_state, dense_out);
-  compute_dense(&vad_output, vad, rnn->vad_gru_state);
-  for (i=0;i<INPUT_DENSE_SIZE;i++) noise_input[i] = dense_out[i];
-  for (i=0;i<VAD_GRU_SIZE;i++) noise_input[i+INPUT_DENSE_SIZE] = rnn->vad_gru_state[i];
-  for (i=0;i<INPUT_SIZE;i++) noise_input[i+INPUT_DENSE_SIZE+VAD_GRU_SIZE] = input[i];
-  compute_gru(&noise_gru, rnn->noise_gru_state, noise_input);
+  compute_dense(rnn->model->input_dense, dense_out, input);
+  compute_gru(rnn->model->vad_gru, rnn->vad_gru_state, dense_out);
+  compute_dense(rnn->model->vad_output, vad, rnn->vad_gru_state);
+  for (i=0;i<rnn->model->input_dense_size;i++) noise_input[i] = dense_out[i];
+  for (i=0;i<rnn->model->vad_gru_size;i++) noise_input[i+rnn->model->input_dense_size] = rnn->vad_gru_state[i];
+  for (i=0;i<INPUT_SIZE;i++) noise_input[i+rnn->model->input_dense_size+rnn->model->vad_gru_size] = input[i];
+  compute_gru(rnn->model->noise_gru, rnn->noise_gru_state, noise_input);
 
-  for (i=0;i<VAD_GRU_SIZE;i++) denoise_input[i] = rnn->vad_gru_state[i];
-  for (i=0;i<NOISE_GRU_SIZE;i++) denoise_input[i+VAD_GRU_SIZE] = rnn->noise_gru_state[i];
-  for (i=0;i<INPUT_SIZE;i++) denoise_input[i+VAD_GRU_SIZE+NOISE_GRU_SIZE] = input[i];
-  compute_gru(&denoise_gru, rnn->denoise_gru_state, denoise_input);
-  compute_dense(&denoise_output, gains, rnn->denoise_gru_state);
+  for (i=0;i<rnn->model->vad_gru_size;i++) denoise_input[i] = rnn->vad_gru_state[i];
+  for (i=0;i<rnn->model->noise_gru_size;i++) denoise_input[i+rnn->model->vad_gru_size] = rnn->noise_gru_state[i];
+  for (i=0;i<INPUT_SIZE;i++) denoise_input[i+rnn->model->vad_gru_size+rnn->model->noise_gru_size] = input[i];
+  compute_gru(rnn->model->denoise_gru, rnn->denoise_gru_state, denoise_input);
+  compute_dense(rnn->model->denoise_output, gains, rnn->denoise_gru_state);
 }
diff --git a/src/rnn.h b/src/rnn.h
index 9e08b44..e763895 100644
--- a/src/rnn.h
+++ b/src/rnn.h
@@ -56,6 +56,7 @@ typedef struct {
   int activation;
 } GRULayer;
 
+typedef struct RNNModel RNNModel;
 typedef struct RNNState RNNState;
 
 void compute_dense(const DenseLayer *layer, float *output, const float *input);
diff --git a/src/rnn_data.c b/src/rnn_data.c
index 8f6c99b..22c5316 100644
--- a/src/rnn_data.c
+++ b/src/rnn_data.c
@@ -5,6 +5,7 @@
 #endif
 
 #include "rnn.h"
+#include "rnn_data.h"
 
 static const rnn_weight input_dense_weights[1008] = {
    -10, 0, -3, 1, -8, -6, 3, -13,
@@ -141,7 +142,7 @@ static const rnn_weight input_dense_bias[24] = {
    -126, 28, 127, 125, -30, 127, -89, -20
 };
 
-const DenseLayer input_dense = {
+static const DenseLayer input_dense = {
    input_dense_bias,
    input_dense_weights,
    42, 24, ACTIVATION_TANH
@@ -597,7 +598,7 @@ static const rnn_weight vad_gru_bias[72] = {
    -29, 127, 34, -66, 49, 53, 27, 62
 };
 
-const GRULayer vad_gru = {
+static const GRULayer vad_gru = {
    vad_gru_bias,
    vad_gru_weights,
    vad_gru_recurrent_weights,
@@ -3115,7 +3116,7 @@ static const rnn_weight noise_gru_bias[144] = {
    -23, -64, 31, 86, -50, 2, -38, 7
 };
 
-const GRULayer noise_gru = {
+static const GRULayer noise_gru = {
    noise_gru_bias,
    noise_gru_weights,
    noise_gru_recurrent_weights,
@@ -10727,7 +10728,7 @@ static const rnn_weight denoise_gru_bias[288] = {
    -21, 25, 18, -58, 25, 126, -84, 127
 };
 
-const GRULayer denoise_gru = {
+static const GRULayer denoise_gru = {
    denoise_gru_bias,
    denoise_gru_weights,
    denoise_gru_recurrent_weights,
@@ -11007,7 +11008,7 @@ static const rnn_weight denoise_output_bias[22] = {
    -126, -105, -53, -49, -18, -9
 };
 
-const DenseLayer denoise_output = {
+static const DenseLayer denoise_output = {
    denoise_output_bias,
    denoise_output_weights,
    96, 22, ACTIVATION_SIGMOID
@@ -11023,9 +11024,28 @@ static const rnn_weight vad_output_bias[1] = {
    -50
 };
 
-const DenseLayer vad_output = {
+static const DenseLayer vad_output = {
    vad_output_bias,
    vad_output_weights,
    24, 1, ACTIVATION_SIGMOID
 };
 
+const struct RNNModel rnnoise_model_orig = {
+    24,
+    &input_dense,
+
+    24,
+    &vad_gru,
+
+    48,
+    &noise_gru,
+
+    96,
+    &denoise_gru,
+
+    22,
+    &denoise_output,
+
+    1,
+    &vad_output
+};
diff --git a/src/rnn_data.h b/src/rnn_data.h
index 5610980..2aa41f9 100644
--- a/src/rnn_data.h
+++ b/src/rnn_data.h
@@ -1,32 +1,33 @@
-/*This file is automatically generated from a Keras model*/
-
 #ifndef RNN_DATA_H
 #define RNN_DATA_H
 
 #include "rnn.h"
 
-#define INPUT_DENSE_SIZE 24
-extern const DenseLayer input_dense;
+struct RNNModel {
+  int input_dense_size;
+  DenseLayer *input_dense;
 
-#define VAD_GRU_SIZE 24
-extern const GRULayer vad_gru;
+  int vad_gru_size;
+  GRULayer *vad_gru;
 
-#define NOISE_GRU_SIZE 48
-extern const GRULayer noise_gru;
+  int noise_gru_size;
+  GRULayer *noise_gru;
 
-#define DENOISE_GRU_SIZE 96
-extern const GRULayer denoise_gru;
+  int denoise_gru_size;
+  GRULayer *denoise_gru;
 
-#define DENOISE_OUTPUT_SIZE 22
-extern const DenseLayer denoise_output;
+  int denoise_output_size;
+  DenseLayer *denoise_output;
 
-#define VAD_OUTPUT_SIZE 1
-extern const DenseLayer vad_output;
+  int vad_output_size;
+  DenseLayer *vad_output;
+};
 
 struct RNNState {
-  float vad_gru_state[VAD_GRU_SIZE];
-  float noise_gru_state[NOISE_GRU_SIZE];
-  float denoise_gru_state[DENOISE_GRU_SIZE];
+  const RNNModel *model;
+  float *vad_gru_state;
+  float *noise_gru_state;
+  float *denoise_gru_state;
 };