From 4c954d8dea23e87d4ed86ae220587babc9cdc7c1 Mon Sep 17 00:00:00 2001 From: David Date: Sun, 27 Dec 2020 15:52:38 +1030 Subject: [PATCH] refactoring into vqvae_util.py --- vq_vae_kmeans_conv1d.py | 109 ++++++------------------------------ vq_vae_kmeans_conv1d_out.py | 75 +++++-------------------- vqvae_util.py | 84 +++++++++++++++++++++++++++ 3 files changed, 113 insertions(+), 155 deletions(-) create mode 100644 vqvae_util.py diff --git a/vq_vae_kmeans_conv1d.py b/vq_vae_kmeans_conv1d.py index d6a0e28..149570e 100755 --- a/vq_vae_kmeans_conv1d.py +++ b/vq_vae_kmeans_conv1d.py @@ -22,6 +22,7 @@ logging.getLogger('tensorflow').setLevel(logging.FATAL) import tensorflow as tf from vqvae_twostage import * +from vqvae_util import * # Constants ------------------------------------------------- @@ -39,9 +40,9 @@ parser.add_argument('--nb_samples', type=int, default=1000000, help='Number of f parser.add_argument('--embedding_dim', type=int, default=16, help='dimension of embedding vectors') parser.add_argument('--num_embedding', type=int, default=2048, help='number of embedded vectors') parser.add_argument('--scale', type=float, default=0.125, help='apply this gain to features when read in') -parser.add_argument('--nnout', type=str, default="", help='Name of output NN we have trained') +parser.add_argument('--nnout', type=str, help='Name of output NN we have trained') parser.add_argument('--mean', action='store_true', help='Extract mean from each chunk') -parser.add_argument('--mean_thresh', type=float, default=0, help='Discard chunks with less than this mean threshold') +parser.add_argument('--mean_thresh', type=float, default=0.0, help='Discard chunks with less than this mean threshold') args = parser.parse_args() dim = args.embedding_dim nb_samples = args.nb_samples @@ -133,24 +134,15 @@ vqvae.get_layer('vq2').set_vq(vq_initial) history = vqvae.fit(train, train_target, batch_size=batch_size, epochs=args.epochs, validation_split=validation_split,callbacks=[CustomCallback()]) -with open(args.nnout, 'wb') as f: - np.save(f, vqvae.get_layer("conv1d_a").get_weights(), allow_pickle=True) - np.save(f, vqvae.get_layer("conv1d_b").get_weights(), allow_pickle=True) - np.save(f, vqvae.get_layer("vq1").get_vq(), allow_pickle=True) - np.save(f, vqvae.get_layer("vq2").get_vq(), allow_pickle=True) - np.save(f, vqvae.get_layer("conv1d_c").get_weights(), allow_pickle=True) - np.save(f, vqvae.get_layer("conv1d_d").get_weights(), allow_pickle=True) - np.save(f, vqvae.get_layer("conv1d_e").get_weights(), allow_pickle=True) -''' -with open('test.npy', 'rb') as f: - vqvae.get_layer("conv1d_a").set_weights(np.load(f, allow_pickle=True)) - vqvae.get_layer("conv1d_b").set_weights(np.load(f, allow_pickle=True)) - vqvae.get_layer("vq1").set_vq(np.load(f, allow_pickle=True)) - vqvae.get_layer("vq2").set_vq(np.load(f, allow_pickle=True)) - vqvae.get_layer("conv1d_c").set_weights(np.load(f, allow_pickle=True)) - vqvae.get_layer("conv1d_d").set_weights(np.load(f, allow_pickle=True)) - vqvae.get_layer("conv1d_e").set_weights(np.load(f, allow_pickle=True)) -''' +if args.nnout is not None: + with open(args.nnout, 'wb') as f: + np.save(f, vqvae.get_layer("conv1d_a").get_weights(), allow_pickle=True) + np.save(f, vqvae.get_layer("conv1d_b").get_weights(), allow_pickle=True) + np.save(f, vqvae.get_layer("vq1").get_vq(), allow_pickle=True) + np.save(f, vqvae.get_layer("vq2").get_vq(), allow_pickle=True) + np.save(f, vqvae.get_layer("conv1d_c").get_weights(), allow_pickle=True) + np.save(f, vqvae.get_layer("conv1d_d").get_weights(), allow_pickle=True) + np.save(f, vqvae.get_layer("conv1d_e").get_weights(), allow_pickle=True) vq_weights = vqvae.get_layer('vq1').get_vq() @@ -184,42 +176,23 @@ plt.show(block=False) # Calculate total mean square error and mse per frame -def calc_mse(train, train_est, nb_samples, nb_features, dec): - msepf = np.zeros(nb_samples-dec) - e1 = 0; n = 0 - for i in range(nb_samples-dec): - e = (train_est[i,:] - train[i,:])**2 - msepf[i] = np.mean(e) - e1 += np.sum(e); n += nb_features - mse = e1/n - return mse, msepf - print("mse",train_target.shape, train_est.shape) mse,msepf = calc_mse(train_target, train_est, nb_samples, nb_features, 1) print("mse: %4.2f dB*dB" % (mse)) +worst_fr = np.argsort(-msepf); +worst_mse = np.sort(-msepf); +print(worst_fr[:10], worst_mse[:10]); plt.figure(3) plt.plot(msepf) plt.title('Spectral Distortion dB*dB per frame') plt.show(block=False) -def reject_outliers(data, m=2): - return data[abs(data - np.mean(data)) < m * np.std(data)] plt.figure(4) plt.title('Histogram of Spectral Distortion dB*dB out to 2*sigma') plt.hist(reject_outliers(msepf), bins='fd') plt.show(block=False) -# Count how many times each vector is used -def vector_count(x, vq, dim, nb_vecs): - # VQ search outside of Keras Backend - flat_inputs = np.reshape(x, (-1, dim)) - distances = np.sum(flat_inputs**2, axis=1, keepdims=True) - 2* np.dot(flat_inputs, vq.T) + np.sum(vq.T ** 2, axis=0, keepdims=True) - encoding_indices = np.argmax(-distances, axis=1) - count = np.zeros(nb_vecs, dtype="int") - count[encoding_indices] += 1 - return count - count = np.zeros(args.num_embedding, dtype="int") for i in range(0, nb_samples, batch_size): count += vector_count(encoder_out[i:i+batch_size], vq_weights, dim, args.num_embedding) @@ -235,23 +208,6 @@ plt.hist(train_mean, bins='fd') plt.show(block=False) plt.title('Mean of each chunk') -# use PCA to plot encoder space and VQ in 2D ----------------------------------------- - -# https://towardsdatascience.com/principal-component-analysis-pca-from-scratch-in-python-7f3e2a540c51 -def find_pca(A): - # calculate the mean of each column - M = np.mean(A.T, axis=1) - # center columns by subtracting column means - C = A - M - # calculate covariance matrix of centered matrix - V = np.cov(C.T) - # eigendecomposition of covariance matrix - values, vectors = np.linalg.eig(V) - #print(vectors) - #print(values) - P = vectors.T.dot(C.T) - return P.T - fig,ax = plt.subplots() encoder_pca=find_pca(encoder_out) ax.hist2d(encoder_pca[:,0],encoder_pca[:,1], bins=(50,50)) @@ -264,37 +220,4 @@ print("Press any key to start VQ pager....") key = getch.getch() plt.close('all') -# VQ Pager - plot input/output spectra to sanity check - -nb_plots = 8 -fs = 100; -key = ' ' -while key != 'q': - frames=range(fs,fs+nb_plots) - nb_plotsy = np.floor(np.sqrt(nb_plots)); nb_plotsx=nb_plots/nb_plotsy; - plt.figure(8) - plt.clf() - plt.tight_layout() - plt.title('Rate K Amplitude Spectra') - for r in range(nb_plots): - plt.subplot(nb_plotsy,nb_plotsx,r+1) - f = frames[r]; - plt.plot(train_target[f,:],'g') - plt.plot(train_est[f,:],'r') - plt.ylim(0,80) - a_mse = np.mean((train_target[f,:]-train_est[f,:])**2) - t = "f: %d %3.1f" % (f, a_mse) - plt.title(t) - plt.show(block=False) - plt.pause(0.0001) - print("n-next b-back s-save_png q-quit", end='\r', flush=True); - key = getch.getch() - if key == 'n': - fs += nb_plots - if key == 'b': - fs -= nb_plots - if key == 's': - plt.savefig('vqvae_spectra.png') - -plt.close() - +vqvae_pager(8,0,train_target,train_est,worst_fr) diff --git a/vq_vae_kmeans_conv1d_out.py b/vq_vae_kmeans_conv1d_out.py index 6f6516a..1cf4796 100755 --- a/vq_vae_kmeans_conv1d_out.py +++ b/vq_vae_kmeans_conv1d_out.py @@ -7,7 +7,7 @@ $ ./vq_vae_kmeans_conv1d.py dev-clean-8k-K20.f32 --eband_K 20 --epochs 5 --scale 0.005 --nnout test.npy -> 11.48dB*dB $ sox -t .sw -r 8000 ~/Downloads/train_8k.sw -t .sw - trim 0 2.5 | ~/codec2/build_linux/src/c2sim - --rateK --rateKout test.f32 - $ ./vq_vae_kmeans_conv1d_out.py test.npy test.f32 test_out.f32 --eband_K 20 --scale 0.005 + $ ./vq_vae_kmeans_conv1d_out.py test.npy test.f32 --featurefile_out test_out.f32 --eband_K 20 --scale 0.005 -> 21.6 dB*dB $ sox -t .sw -c 1 -r 8000 ~/Downloads/train_8k.sw -t .sw - trim 0 2.5 | ~/codec2/build_linux/src/c2sim - --rateK --rateKin test_out.f32 -o test1.raw @@ -24,6 +24,7 @@ logging.getLogger('tensorflow').setLevel(logging.FATAL) import tensorflow as tf from vqvae_twostage import * +from vqvae_util import * # Constants ------------------------------------------------- @@ -36,7 +37,7 @@ nb_timesteps = 4 parser = argparse.ArgumentParser(description='Two stage VQ-VAE for rate K vectors') parser.add_argument('ampnn', help='model weights and VQs in .npy') parser.add_argument('featurefile', help='input f32 file of spectral mag vectors, each element is 10*log10(energy), i.e. dB') -parser.add_argument('featurefile_out', help='output f32 file of spectral mag vectors, each element is 10*log10(energy), i.e. dB') +parser.add_argument('--featurefile_out', help='output f32 file of spectral mag vectors, each element is 10*log10(energy), i.e. dB') parser.add_argument('--nb_samples', type=int, default=1000000, help='Number of frames to train on') parser.add_argument('--eband_K', type=int, default=14, help='Length of eband vector') parser.add_argument('--embedding_dim', type=int, default=16, help='dimension of embedding vectors') @@ -122,29 +123,23 @@ encoder_out = encoder_out.reshape(-1, dim) print("target_est", target_est.shape, nb_samples) # make output file the same size despite chunking -target_est = np.concatenate((target_est, np.zeros((nb_samples_file-nb_samples,nb_features)))) -print(target_est.shape, nb_samples_file) -target_est_out = target_est.astype(np.float32); -print(features.shape, target_est_out.shape, target_est_out.dtype); -target_est_out.tofile(args.featurefile_out) +if args.featurefile_out is not None: + target_est = np.concatenate((target_est, np.zeros((nb_samples_file-nb_samples,nb_features)))) + print(target_est.shape, nb_samples_file) + target_est_out = target_est.astype(np.float32); + print(features.shape, target_est_out.shape, target_est_out.dtype); + target_est_out.tofile(args.featurefile_out) # Plot training results ------------------------- # Calculate total mean square error and mse per frame -def calc_mse(train, train_est, nb_samples, nb_features, dec): - msepf = np.zeros(nb_samples-dec) - e1 = 0; n = 0 - for i in range(nb_samples-dec): - e = (train_est[i,:] - train[i,:])**2 - msepf[i] = np.mean(e) - e1 += np.sum(e); n += nb_features - mse = e1/n - return mse, msepf - print("mse",target.shape, target_est.shape) mse,msepf = calc_mse(target, target_est, nb_samples, nb_features, 1) print("mse: %4.2f dB*dB" % (mse)) +worst_fr = np.argsort(-msepf); +worst_mse = np.sort(-msepf); +print(worst_fr[:10], worst_mse[:10]); if args.plots == False: quit(); @@ -153,23 +148,11 @@ plt.plot(msepf) plt.title('Spectral Distortion dB*dB per frame') plt.show(block=False) -def reject_outliers(data, m=2): - return data[abs(data - np.mean(data)) < m * np.std(data)] plt.figure(2) plt.title('Histogram of Spectral Distortion dB*dB out to 2*sigma') plt.hist(reject_outliers(msepf), bins='fd') plt.show(block=False) -# Count how many times each vector is used -def vector_count(x, vq, dim, nb_vecs): - # VQ search outside of Keras Backend - flat_inputs = np.reshape(x, (-1, dim)) - distances = np.sum(flat_inputs**2, axis=1, keepdims=True) - 2* np.dot(flat_inputs, vq.T) + np.sum(vq.T ** 2, axis=0, keepdims=True) - encoding_indices = np.argmax(-distances, axis=1) - count = np.zeros(nb_vecs, dtype="int") - count[encoding_indices] += 1 - return count - count = np.zeros(args.num_embedding, dtype="int") for i in range(0, nb_samples, batch_size): count += vector_count(encoder_out[i:i+batch_size], vq_weights, dim, args.num_embedding) @@ -185,37 +168,5 @@ plt.hist(target_mean, bins='fd') plt.show(block=False) plt.title('Mean of each chunk') -# VQ Pager - plot input/output spectra to sanity check - -nb_plots = 8 -fs = 0; -key = ' ' -while key != 'q': - frames=range(fs,fs+nb_plots) - nb_plotsy = np.floor(np.sqrt(nb_plots)); nb_plotsx=nb_plots/nb_plotsy; - plt.figure(5) - plt.clf() - plt.tight_layout() - plt.title('Rate K Amplitude Spectra') - for r in range(nb_plots): - plt.subplot(nb_plotsy,nb_plotsx,r+1) - f = frames[r]; - plt.plot(target[f,:],'g') - plt.plot(target_est[f,:],'r') - plt.ylim(0,80) - a_mse = np.mean((target[f,:]-target_est[f,:])**2) - t = "f: %d %3.1f" % (f, a_mse) - plt.title(t) - plt.show(block=False) - plt.pause(0.0001) - print("n-next b-back s-save_png q-quit", end='\r', flush=True); - key = getch.getch() - if key == 'n': - fs += nb_plots - if key == 'b': - fs -= nb_plots - if key == 's': - plt.savefig('vqvae_spectra.png') - -plt.close() +vqvae_pager(5,0,target,target_est,worst_fr) diff --git a/vqvae_util.py b/vqvae_util.py new file mode 100644 index 0000000..d51c69a --- /dev/null +++ b/vqvae_util.py @@ -0,0 +1,84 @@ +''' + VQVAE utility functions +''' + +import numpy as np +from matplotlib import pyplot as plt +import getch + +# VQ Pager - plot input/output spectra to sanity check +def vqvae_pager(fig,fs,target,est,worst_fr): + nb_plots = 8 + w = 0; + key = ' ' + while key != 'q': + frames=range(fs,fs+nb_plots) + nb_plotsy = np.floor(np.sqrt(nb_plots)); nb_plotsx=nb_plots/nb_plotsy; + plt.figure(fig) + plt.clf() + plt.tight_layout() + plt.title('Rate K Amplitude Spectra') + for r in range(nb_plots): + plt.subplot(nb_plotsy,nb_plotsx,r+1) + f = frames[r]; + plt.plot(target[f,:],'g') + plt.plot(est[f,:],'r') + plt.ylim(0,80) + a_mse = np.mean((target[f,:]-est[f,:])**2) + t = "f: %d %3.1f" % (f, a_mse) + plt.title(t) + plt.show(block=False) + plt.pause(0.0001) + print("n-next b-back w-worst s-save_png q-quit", end='\r', flush=True); + key = getch.getch() + if key == 'n': + fs += nb_plots + if key == 'b': + fs -= nb_plots + if key == 's': + plt.savefig('vqvae_spectra.png') + if key == 's': + plt.savefig('vqvae_spectra.png') + if key == 'w': + fs = min(worst_fr[w],target.shape[0]-nb_plots) + w += 1; + plt.close() + +# Calculate total mean square error and mse per frame +def calc_mse(train, train_est, nb_samples, nb_features, dec): + msepf = np.zeros(nb_samples-dec) + e1 = 0; n = 0 + for i in range(nb_samples-dec): + e = (train_est[i,:] - train[i,:])**2 + msepf[i] = np.mean(e) + e1 += np.sum(e); n += nb_features + mse = e1/n + return mse, msepf + +def reject_outliers(data, m=2): + return data[abs(data - np.mean(data)) < m * np.std(data)] + +# Count how many times each vector is used +def vector_count(x, vq, dim, nb_vecs): + # VQ search outside of Keras Backend + flat_inputs = np.reshape(x, (-1, dim)) + distances = np.sum(flat_inputs**2, axis=1, keepdims=True) - 2* np.dot(flat_inputs, vq.T) + np.sum(vq.T ** 2, axis=0, keepdims=True) + encoding_indices = np.argmax(-distances, axis=1) + count = np.zeros(nb_vecs, dtype="int") + count[encoding_indices] += 1 + return count + +# use PCA to plot encoder space and VQ in 2D +# https://towardsdatascience.com/principal-component-analysis-pca-from-scratch-in-python-7f3e2a540c51 +def find_pca(A): + # calculate the mean of each column + M = np.mean(A.T, axis=1) + # center columns by subtracting column means + C = A - M + # calculate covariance matrix of centered matrix + V = np.cov(C.T) + # eigendecomposition of covariance matrix + values, vectors = np.linalg.eig(V) + P = vectors.T.dot(C.T) + return P.T +