#!/usr/bin/env python

import theano

import numpy

from theano import tensor

from blocks.bricks import Linear, Tanh, Rectifier

from blocks.bricks.conv import Convolutional, MaxPooling

from blocks.bricks.lookup import LookupTable

from blocks.bricks.recurrent import SimpleRecurrent, LSTM

from blocks.initialization import IsotropicGaussian, Constant

from blocks.algorithms import (GradientDescent, Scale, AdaDelta, RemoveNotFinite, RMSProp, BasicMomentum,

StepClipping, CompositeRule, Momentum)

from blocks.graph import ComputationGraph

from blocks.model import Model

from blocks.main_loop import MainLoop

from blocks.filter import VariableFilter

from blocks.roles import WEIGHT, BIAS

from blocks.graph import ComputationGraph, apply_dropout, apply_noise

from blocks.extensions import ProgressBar

from blocks.extensions.monitoring import TrainingDataMonitoring, DataStreamMonitoring

from blocks.extensions import FinishAfter, Printing

from blocks.extras.extensions.plot import Plot

from ctc import CTC

from timit import setup_datastream

from edit_distance import batch_edit_distance

from ext_param_info import ParamInfo

# ==========================================================================================

# THE HYPERPARAMETERS

# ==========================================================================================

# Stop after this many epochs

n_epochs = 10000

# How often (number of batches) to print / plot

monitor_freq = 50

sort_batch_count = 50

batch_size = 100

# The convolutionnal layers. Parameters:

# nfilter the number of filters

# filter_size the size of the filters (number of timesteps)

# stride the stride on which to apply the filter (non-1 stride are not optimized, runs very slowly with current Theano)

# pool_stride the block size for max pooling

# normalize do we normalize the values before applying the activation function?

# activation a brick for the activation function

# dropout dropout applied after activation function

# skip do we introduce skip connections from previous layer(s) to next layer(s) ?

convs = [

]

# recurrent layers. Parameters:

# type type of the layer (simple, lstm, blstm)

# dim size of the state

# normalize do we normalize the values after the RNN ?

# dropout dropout after the RNN

# skip do we introduce skip connections from previous layer(s) to next layer(s) ?

recs = [

]

# do we normalize the activations just before the softmax layer ?

normalize_out = True

# regularization : noise on the weights

weight_noise = 0.01

# regularization : L2 penalization

l2_output_bias = 0.

l2_output_weight = 0.

l2_all_bias = 0.0

l2_all_weight = 0.

# number of phonemes in timit, a constant

num_output_classes = 61

# the step rule (uncomment your favorite choice)

step_rule = CompositeRule([AdaDelta(), RemoveNotFinite()])

#step_rule = CompositeRule([Momentum(learning_rate=0.00001, momentum=0.99), RemoveNotFinite()])

#step_rule = CompositeRule([Momentum(learning_rate=0.001, momentum=0.9), RemoveNotFinite()])

#step_rule = CompositeRule([AdaDelta(), Scale(0.01), RemoveNotFinite()])

#step_rule = CompositeRule([RMSProp(learning_rate=0.1, decay_rate=0.95),

# RemoveNotFinite()])

#step_rule = CompositeRule([RMSProp(learning_rate=0.0001, decay_rate=0.95),

# BasicMomentum(momentum=0.9),

# RemoveNotFinite()])

# How the weights are initialized

weights_init = IsotropicGaussian(0.01)

biases_init = Constant(0.001)

# ==========================================================================================

# THE MODEL

# ==========================================================================================

print('Building model ...')

# THEANO INPUT VARIABLES

inputt = tensor.matrix('input')

input_mask = tensor.matrix('input_mask')

y = tensor.lmatrix('output').T

y_mask = tensor.matrix('output_mask').T

y_len = y_mask.sum(axis=0)

L = y.shape[0]

B = y.shape[1]

# inputt : B x T

# input_mask : B x T

# y : L x B

# y_mask : L x B

# NORMALIZE THE INPUTS

inputt = inputt / (inputt**2).mean()

dropout_locs = []

# CONVOLUTION LAYERS

conv_in = inputt[:, None, :, None]

conv_in_channels = 1

conv_in_mask = input_mask

cb = []

for i, p in enumerate(convs):

# Convolution bricks

conv = Convolutional(filter_size=(p['filter_size'],1),

# step=(p['stride'],1),

num_filters=p['nfilter'],

num_channels=conv_in_channels,

batch_size=batch_size,

border_mode='valid',

tied_biases=True,

name='conv%d'%i)

cb.append(conv)

maxpool = MaxPooling(pooling_size=(p['pool_stride'], 1), name='mp%d'%i)

conv_out = conv.apply(conv_in)[:, :, ::p['stride'], :]

conv_out = maxpool.apply(conv_out)

if p['normalize']:

conv_out_mean = conv_out.mean(axis=2).mean(axis=0)

conv_out_var = ((conv_out - conv_out_mean[None, :, None, :])**2).mean(axis=2).mean(axis=0).sqrt()

conv_out = (conv_out - conv_out_mean[None, :, None, :]) / conv_out_var[None, :, None, :]

if p['activation'] is not None:

conv_out = p['activation'].apply(conv_out)

if p['dropout'] > 0:

b = [p['activation'] if p['activation'] is not None else conv]

dropout_locs.append((VariableFilter(bricks=b, name='output'), p['dropout']))

if p['skip'] is not None and len(p['skip'])>0:

maxpooladd = MaxPooling(pooling_size=(p['stride']*p['pool_stride'], 1), name='Mp%d'%i)

skip = []

if 'max' in p['skip']:

skip.append(maxpooladd.apply(conv_in)[:, :, :conv_out.shape[2], :])

if 'min' in p['skip']:

skip.append(maxpooladd.apply(-conv_in)[:, :, :conv_out.shape[2], :])

if 'subsample' in p['skip']:

skip.append(conv_in[:, :, ::(p['stride']*p['pool_stride']), :][:, :, :conv_out.shape[2], :])

conv_out = tensor.concatenate([conv_out] + skip, axis=1)

conv_out_channels = p['nfilter'] + len(p['skip']) * conv_in_channels

else:

conv_out_channels = p['nfilter']

conv_out_mask = conv_in_mask[:, ::(p['stride']*p['pool_stride'])][:, :conv_out.shape[2]]

conv_in = conv_out

conv_in_channels = conv_out_channels

conv_in_mask = conv_out_mask

# RECURRENT LAYERS

rec_mask = conv_out_mask.dimshuffle(1, 0)

rec_in = conv_out[:, :, :, 0].dimshuffle(2, 0, 1)

rec_in_dim = conv_out_channels

rb = []

for i, p in enumerate(recs):

# RNN bricks

if p['type'] == 'lstm':

pre_rec = Linear(input_dim=rec_in_dim, output_dim=4*p['dim'], name='rnn_linear%d'%i)

rec = LSTM(activation=Tanh(), dim=p['dim'], name="rnn%d"%i)

rb = rb + [pre_rec, rec]

rnn_in = pre_rec.apply(rec_in)

rec_out, _ = rec.apply(inputs=rnn_in, mask=rec_mask)

dropout_b = [rec]

rec_out_dim = p['dim']

elif p['type'] == 'simple':

pre_rec = Linear(input_dim=rec_in_dim, output_dim=p['dim'], name='rnn_linear%d'%i)

rec = SimpleRecurrent(activation=Tanh(), dim=p['dim'], name="rnn%d"%i)

rb = rb + [pre_rec, rec]

rnn_in = pre_rec.apply(rec_in)

rec_out = rec.apply(inputs=rnn_in, mask=rec_mask)

dropout_b = [rec]

rec_out_dim = p['dim']

elif p['type'] == 'blstm':

pre_frec = Linear(input_dim=rec_in_dim, output_dim=4*p['dim'], name='frnn_linear%d'%i)

pre_brec = Linear(input_dim=rec_in_dim, output_dim=4*p['dim'], name='brnn_linear%d'%i)

frec = LSTM(activation=Tanh(), dim=p['dim'], name="frnn%d"%i)

brec = LSTM(activation=Tanh(), dim=p['dim'], name="brnn%d"%i)

rb = rb + [pre_frec, pre_brec, frec, brec]

frnn_in = pre_frec.apply(rec_in)

frnn_out, _ = frec.apply(inputs=frnn_in, mask=rec_mask)

brnn_in = pre_brec.apply(rec_in)

brnn_out, _ = brec.apply(inputs=brnn_in, mask=rec_mask)

rec_out = tensor.concatenate([frnn_out, brnn_out], axis=2)

dropout_b = [frec, brec]

rec_out_dim = 2*p['dim']

else:

assert False

if p['normalize']:

rec_out_mean = rec_out.mean(axis=1).mean(axis=0)

rec_out_var = ((rec_out - rec_out_mean[None, None, :])**2).mean(axis=1).mean(axis=0).sqrt()

rec_out = (rec_out - rec_out_mean[None, None, :]) / rec_out_var[None, None, :]

if p['dropout'] > 0:

dropout_locs.append((VariableFilter(bricks=dropout_b, name='output'), p['dropout']))

if p['skip']:

rec_out = tensor.concatenate([rec_in, rec_out], axis=2)

rec_out_dim = rec_in_dim + rec_out_dim

rec_in = rec_out

rec_in_dim = rec_out_dim

# LINEAR FOR THE OUTPUT

rec_to_o = Linear(name='rec_to_o',

input_dim=rec_out_dim,

output_dim=num_output_classes + 1)

y_hat_pre = rec_to_o.apply(rec_out)

# y_hat_pre : T x B x C+1

if normalize_out:

y_hat_pre_mean = y_hat_pre.mean(axis=1).mean(axis=0)

y_hat_pre_var = ((y_hat_pre - y_hat_pre_mean[None, None, :])**2).mean(axis=1).mean(axis=0).sqrt()

y_hat_pre = (y_hat_pre - y_hat_pre_mean[None, None, :]) / y_hat_pre_var[None, None, :]

# y_hat : T x B x C+1

y_hat = tensor.nnet.softmax(

y_hat_pre.reshape((-1, num_output_classes + 1))

).reshape((y_hat_pre.shape[0], y_hat_pre.shape[1], -1))

y_hat.name = 'y_hat'

y_hat_mask = rec_mask

# CTC COST AND ERROR MEASURE

cost = CTC().apply_log_domain(y, y_hat, y_len, y_hat_mask).mean()

cost.name = 'CTC'

dl, dl_length = CTC().best_path_decoding(y_hat, y_hat_mask)

dl = dl[:L, :]

dl_length = tensor.minimum(dl_length, L)

edit_distances = batch_edit_distance(dl.T.astype('int32'), dl_length.astype('int32'),

y.T.astype('int32'), y_len.astype('int32'))

edit_distance = edit_distances.mean()

edit_distance.name = 'edit_distance'

errors_per_char = (edit_distances / y_len).mean()

errors_per_char.name = 'errors_per_char'

is_error = tensor.neq(dl, y) * tensor.lt(tensor.arange(L)[:,None], y_len[None,:])

is_error = tensor.switch(is_error.sum(axis=0), tensor.ones((B,)), tensor.neq(y_len, dl_length))

error_rate = is_error.mean()

error_rate.name = 'error_rate'

# REGULARIZATION

cg = ComputationGraph([cost, error_rate])

if weight_noise > 0:

noise_vars = VariableFilter(roles=[WEIGHT])(cg)

cg = apply_noise(cg, noise_vars, weight_noise)

for vfilter, p in dropout_locs:

cg = apply_dropout(cg, vfilter(cg), p)

[cost_reg, error_rate_reg] = cg.outputs

ctc_reg = cost_reg + 1e-24

ctc_reg.name = 'CTC'

if l2_output_bias > 0:

cost_reg += l2_output_bias * sum(x.norm(2) for x in VariableFilter(roles=[BIAS], bricks=[rec_to_o])(cg))

if l2_output_weight > 0:

cost_reg += l2_output_weight * sum(x.norm(2) for x in VariableFilter(roles=[WEIGHT], bricks=[rec_to_o])(cg))

if l2_all_bias > 0:

cost_reg += l2_all_bias * sum(x.norm(2) for x in VariableFilter(roles=[BIAS])(cg))

if l2_all_weight > 0:

cost_reg += l2_all_weight * sum(x.norm(2) for x in VariableFilter(roles=[WEIGHT])(cg))

cost_reg.name = 'cost'

# INITIALIZATION

for brick in [rec_to_o] + cb + rb:

brick.weights_init = weights_init

brick.biases_init = biases_init

brick.initialize()

# ==========================================================================================

# THE INFRASTRUCTURE

# ==========================================================================================

# SET UP THE DATASTREAM

print('Bulding DataStream ...')

ds, stream = setup_datastream('/home/lx.nobackup/datasets/timit/readable',

batch_size=batch_size,

sort_batch_count=sort_batch_count)

valid_ds, valid_stream = setup_datastream('/home/lx.nobackup/datasets/timit/readable',

batch_size=batch_size,

sort_batch_count=sort_batch_count,

valid=True)

# SET UP THE BLOCKS ALGORITHM WITH EXTENSIONS

print('Bulding training process...')

algorithm = GradientDescent(cost=cost_reg,

parameters=ComputationGraph(cost).parameters,

step_rule=step_rule)

monitor_cost = TrainingDataMonitoring([ctc_reg, cost_reg, error_rate_reg],

prefix="train",

every_n_batches=monitor_freq,

after_epoch=False)

monitor_valid = DataStreamMonitoring([cost, error_rate, edit_distance, errors_per_char],

data_stream=valid_stream,

prefix="valid",

after_epoch=True)

plot = Plot(document='CTC_timit_%s%s%s%s_%s'%

after_epoch=True)

model = Model(cost)

main_loop = MainLoop(data_stream=stream, algorithm=algorithm,

model=model)

# NOW WE FINALLY CAN TRAIN OUR MODEL

print('Starting training ...')

main_loop.run()

# vim: set sts=4 ts=4 sw=4 tw=0 et: