# Spectrogram auto-encoder

# Dan Stowell 2016.

#

# Unusual things about this implementation:

# * Data is not pre-whitened, instead we use a custom layer (NormalisationLayer) to normalise the mean-and-variance of the data for us. This is because I want the spectrogram to be normalised when it is input but not normalised when it is output.

# * It's a convolutional net but only along the time axis; along the frequency axis it's fully-connected.

import numpy as np

import theano

import theano.tensor as T

import lasagne

#import downhill

from lasagne.nonlinearities import rectify, leaky_rectify, very_leaky_rectify

from numpy import float32

try:

from lasagne.layers import InverseLayer as _

use_maxpool = True

except ImportError:

print("""**********************

WARNING: InverseLayer not found in Lasagne. Please use a more recent version of Lasagne.

WARNING: We'll deactivate the maxpooling part of the network (since we can't use InverseLayer to undo it)""")

use_maxpool = False

import matplotlib

#matplotlib.use('PDF') # http://www.astrobetter.com/plotting-to-a-file-in-python/

import matplotlib.pyplot as plt

import matplotlib.cm as cm

from matplotlib.backends.backend_pdf import PdfPages

plt.rcParams.update({'font.size': 6})

from userconfig import *

import util

from layers_custom import *

###################################################################################################################

# create Theano variables for input minibatch

input_var = T.tensor4('X')

# note that in general, the main data tensors will have these axes:

# - minibatchsize

# - numchannels (always 1 for us, since spectrograms)

# - numfilts (or specbinnum for input)

# - numtimebins

if example_is_audio:

# load our example audio file as a specgram

foregroundgram = util.standard_specgram((util.load_soundfile(foregroundwavpath, 0)))

backgroundgram = util.standard_specgram((util.load_soundfile(backgroundwavpath, 0)))

print("foregroundgram is of shape %s" % str(np.shape(foregroundgram)))

print("backgroundgram is of shape %s" % str(np.shape(backgroundgram)))

###################################################################################################################

# here we define our "semi-convolutional" autoencoder

# NOTE: lasagne assumes pooling is on the TRAILING axis of the tensor, so we always use time as the trailing axis

def make_custom_convlayer(network, in_num_chans, out_num_chans):

return network, filters

network = lasagne.layers.InputLayer((None, 1, specbinnum, numtimebins), input_var)

print("shape after input layer: %s" % str(network.output_shape))

#

# normalisation layer

# -- note that we deliberately normalise the input but do not undo that at the output.

# -- note that the normalisation params are not set by the training procedure, they need to be set before training begins.

network = NormalisationLayer(network, specbinnum)

normlayer = network # we need to remember this one so we can set its parameters

#

network, filters_enc = make_custom_convlayer(network, in_num_chans=specbinnum, out_num_chans=numfilters)

#

# NOTE: here we're using max-pooling, along the time axis only, and then

# using Lasagne's "InverseLayer" to undo the maxpooling in one-hot fashion.

# There's a side-effect of this: if you use *overlapping* maxpooling windows,

# the InverseLayer may behave slightly unexpectedly, adding some points with

# double magnitude. It's OK here since we're not overlapping the windows

if use_maxpool:

network = lasagne.layers.MaxPool2DLayer(network, pool_size=(1,2), stride=(1,2))

maxpool_layer = network # store a pointer to this one

# NOTE: HERE is the "middle" of the autoencoder!

latents = network # we remember the "latents" at the midpoint of the net, since we'll want to inspect them, and maybe regularise them too

if use_maxpool:

network = lasagne.layers.InverseLayer(network, maxpool_layer)

network, filters_dec = make_custom_convlayer(network, in_num_chans=numfilters, out_num_chans=specbinnum)

network = lasagne.layers.NonlinearityLayer(network, nonlinearity=rectify) # finally a standard rectify since nonneg (specgram) target

###################################################################################################################

# define simple L2 loss function with a mild touch of regularisation

prediction = lasagne.layers.get_output(network)

loss = T.mean(lasagne.objectives.squared_error(prediction, input_var))

# loss = loss.mean() + 1e-4 * lasagne.regularization.regularize_network_params(network, lasagne.regularization.l2)

lamb = 0.75

# loss =

###################################################################################################################

plot_probedata_data = None

def plot_probedata(outpostfix, plottitle=None):

pdf.close()

plot_probedata('init')

###################################################################################################################

if True:

###################################

# here we set up some training data. this is ALL A BIT SIMPLE - for a proper experiment we'd prepare a full dataset, and it might be too big to be all in memory.

training_data_size=100

training_data = np.zeros((training_data_size, minibatchsize, 1, specbinnum, numtimebins), dtype=float32)

if example_is_audio:

# manually grab a load of random subsets of the training audio

training_data_size=100

for which_training_batch in range(training_data_size):

for which_training_datum in range(minibatchsize):

startindex = np.random.randint(foregroundgram.shape[1]-numtimebins)

training_data[which_training_batch, which_training_datum, :, :, :] = foregroundgram[:, startindex:startindex+numtimebins]

else:

# make some simple (sparse) data that we can train with

for which_training_batch in range(training_data_size):

for which_training_datum in range(minibatchsize):

for _ in range(5):

training_data[which_training_batch, which_training_datum, :, np.random.randint(specbinnum), np.random.randint(numtimebins)] = 1

###################################

# pre-training setup

# set the normalisation parameters manually as an estimate from the training data

normlayer.set_normalisation(training_data)

###################################

# training

# compile training function that updates parameters and returns training loss

params = lasagne.layers.get_all_params(network, trainable=True)

updates = lasagne.updates.nesterov_momentum(loss, params, learning_rate=0.01, momentum=0.9)

train_fn = theano.function([input_var], loss, updates=updates)

# train network

numepochs = 2048 # 3 # 100 # 5000

for epoch in range(numepochs):

loss = 0

for input_batch in training_data:

loss += train_fn(input_batch)

if epoch==0 or epoch==numepochs-1 or (2 ** int(np.log2(epoch)) == epoch):

lossreadout = loss / len(training_data)

infostring = "Epoch %d/%d: Loss %g" % (epoch, numepochs, lossreadout)

print(infostring)

plot_probedata('progress', plottitle="progress (%s)" % infostring)

plot_probedata('trained', plottitle="trained (%d epochs; Loss %g)" % (numepochs, lossreadout))