In [ ]:
import copy as pcopy
from math import floor, sin
from random import random
import functools
import theano
from pylearn2.datasets.dataset import Dataset
from pylearn2.datasets.vector_spaces_dataset import VectorSpacesDataset
from pylearn2.sandbox.rnn.utils.iteration import SequenceDatasetIterator
from pylearn2.space import VectorSpace, CompositeSpace, IndexSpace
from pylearn2.sandbox.rnn.space import SequenceSpace, SequenceDataSpace, SequenceMaskSpace
from pylearn2.utils.iteration import resolve_iterator_class, FiniteDatasetIterator
from pylearn2.sandbox.rnn.models.rnn import RNN, Recurrent, LSTM
from pylearn2.models.mlp import Linear
from pylearn2.train import Train
from pylearn2.training_algorithms.sgd import SGD
from pylearn2.termination_criteria import EpochCounter
from pylearn2.costs.mlp import Default

figsize(40, 20)

# Sine Data Generation
n_train = 1000
len_min = 200
len_max = 300
step_min = 0.05*pi
step_max = 0.2*pi
n_target = 10

# RNN Parameters
n_hidden = 20
alpha = 0.01
batch_size = 50

# Visualization
n_samples = 10

In [ ]:
def sine_seq(n, len_min, len_max, step_min, step_max):
    r = []
    for i in range(n):
        l = int(len_min + floor(random() * (len_max - len_min)))
        start = random() * 2 * pi
        step = step_min + random() * (step_max - step_min)
        data = map(sin, linspace(start, start + l * step, l))
        r.append(data)
    return r

seqs = sine_seq(n_samples, len_min, len_max, step_min, step_max)

_, axes = subplots(n_samples, 1)

for i in range(n_samples):
    inp = seqs[i]
    axes[i].set_xlim([-1, len_max + 1])
    axes[i].set_ylim([-1.1, 1.1])
    axes[i].scatter(arange(len(inp)), inp)

In [ ]:
class SineSequences(VectorSpacesDataset):
    def __init__(self, n, len_min, len_max, len_target, step_min, step_max):
        self.num_examples = n
   
        space = CompositeSpace([
          SequenceDataSpace(VectorSpace(dim=1)),
          SequenceMaskSpace(),
          SequenceDataSpace(VectorSpace(dim=len_target))
        ])
        source = ('features', 'features_mask', 'targets')
        self.data_specs = (space, source)
        
        features = zeros((n, len_max, 1), dtype='float32')                
        features_mask = zeros((n, len_max, 1), dtype='float32')
        targets = zeros((n, len_max, len_target), dtype='float32')
        
        seqs = sine_seq(n, len_min, len_max, step_min, step_max)
        for i in range(len(seqs)):
            seq = seqs[i]
            features[i,:len(seq)-len_target,0] = seq[:-len_target]
            features_mask[i,:len(seq)-len_target,0] = 1
            for j in range(len(seq)-len_target):
                targets[i,j,:] = seq[j+1:j+1+len_target]
            
        self.data = (features, features_mask, targets)
        
    def get_data_specs(self):
        return self.data_specs
        
    def get_num_examples(self):
        return self.num_examples
    
    @functools.wraps(Dataset.iterator)
    def iterator(self, mode='shuffled_sequential', batch_size=None, num_batches=None, rng=None, data_specs=None, return_tuple=False):
        mode = resolve_iterator_class(mode)
        
        return SequenceDatasetIterator(self,
                                       data_specs,
                                       mode(self.num_examples, batch_size, num_batches, rng),
                                       return_tuple=return_tuple)

In [ ]:
def visualize(batch, train, net, epochs, axes):
    bkup = pcopy.copy(sys.stdout)
    f = open('/dev/null', 'w')
    sys.stdout = f
    
    idxs_color = linspace(0, 1, len(epochs))
    
    X = net.get_input_space().make_theano_batch()
    Y = net.fprop(X)
    f = theano.function(X, Y)
    
    examples = []
    for i in range(n_samples):
        example = asarray(batch[i][0]).flatten()
        example = example[example!=0]
        examples.append(example)        
        
        axes[i].set_xlim([-1, len_max+n_target])
        axes[i].set_ylim([-1.1, 1.1])
        axes[i].scatter(arange(len(example)), example)   

    for j,e in enumerate(epochs):
        train.algorithm.termination_criterion = EpochCounter(e)
        train.main_loop()
        
        for i in range(n_samples):
            y_pred = f(batch[i][0], batch[i][1])[0]
            axes[i].scatter(len(examples[i]) + arange(n_target), y_pred[len(examples[i])-1], color=cm.cool(idxs_color[j]))
            
    sys.stdout = bkup

In [ ]:
trainset = SineSequences(n_train, len_min, len_max, n_target, step_min, step_max)

rnn = RNN(
  input_space=SequenceSpace(VectorSpace(dim=1)),
  layers=[
    Recurrent(dim=n_hidden, layer_name='recurrent', irange=0.01),
    Linear(dim=n_target, layer_name='linear', irange=0.01)
  ]
)

nw_rnn = n_hidden + n_hidden**2 + n_hidden * n_target
print '# weights rnn =', nw_rnn

def nw_lstm(n):
    return n**2 * 4 + n * 4 + n * n_target + n * 3
nws_lstm = array(map(nw_lstm, range(20)))
idx = argmin(abs(nws_lstm-nw_rnn))
nw_lstm = nws_lstm[idx]
print '# memory units lstm =', idx, 'equiv', nw_lstm, 'weights'

# !!! Line 454 in rnn.py has to be fixed (mask was missing from params) !!! 
# PR is sent to pylearn2
lstm = RNN(
  input_space=SequenceSpace(VectorSpace(dim=1)),
  layers=[
    LSTM(dim=idx, layer_name='recurrent', irange=0.01),
    Linear(dim=n_target, layer_name='linear', irange=0.01)
  ]
)

alg = SGD(alpha, Default(), batch_size)
rnn_train = Train(trainset, rnn, alg)
lstm_train = Train(trainset, lstm, alg)

testset = SineSequences(n_samples, len_min, len_max, n_target, step_min, step_max)
it = testset.iterator(mode='shuffled_sequential', batch_size=1, data_specs=(rnn.get_input_space(), rnn.get_input_source()))
batch = [it.next() for i in range(n_samples)]

_, axes = subplots(n_samples, 2)
visualize(batch, rnn_train, rnn, [10,40,50], axes[:,0])
visualize(batch, lstm_train, lstm, [10,40,50], axes[:,1])