"""

RNN Speech Generation Model

Ishaan Gulrajani

"""

import os, sys

sys.path.append(os.getcwd())

try: # This only matters on Ishaan's computer

import gpu_queue

# gpu_queue.delay(60*60*3)

gpu_queue.wait_for_gpu(high_priority=False)

except ImportError:

pass

import numpy

numpy.random.seed(123)

import random

random.seed(123)

import dataset

import theano

import theano.tensor as T

import theano.ifelse

import lib

import lasagne

import scipy.io.wavfile

import time

import functools

import itertools

# Hyperparams

BATCH_SIZE = 128

N_FRAMES = 64 # How many 'frames' to include in each truncated BPTT pass

FRAME_SIZE = 4 # How many samples per frame

DIM = 512 # Model dimensionality. 512 is sufficient for model development; 1024 if you want good samples.

N_GRUS = 3 # How many GRUs to stack in the frame-level model

Q_LEVELS = 256 # How many levels to use when discretizing samples. e.g. 256 = 8-bit scalar quantization

GRAD_CLIP = 1 # Elementwise grad clip threshold

# Dataset

DATA_PATH = '/media/seagate/blizzard/parts'

N_FILES = 141703

BITRATE = 16000

# Other constants

TRAIN_MODE = 'iters' # 'iters' to use PRINT_ITERS and STOP_ITERS, 'time' to use PRINT_TIME and STOP_TIME

PRINT_ITERS = 10000 # Print cost, generate samples, save model checkpoint every N iterations.

STOP_ITERS = 100000 # Stop after this many iterations

PRINT_TIME = 60*60 # Print cost, generate samples, save model checkpoint every N seconds.

STOP_TIME = 60*60*12 # Stop after this many seconds of actual training (not including time req'd to generate samples etc.)

TEST_SET_SIZE = 128 # How many audio files to use for the test set

SEQ_LEN = N_FRAMES * FRAME_SIZE # Total length (# of samples) of each truncated BPTT sequence

Q_ZERO = numpy.int32(Q_LEVELS//2) # Discrete value correponding to zero amplitude

def frame_level_rnn(input_sequences, h0, reset):

return (output, last_hidden)

def sample_level_predictor(frame_level_outputs, prev_samples):

return lib.ops.Linear('SampleLevel.Output', DIM, Q_LEVELS, out)

sequences = T.imatrix('sequences')

h0 = T.tensor3('h0')

reset = T.iscalar('reset')

input_sequences = sequences[:, :-FRAME_SIZE]

target_sequences = sequences[:, FRAME_SIZE:]

frame_level_outputs, new_h0 = frame_level_rnn(input_sequences, h0, reset)

prev_samples = sequences[:, :-1]

prev_samples = prev_samples.reshape((1, BATCH_SIZE, 1, -1))

prev_samples = T.nnet.neighbours.images2neibs(prev_samples, (1, FRAME_SIZE), neib_step=(1, 1), mode='valid')

prev_samples = prev_samples.reshape((BATCH_SIZE * SEQ_LEN, FRAME_SIZE))

sample_level_outputs = sample_level_predictor(

frame_level_outputs.reshape((BATCH_SIZE * SEQ_LEN, DIM)),

prev_samples

)

cost = T.nnet.categorical_crossentropy(

T.nnet.softmax(sample_level_outputs),

target_sequences.flatten()

).mean()

# By default we report cross-entropy cost in bits.

# Switch to nats by commenting out this line:

cost = cost * lib.floatX(1.44269504089)

params = lib.search(cost, lambda x: hasattr(x, 'param'))

lib._train.print_params_info(cost, params)

grads = T.grad(cost, wrt=params, disconnected_inputs='warn')

grads = [T.clip(g, lib.floatX(-GRAD_CLIP), lib.floatX(GRAD_CLIP)) for g in grads]

updates = lasagne.updates.adam(grads, params)

train_fn = theano.function(

)

frame_level_generate_fn = theano.function(

)

frame_level_outputs = T.matrix('frame_level_outputs')

prev_samples = T.imatrix('prev_samples')

sample_level_generate_fn = theano.function(

)

def generate_and_save_samples(tag):

write_audio_file("sample_{}_{}".format(tag, i), samples[i])

print "Training!"

total_iters = 0

total_time = 0.

last_print_time = 0.

last_print_iters = 0

for epoch in itertools.count():

h0 = numpy.zeros((BATCH_SIZE, N_GRUS, DIM), dtype='float32')

costs = []

data_feeder = dataset.feed_epoch(DATA_PATH, N_FILES, BATCH_SIZE, SEQ_LEN, FRAME_SIZE, Q_LEVELS, Q_ZERO)

for seqs, reset in data_feeder:

start_time = time.time()

cost, h0 = train_fn(seqs, h0, reset)

total_time += time.time() - start_time

total_iters += 1

costs.append(cost)

if (TRAIN_MODE=='iters' and total_iters-last_print_iters == PRINT_ITERS) or \

(TRAIN_MODE=='time' and total_time-last_print_time >= PRINT_TIME):

print "epoch:{}\ttotal iters:{}\ttrain cost:{}\ttotal time:{}\ttime per iter:{}".format(

epoch,

total_iters,

numpy.mean(costs),

total_time,

total_time / total_iters

)

tag = "iters{}_time{}".format(total_iters, total_time)

generate_and_save_samples(tag)

lib.save_params('params_{}.pkl'.format(tag))

costs = []

last_print_time += PRINT_TIME

last_print_iters += PRINT_ITERS

if (TRAIN_MODE=='iters' and total_iters == STOP_ITERS) or \

(TRAIN_MODE=='time' and total_time >= STOP_TIME):

print "Done!"

sys.exit()