Train rnn with microsoft SNSD

Makefile.am

@@ -4,9 +4,9 @@ AM_CFLAGS = -I$(top_srcdir)/include $(DEPS_CFLAGS)
 
 dist_doc_DATA = COPYING AUTHORS README
 
-include_HEADERS = include/rnnoise.h
+include_HEADERS = include/rnnoise-nu.h
 
-lib_LTLIBRARIES = librnnoise.la
+lib_LTLIBRARIES = librnnoise-nu.la
 noinst_HEADERS = src/arch.h  \
 		 src/celt_lpc.h  \
 		 src/common.h  \
@@ -18,42 +18,43 @@ noinst_HEADERS = src/arch.h  \
 		 src/rnn.h  \
 		 src/tansig_table.h
 
-librnnoise_la_SOURCES = \
+librnnoise_nu_la_SOURCES = \
 	src/denoise.c \
 	src/rnn.c \
 	src/rnn_data.c \
+	src/rnn_reader.c \
 	src/pitch.c \
 	src/kiss_fft.c \
 	src/celt_lpc.c
 
-librnnoise_la_LIBADD = $(DEPS_LIBS) $(lrintf_lib) $(LIBM)
-librnnoise_la_LDFLAGS = -no-undefined \
+librnnoise_nu_la_LIBADD = $(DEPS_LIBS) $(lrintf_lib) $(LIBM)
+librnnoise_nu_la_LDFLAGS = -no-undefined \
  -version-info @OP_LT_CURRENT@:@OP_LT_REVISION@:@OP_LT_AGE@
 
 if OP_ENABLE_EXAMPLES
-noinst_PROGRAMS = examples/rnnoise_demo
+noinst_PROGRAMS = examples/rnnoise_nu_demo
 endif
 
-examples_rnnoise_demo_SOURCES = examples/rnnoise_demo.c
-examples_rnnoise_demo_LDADD = librnnoise.la
+examples_rnnoise_nu_demo_SOURCES = examples/rnnoise_demo.c
+examples_rnnoise_nu_demo_LDADD = librnnoise-nu.la
 
 pkgconfigdir = $(libdir)/pkgconfig
-pkgconfig_DATA = rnnoise.pc
+pkgconfig_DATA = rnnoise-nu.pc
 
 debug:
 	$(MAKE) CFLAGS="${CFLAGS} -O0 -ggdb -DOP_ENABLE_ASSERTIONS" all
 
 EXTRA_DIST = \
- rnnoise.pc.in \
- rnnoise-uninstalled.pc.in \
+ rnnoise-nu.pc.in \
+ rnnoise-nu-uninstalled.pc.in \
  doc/Doxyfile.in \
  doc/Makefile
 
 # Targets to build and install just the library without the docs
-librnnoise install-librnnoise: NO_DOXYGEN = 1
+librnnoise-nu install-librnnoise-nu: NO_DOXYGEN = 1
 
-rnnoise: all
-install-rnnoise: install
+rnnoise-nu: all
+install-rnnoise-nu: install
 
 # Or just the docs
 docs: doc/doxygen-build.stamp
@@ -126,4 +127,4 @@ dist-hook:
 	echo 'PACKAGE_VERSION="$(PACKAGE_VERSION)"' > $(top_distdir)/package_version
 
 
-.PHONY: rnnoise install-rnnoise docs install-docs
+.PHONY: rnnoise-nu install-rnnoise-nu docs install-docs

configure.ac

@@ -31,16 +31,14 @@ m4_ifdef([AM_SILENT_RULES], [AM_SILENT_RULES([yes])])
 AC_DEFINE([RNNOISE_BUILD], [], [This is a build of the library])
 
 dnl Library versioning for libtool.
-dnl Please update these for releases.
-dnl CURRENT, REVISION, AGE
-dnl - library source changed -> increment REVISION
-dnl - interfaces added/removed/changed -> increment CURRENT, REVISION = 0
-dnl - interfaces added -> increment AGE
-dnl - interfaces removed -> AGE = 0
+dnl For some inane reason, in .so.major.minor.revision:
+dnl - major = CURRENT-AGE
+dnl - minor = AGE
+dnl - revision = REVISION
 
-OP_LT_CURRENT=4
-OP_LT_REVISION=1
-OP_LT_AGE=4
+OP_LT_CURRENT=1
+OP_LT_REVISION=0
+OP_LT_AGE=0
 
 AC_SUBST(OP_LT_CURRENT)
 AC_SUBST(OP_LT_REVISION)
@@ -108,8 +106,8 @@ AM_CONDITIONAL([HAVE_DOXYGEN], [test "$HAVE_DOXYGEN" = "yes"])
 
 AC_CONFIG_FILES([
   Makefile
-  rnnoise.pc
-  rnnoise-uninstalled.pc
+  rnnoise-nu.pc
+  rnnoise-nu-uninstalled.pc
   doc/Doxyfile
 ])
 AC_CONFIG_HEADERS([config.h])

doc/Doxyfile.in

@@ -3,7 +3,7 @@
 PROJECT_NAME           = @PACKAGE_NAME@
 PROJECT_NUMBER         = @PACKAGE_VERSION@
 PROJECT_BRIEF          = "RNN-based noise suppressor."
-INPUT                  = @top_srcdir@/include/rnnoise.h
+INPUT                  = @top_srcdir@/include/rnnoise-nu.h
 OPTIMIZE_OUTPUT_FOR_C  = YES
 
 QUIET                  = YES

examples/rnnoise_demo.c

@@ -25,7 +25,7 @@
 */
 
 #include <stdio.h>
-#include "rnnoise.h"
+#include "rnnoise-nu.h"
 
 #define FRAME_SIZE 480
 
@@ -35,7 +35,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;

include/rnnoise.h → include/rnnoise-nu.h

@@ -24,6 +24,12 @@
    SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
 
+#ifndef RNNOISE_NU_H
+#define RNNOISE_NU_H 1
+
+#include <stdio.h>
+
+
 #ifndef RNNOISE_EXPORT
 # if defined(WIN32)
 #  if defined(RNNOISE_BUILD) && defined(DLL_EXPORT)
@@ -38,15 +44,21 @@
 # endif
 #endif
 
-
 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);
 
 RNNOISE_EXPORT float rnnoise_process_frame(DenoiseState *st, float *out, const float *in);
+
+RNNOISE_EXPORT RNNModel *rnnoise_model_from_file(FILE *f);
+
+RNNOISE_EXPORT void rnnoise_model_free(RNNModel *model);
+
+#endif

rnnoise-uninstalled.pc.in → rnnoise-nu-uninstalled.pc.in

@@ -5,9 +5,9 @@ exec_prefix=
 libdir=${pcfiledir}/.libs
 includedir=${pcfiledir}/@top_srcdir@/include
 
-Name: rnnoise uninstalled
-Description: RNN-based noise suppression (not installed)
+Name: rnnoise-nu uninstalled
+Description: RNN-based noise suppression, improved (not installed)
 Version: @PACKAGE_VERSION@
 Conflicts:
-Libs: ${libdir}/librnnoise.la @lrintf_lib@
+Libs: ${libdir}/librnnoise-nu.la @lrintf_lib@
 Cflags: -I${includedir}

rnnoise.pc.in → rnnoise-nu.pc.in

@@ -5,10 +5,10 @@ exec_prefix=@exec_prefix@
 libdir=@libdir@
 includedir=@includedir@
 
-Name: rnnoise
-Description: RNN-based noise suppression
+Name: rnnoise-nu
+Description: RNN-based noise suppression, improved
 Version: @PACKAGE_VERSION@
 Conflicts:
-Libs: -L${libdir} -lrnnoise
+Libs: -L${libdir} -lrnnoise-nu
 Libs.private: @lrintf_lib@
 Cflags: -I${includedir}/

src/denoise.c

@@ -34,7 +34,7 @@
 #include "kiss_fft.h"
 #include "common.h"
 #include <math.h>
-#include "rnnoise.h"
+#include "rnnoise-nu.h"
 #include "pitch.h"
 #include "arch.h"
 #include "rnn.h"
@@ -70,6 +70,10 @@
 #define TRAINING 0
 #endif
 
+
+extern const struct RNNModel 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,15 +288,19 @@ 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 = &model_orig;
   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;
 }
 
@@ -538,20 +546,21 @@ int main(int argc, char **argv) {
   int vad_cnt=0;
   int gain_change_count=0;
   float speech_gain = 1, noise_gain = 1;
-  FILE *f1, *f2, *fout;
+  FILE *f1, *f2;
+  int maxCount;
   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]);
+    fprintf(stderr, "usage: %s <speech> <noise> <count>\n", argv[0]);
     return 1;
   }
   f1 = fopen(argv[1], "r");
   f2 = fopen(argv[2], "r");
-  fout = fopen(argv[3], "w");
+  maxCount = atoi(argv[3]);
   for(i=0;i<150;i++) {
     short tmp[FRAME_SIZE];
     fread(tmp, sizeof(short), FRAME_SIZE, f2);
@@ -563,12 +572,11 @@ int main(int argc, char **argv) {
     float Ln[NB_BANDS];
     float features[NB_FEATURES];
     float g[NB_BANDS];
-    float gf[FREQ_SIZE]={1};
     short tmp[FRAME_SIZE];
     float vad=0;
-    float vad_prob;
     float E=0;
-    if (count==50000000) break;
+    if (count==maxCount) break;
+    if ((count%1000)==0) fprintf(stderr, "%d\r", count);
     if (++gain_change_count > 2821) {
       speech_gain = pow(10., (-40+(rand()%60))/20.);
       noise_gain = pow(10., (-30+(rand()%50))/20.);
@@ -643,37 +651,16 @@ int main(int argc, char **argv) {
       if (vad==0 && noise_gain==0) g[i] = -1;
     }
     count++;
-#if 0
-    for (i=0;i<NB_FEATURES;i++) printf("%f ", features[i]);
-    for (i=0;i<NB_BANDS;i++) printf("%f ", g[i]);
-    for (i=0;i<NB_BANDS;i++) printf("%f ", Ln[i]);
-    printf("%f\n", vad);
-#endif
 #if 1
     fwrite(features, sizeof(float), NB_FEATURES, stdout);
     fwrite(g, sizeof(float), NB_BANDS, stdout);
     fwrite(Ln, sizeof(float), NB_BANDS, stdout);
     fwrite(&vad, sizeof(float), 1, stdout);
-#endif
-#if 0
-    compute_rnn(&noisy->rnn, g, &vad_prob, features);
-    interp_band_gain(gf, g);
-#if 1
-    for (i=0;i<FREQ_SIZE;i++) {
-      X[i].r *= gf[i];
-      X[i].i *= gf[i];
-    }
-#endif
-    frame_synthesis(noisy, xn, X);
-
-    for (i=0;i<FRAME_SIZE;i++) tmp[i] = xn[i];
-    fwrite(tmp, sizeof(short), FRAME_SIZE, fout);
 #endif
   }
   fprintf(stderr, "matrix size: %d x %d\n", count, NB_FEATURES + 2*NB_BANDS + 1);
   fclose(f1);
   fclose(f2);
-  fclose(fout);
   return 0;
 }
 

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);
 }

src/rnn.h

@@ -27,6 +27,8 @@
 #ifndef RNN_H_
 #define RNN_H_
 
+#include "rnnoise-nu.h"
+
 #include "opus_types.h"
 
 #define WEIGHTS_SCALE (1.f/256)