Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 1 | #!/usr/bin/python |
| 2 | |
| 3 | from __future__ import print_function |
| 4 | |
| 5 | import keras |
| 6 | from keras.models import Sequential |
| 7 | from keras.models import Model |
| 8 | from keras.layers import Input |
| 9 | from keras.layers import Dense |
| 10 | from keras.layers import LSTM |
| 11 | from keras.layers import GRU |
| 12 | from keras.layers import SimpleRNN |
| 13 | from keras.layers import Dropout |
| 14 | from keras.layers import concatenate |
| 15 | from keras import losses |
| 16 | from keras import regularizers |
| 17 | import h5py |
| 18 | |
| 19 | from keras import backend as K |
| 20 | import numpy as np |
| 21 | |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 22 | #import tensorflow as tf |
| 23 | #from keras.backend.tensorflow_backend import set_session |
| 24 | #config = tf.ConfigProto() |
| 25 | #config.gpu_options.per_process_gpu_memory_fraction = 0.42 |
| 26 | #set_session(tf.Session(config=config)) |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 27 | |
| 28 | |
| 29 | def my_crossentropy(y_true, y_pred): |
| 30 | return K.mean(2*K.abs(y_true-0.5) * K.binary_crossentropy(y_pred, y_true), axis=-1) |
| 31 | |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 32 | def mymask(y_true): |
| 33 | return K.minimum(y_true+1., 1.) |
| 34 | |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 35 | def msse(y_true, y_pred): |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 36 | return K.mean(mymask(y_true) * K.square(K.sqrt(y_pred) - K.sqrt(y_true)), axis=-1) |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 37 | |
| 38 | def mycost(y_true, y_pred): |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 39 | return K.mean(mymask(y_true) * (K.square(K.sqrt(y_pred) - K.sqrt(y_true)) + 0.01*K.binary_crossentropy(y_pred, y_true)), axis=-1) |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 40 | |
| 41 | def my_accuracy(y_true, y_pred): |
| 42 | return K.mean(2*K.abs(y_true-0.5) * K.equal(y_true, K.round(y_pred)), axis=-1) |
| 43 | |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 44 | reg = 0.000001 |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 45 | |
| 46 | print('Build model...') |
| 47 | main_input = Input(shape=(None, 42), name='main_input') |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 48 | tmp = Dense(24, activation='tanh', name='input_dense')(main_input) |
| 49 | vad_gru = GRU(24, activation='tanh', recurrent_activation='sigmoid', return_sequences=True, name='vad_gru', kernel_regularizer=regularizers.l2(reg), recurrent_regularizer=regularizers.l2(reg))(tmp) |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 50 | vad_output = Dense(1, activation='sigmoid', name='vad_output')(vad_gru) |
| 51 | noise_input = keras.layers.concatenate([tmp, vad_gru, main_input]) |
| 52 | noise_gru = GRU(48, activation='relu', recurrent_activation='sigmoid', return_sequences=True, name='noise_gru', kernel_regularizer=regularizers.l2(reg), recurrent_regularizer=regularizers.l2(reg))(noise_input) |
| 53 | denoise_input = keras.layers.concatenate([vad_gru, noise_gru, main_input]) |
| 54 | |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 55 | denoise_gru = GRU(96, activation='tanh', recurrent_activation='sigmoid', return_sequences=True, name='denoise_gru', kernel_regularizer=regularizers.l2(reg), recurrent_regularizer=regularizers.l2(reg))(denoise_input) |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 56 | |
| 57 | denoise_output = Dense(22, activation='sigmoid', name='denoise_output')(denoise_gru) |
| 58 | |
| 59 | model = Model(inputs=main_input, outputs=[denoise_output, vad_output]) |
| 60 | |
| 61 | model.compile(loss=[mycost, my_crossentropy], |
| 62 | metrics=[msse], |
| 63 | optimizer='adam', loss_weights=[10, 0.5]) |
| 64 | |
| 65 | |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 66 | batch_size = 32 |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 67 | |
| 68 | print('Loading data...') |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 69 | with h5py.File('denoise_data6.h5', 'r') as hf: |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 70 | all_data = hf['data'][:] |
| 71 | print('done.') |
| 72 | |
| 73 | window_size = 2000 |
| 74 | |
| 75 | nb_sequences = len(all_data)//window_size |
| 76 | print(nb_sequences, ' sequences') |
| 77 | x_train = all_data[:nb_sequences*window_size, :42] |
| 78 | x_train = np.reshape(x_train, (nb_sequences, window_size, 42)) |
| 79 | |
| 80 | y_train = np.copy(all_data[:nb_sequences*window_size, 42:64]) |
| 81 | y_train = np.reshape(y_train, (nb_sequences, window_size, 22)) |
| 82 | |
| 83 | noise_train = np.copy(all_data[:nb_sequences*window_size, 64:86]) |
| 84 | noise_train = np.reshape(noise_train, (nb_sequences, window_size, 22)) |
| 85 | |
| 86 | vad_train = np.copy(all_data[:nb_sequences*window_size, 86:87]) |
| 87 | vad_train = np.reshape(vad_train, (nb_sequences, window_size, 1)) |
| 88 | |
| 89 | all_data = 0; |
| 90 | #x_train = x_train.astype('float32') |
| 91 | #y_train = y_train.astype('float32') |
| 92 | |
| 93 | print(len(x_train), 'train sequences. x shape =', x_train.shape, 'y shape = ', y_train.shape) |
| 94 | |
| 95 | print('Train...') |
| 96 | model.fit(x_train, [y_train, vad_train], |
| 97 | batch_size=batch_size, |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 98 | epochs=60, |
Jean-Marc Valin | cf473ce | 2017-08-03 15:26:05 -0400 | [diff] [blame] | 99 | validation_split=0.1) |
Jean-Marc Valin | 54eeea7 | 2017-08-08 11:20:29 -0400 | [diff] [blame^] | 100 | model.save("newweights6a2a.hdf5") |