import os
import numpy as np
# For reproducibility
np.random.seed(42)
# force cuda device (empty for CPU)
os.environ["CUDA_VISIBLE_DEVICES"]=""
import sys
from keras import backend as K
# keras imports
from keras.models import Model
from keras.layers import Dense, Reshape, Input
from keras.layers.merge import add, concatenate, multiply
from keras.layers.core import Activation, Lambda
from keras.layers.recurrent import GRU
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.normalization import BatchNormalization
from keras.layers.convolutional import Convolution1D
from keras.optimizers import SGD, adam
import argparse
import math
# sklearn imports
from sklearn import preprocessing
from sklearn.externals import joblib
# netcdf for reading packaged data
from scipy.io import netcdf
import matplotlib as mpl
mpl.use('Agg') # no need for X-server
from matplotlib import pyplot as plt
from models import fft_model, time_glot_model, discriminator, generator, gan_container
from data_utils import nc_data_provider, norm_stats
# edge smoothing window
gen_filtwidths = np.asarray([15, 15, 15])
edgelen = sum(gen_filtwidths-1)
hannwin = np.hanning(edgelen)
smoothwin = np.concatenate((hannwin[:edgelen//2], np.ones(400-edgelen), hannwin[edgelen//2:]))
def plot_feats(generated_feats, epoch, index, ext='', fig_dir="./figures", fig_type=""):
plt.figure()
for row in generated_feats:
plt.plot(row)
plt.savefig(fig_dir + '/' + fig_type +'_epoch{}_index{}'.format(epoch, index) + ext + '.png')
plt.close()
def train_pls_model(BATCH_SIZE, data_dir, file_list, context_len=32, max_files=30):
no_epochs = 20
max_epochs_no_improvement = 5
timesteps = context_len
optim = adam(lr=0.0001)
pls_model = time_glot_model(timesteps=timesteps)
pls_model.compile(loss=['mse', 'mse'], loss_weights=[1.0, 0.0], optimizer=optim) # disregard fft loss
fft_mod = fft_model()
patience = max_epochs_no_improvement
best_val_loss = 1e20
for epoch in range(no_epochs):
print("Pre-train epoch is", epoch)
epoch_error = [0.0, 0.0]
total_batches = 0
val_data = []
for data in nc_data_provider(file_list, data_dir,
max_files=max_files, context_len=timesteps):
if len(val_data) == 0:
val_data = data
print("using data subset for validation")
continue
X_train = data[0]
Y_train = data[1]
no_batches = int(X_train.shape[0] / BATCH_SIZE)
print("Number of batches", int(X_train.shape[0] / BATCH_SIZE))
# shuffle data
ind = np.random.permutation(X_train.shape[0])
X_train = X_train[ind]
Y_train = Y_train[ind]
for index in range(int(X_train.shape[0] / BATCH_SIZE)):
x_feats_batch = X_train[
index * BATCH_SIZE:(index + 1) * BATCH_SIZE]
y_feats_batch = Y_train[
index * BATCH_SIZE:(index + 1) * BATCH_SIZE]
x_feats_batch_fft = fft_mod.predict(x_feats_batch)
d = pls_model.train_on_batch([y_feats_batch],
[x_feats_batch, x_feats_batch_fft])
epoch_error += d
if (index + total_batches) % 500 == 0:
print("pre-training batch %d, wave loss: %f, spec loss %f" %
(index+total_batches, d[0], d[1]))
wave, spec = pls_model.predict([y_feats_batch])
wav_gen = wave[0,:]
wav_ref = x_feats_batch[0,:]
wavs = np.array([wav_ref, wav_gen])
plot_feats(wavs, epoch, index+total_batches, fig_type='mse', ext='.wave-pls')
spec_gen = spec[0,:]
spec_ref = x_feats_batch_fft[0,:]
specs = np.array([spec_ref, spec_gen])
plot_feats(specs, epoch, index+total_batches, fig_type='mse', ext='.spec-pls')
total_batches += no_batches
epoch_error[0] /= total_batches
epoch_error[1] /= total_batches
val_spec = fft_mod.predict(val_data[0])
val_loss = pls_model.evaluate([val_data[1]],
[val_data[0], val_spec],
batch_size=BATCH_SIZE)
print("epoch %d validation wave loss: %f ,spec loss %f \n" %
(epoch, val_loss[0], val_loss[1]))
print("epoch %d training wave loss: %f, spec loss %f \n" %
(epoch, epoch_error[0], epoch_error[1]))
# only on wave loss
if val_loss[0] < best_val_loss:
best_val_loss = val_loss[0]
patience = max_epochs_no_improvement
pls_model.save_weights('./pls.model')
else:
patience -= 1
if patience == 0:
break
print ("Finished training")
def train_noise_model(BATCH_SIZE, data_dir, file_list, save_weights=False,
context_len=32, max_files=30, stats=None):
no_epochs = 15
timesteps = context_len
optim_container = adam(lr=1e-4)
optim_discriminator = SGD(lr=1e-5)
fft_mod = fft_model()
pls_model = time_glot_model(timesteps=timesteps)
pls_model.compile(loss=['mse','mse'], loss_weights=[1.0, 1.0], optimizer='adam')
pls_model.load_weights("./pls.model")
disc_model = discriminator()
gen_model = generator()
disc_on_gen = gan_container(gen_model, disc_model)
gen_model.compile(loss='mse', optimizer="adam")
# use peek adversarial and peek mse loss for training generator
disc_model.trainable = False
disc_on_gen.compile(loss=['mse','mse'], loss_weights=[1.0, 1.0], optimizer=optim_container)
# don't use peek loss for discriminator
disc_model.trainable = True
disc_model.compile(loss=['mse','mse'], loss_weights=[1.0, 0.0], optimizer=optim_discriminator)
print ("Discriminator model:")
print (disc_model.summary())
print ("Generator model:")
print (gen_model.summary())
print ("Joint model:")
print (disc_on_gen.summary())
label_fake = np.zeros((BATCH_SIZE, 1), dtype=np.float32)
label_real = np.ones((BATCH_SIZE, 1), dtype=np.float32)
# train residual GAN with FFT
for epoch in range(no_epochs):
print("Epoch is", epoch)
epoch_error = 0
total_batches = 0
for data in nc_data_provider(file_list, data_dir,
max_files=max_files, context_len=timesteps):
X_train = data[0]
Y_train = data[1]
pls_len = X_train.shape[1]
no_batches = int(X_train.shape[0] / BATCH_SIZE)
# shuffle data
ind = np.random.permutation(X_train.shape[0])
X_train = X_train[ind]
Y_train = Y_train[ind]
for index in range(int(X_train.shape[0] / BATCH_SIZE)):
x_feats_batch = X_train[
index * BATCH_SIZE:(index + 1) * BATCH_SIZE]
y_feats_batch = Y_train[
index * BATCH_SIZE:(index + 1) * BATCH_SIZE]
x_pred_batch, x_pred_batch_fft = pls_model.predict([y_feats_batch])
pls_pred = x_pred_batch
pls_real = x_feats_batch
# smoothing windows to prevent edge effects
pls_pred *= smoothwin
pls_real *= smoothwin
# evaluate target fft
fft_real = fft_mod.predict(pls_real)
noise = np.random.randn(BATCH_SIZE, pls_len)
# train generator through discriminator
_, peek_real = disc_model.predict([pls_real, fft_real])
disc_model.trainable = False
loss_g = disc_on_gen.train_on_batch([pls_pred, noise], [label_real, peek_real])
noise = np.random.randn(BATCH_SIZE, pls_len)
# train discriminator with real data
disc_model.trainable = True
loss_dr = disc_model.train_on_batch([pls_real, fft_real], [label_real, peek_real])
# train discriminator with fake data
pls_fake, fft_fake = gen_model.predict([pls_pred, noise])
loss_df = disc_model.train_on_batch([pls_fake, fft_fake], [label_fake, peek_real])
if (index + total_batches) % 500 == 0:
print("training batch %d, G loss: %f, D loss (real): %f, D loss (fake): %f" %
(index + total_batches, loss_g[0], loss_dr[0], loss_df[0]))
if (index + total_batches) % 500 == 0:
wav_ref = pls_real[0,:]
wav_gen = pls_pred[0,:]
wav_noised = pls_fake[0,:]
wavs = np.array([wav_ref, wav_gen, wav_noised])
plot_feats(wavs, epoch, index+total_batches, fig_type='gan', ext='.wave')
spec_ref = fft_real[0,:]
spec_noised = fft_fake[0,:]
specs = np.array([spec_ref, spec_noised])
plot_feats(specs, epoch, index+total_batches, fig_type='gan', ext='.spec')
total_batches += no_batches
gen_model.save_weights('./models/noise_gen_epoch' + str(epoch) + '.model')
print ("Finished noise model training")
def generate(file_list, data_dir, output_dir, context_len=32, stats=None,
base_model_path='./pls.model', gan_model_path='./noise_gen.model'):
pulse_model = time_glot_model(timesteps=context_len)
gan_model = generator()
pulse_model.compile(loss='mse', optimizer="adam")
gan_model.compile(loss='mse', optimizer="adam")
pulse_model.load_weights(base_model_path)
gan_model.load_weights(gan_model_path)
for data in nc_data_provider(file_list, data_dir, input_only=True,
context_len=context_len):
for fname, ac_data in data.iteritems():
print (fname)
pls_pred, _ = pulse_model.predict([ac_data])
noise = np.random.randn(pls_pred.shape[0], pls_pred.shape[1])
pls_gan, _ = gan_model.predict([pls_pred, noise])
out_file = os.path.join(args.output_dir, fname + '.pls')
pls_gan.astype(np.float32).tofile(out_file)
out_file = os.path.join(args.output_dir, fname + '.pls_nonoise')
pls_pred.astype(np.float32).tofile(out_file)
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--mode", type=str)
parser.add_argument("--batch_size", type=int, default=128)
parser.add_argument("--data_dir", type=str,
default="./traindata")
parser.add_argument("--testdata_dir", type=str,
default="./testdata")
parser.add_argument("--output_dir", type=str,
default="./output")
parser.add_argument("--rnn_context_len", type=int, default=64)
parser.add_argument("--max_files", type=int, default=100)
parser.set_defaults(nice=False)
parser.add_argument("--gan_model", type=str,
default=None)
args = parser.parse_args()
return args
if __name__ == "__main__":
args = get_args()
if args.mode == "train":
file_list = os.listdir(args.data_dir)
train_pls_model(BATCH_SIZE=args.batch_size, data_dir=args.data_dir,
file_list=file_list, max_files=args.max_files,
context_len=args.rnn_context_len)
stats = norm_stats(file_list[0], args.data_dir)
train_noise_model(BATCH_SIZE=args.batch_size, data_dir=args.data_dir,
file_list=file_list, max_files=args.max_files,
context_len=args.rnn_context_len,
stats=stats)
elif args.mode == "train_pulse_model":
print ("MODE: Training time domain pulse model")
file_list = os.listdir(args.data_dir)
train_pls_model(BATCH_SIZE=args.batch_size, data_dir=args.data_dir,
file_list=file_list, max_files=args.max_files,
context_len=args.rnn_context_len)
elif args.mode == "train_noise_model":
print ("MODE: Training noise model")
file_list = os.listdir(args.data_dir)
stats = norm_stats(file_list[0], args.data_dir)
train_noise_model(BATCH_SIZE=args.batch_size, data_dir=args.data_dir,
file_list=file_list, max_files=args.max_files,
context_len=args.rnn_context_len,
stats=stats)
elif args.mode == "generate":
test_dir = args.testdata_dir
file_list = os.listdir(test_dir)
stats = norm_stats(file_list[0], test_dir)
generate(data_dir=test_dir, file_list=file_list,
output_dir=args.output_dir,
context_len=args.rnn_context_len, stats=stats,
gan_model_path=args.gan_model)
