#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os

import numpy as np
from numpy import linalg as LA
from theano import tensor as T
import theano
from theano.tensor.shared_randomstreams import RandomStreams

from deepy import NeuralClassifier
from deepy.utils import get_activation
from deepy.core.disconnected_grad import disconnected_grad
from deepy.utils.functions import FLOATX
from deepy.layers import NeuralLayer
from examples.attention_models.gaussian_sampler import SampleMultivariateGaussian

class AttentionLayer(NeuralLayer):

    def __init__(self, activation='relu', std=0.1, disable_reinforce=False, random_glimpse=False):
        self.disable_reinforce = disable_reinforce
        self.random_glimpse = random_glimpse
        self.gaussian_std = std
        super(AttentionLayer, self).__init__(10, activation)

    def init(self, config, vars, x, input_n, id="UNKNOWN"):
        self._config = config
        self._vars = vars
        self.input_n = input_n
        self.id = id
        self.x = x
        self._setup_params()
        self._setup_functions()
        self.connected = True

    def _glimpse_sensor(self, x_t, l_p):
        """
        Parameters:
            x_t - 28x28 image
            l_p - 2x1 focus vector
        Returns:
            4x12 matrix
        """
        # Turn l_p to the left-top point of rectangle
        l_p = l_p * 14 + 14 - 2
        l_p = T.cast(T.round(l_p), "int32")

        l_p = l_p * (l_p >= 0)
        l_p = l_p * (l_p < 24) + (l_p >= 24) * 23
        l_p2 = l_p - 2
        l_p2 = l_p2 * (l_p2 >= 0)
        l_p2 = l_p2 * (l_p2 < 20) + (l_p2 >= 20) * 19
        l_p3 = l_p - 6
        l_p3 = l_p3 * (l_p3 >= 0)
        l_p3 = l_p3 * (l_p3 < 16) + (l_p3 >= 16) * 15
        glimpse_1 = x_t[l_p[0]: l_p[0] + 4][:, l_p[1]: l_p[1] + 4]
        glimpse_2 = x_t[l_p2[0]: l_p2[0] + 8][:, l_p2[1]: l_p2[1] + 8]
        glimpse_2 = theano.tensor.signal.downsample.max_pool_2d(glimpse_2, (2,2))
        glimpse_3 = x_t[l_p3[0]: l_p3[0] + 16][:, l_p3[1]: l_p3[1] + 16]
        glimpse_3 = theano.tensor.signal.downsample.max_pool_2d(glimpse_3, (4,4))
        return T.concatenate([glimpse_1, glimpse_2, glimpse_3])

    def _refined_glimpse_sensor(self, x_t, l_p):
        """
        Parameters:
            x_t - 28x28 image
            l_p - 2x1 focus vector
        Returns:
            7*14 matrix
        """
        # Turn l_p to the left-top point of rectangle
        l_p = l_p * 14 + 14 - 4
        l_p = T.cast(T.round(l_p), "int32")

        l_p = l_p * (l_p >= 0)
        l_p = l_p * (l_p < 21) + (l_p >= 21) * 20
        glimpse_1 = x_t[l_p[0]: l_p[0] + 7][:, l_p[1]: l_p[1] + 7]
        # glimpse_2 = theano.tensor.signal.downsample.max_pool_2d(x_t, (4,4))
        # return T.concatenate([glimpse_1, glimpse_2])
        return glimpse_1

    def _multi_gaussian_pdf(self, vec, mean):
        norm2d_var = ((1.0 / T.sqrt((2*np.pi)**2 * self.cov_det_var)) *
                      T.exp(-0.5 * ((vec-mean).T.dot(self.cov_inv_var).dot(vec-mean))))
        return norm2d_var

    def _glimpse_network(self, x_t, l_p):
        """
        Parameters:
            x_t - 28x28 image
            l_p - 2x1 focus vector
        Returns:
            4x12 matrix
        """
        sensor_output = self._refined_glimpse_sensor(x_t, l_p)
        sensor_output = T.flatten(sensor_output)
        h_g = self._relu(T.dot(sensor_output, self.W_g0))
        h_l = self._relu(T.dot(l_p, self.W_g1))
        g = self._relu(T.dot(h_g, self.W_g2_hg) + T.dot(h_l, self.W_g2_hl))
        return g

    def _location_network(self, h_t):
        """
        Parameters:
            h_t - 256x1 vector
        Returns:
            2x1 focus vector
        """
        return T.dot(h_t, self.W_l)

    def _action_network(self, h_t):
        """
        Parameters:
            h_t - 256x1 vector
        Returns:
            10x1 vector
        """
        z = self._relu(T.dot(h_t, self.W_a) + self.B_a)
        return self._softmax(z)

    def _core_network(self, l_p, h_p, x_t):
        """
        Parameters:
            x_t - 28x28 image
            l_p - 2x1 focus vector
            h_p - 256x1 vector
        Returns:
            h_t, 256x1 vector
        """
        g_t = self._glimpse_network(x_t, l_p)
        h_t = self._tanh(T.dot(g_t, self.W_h_g) + T.dot(h_p, self.W_h) + self.B_h)
        l_t = self._location_network(h_t)

        if not self.disable_reinforce:
            sampled_l_t = self._sample_gaussian(l_t, self.cov)
            sampled_pdf = self._multi_gaussian_pdf(disconnected_grad(sampled_l_t), l_t)
            wl_grad = T.grad(T.log(sampled_pdf), self.W_l)
        else:
            sampled_l_t = l_t
            wl_grad = self.W_l

        if self.random_glimpse and self.disable_reinforce:
            sampled_l_t = self.srng.uniform((2,)) * 0.8

        a_t = self._action_network(h_t)

        return sampled_l_t, h_t, a_t, wl_grad

    def _output_func(self):
        self.x = self.x.reshape((28, 28))
        [l_ts, h_ts, a_ts, wl_grads], _ = theano.scan(fn=self._core_network,
                         outputs_info=[self.l0, self.h0, None, None],
                         non_sequences=[self.x],
                         n_steps=5)

        self.positions = l_ts
        self.last_decision = T.argmax(a_ts[-1])
        wl_grad = T.sum(wl_grads, axis=0) / wl_grads.shape[0]
        self.wl_grad = wl_grad
        return a_ts[-1].reshape((1,10))

    def _setup_functions(self):
        self._assistive_params = []
        self._relu = get_activation("tanh")
        self._tanh = get_activation("tanh")
        self._softmax = get_activation("softmax")
        self.output_func = self._output_func()

    def _setup_params(self):
        self.srng = RandomStreams(seed=234)
        self.large_cov = np.array([[0.06,0],[0,0.06]], dtype=FLOATX)
        self.small_cov = np.array([[self.gaussian_std,0],[0,self.gaussian_std]], dtype=FLOATX)
        self.cov = theano.shared(np.array(self.small_cov, dtype=FLOATX))
        self.cov_inv_var = theano.shared(np.array(LA.inv(self.small_cov), dtype=FLOATX))
        self.cov_det_var = theano.shared(np.array(LA.det(self.small_cov), dtype=FLOATX))
        self._sample_gaussian = SampleMultivariateGaussian()

        self.W_g0 = self.create_weight(7 * 7, 128, label="g0")
        self.W_g1 = self.create_weight(2, 128, label="g1")
        self.W_g2_hg = self.create_weight(128, 256, label="g2_hg")
        self.W_g2_hl = self.create_weight(128, 256, label="g2_hl")

        self.W_h_g = self.create_weight(256, 256, label="h_g")
        self.W_h = self.create_weight(256, 256, label="h")
        self.B_h = self.create_bias(256, label="h")
        self.h0 = self.create_vector(256, "h0")
        self.l0 = self.create_vector(2, "l0")
        self.l0.set_value(np.array([-1, -1], dtype=FLOATX))

        self.W_l = self.create_weight(256, 2, label="l")
        self.W_l.set_value(self.W_l.get_value() / 10)
        self.B_l = self.create_bias(2, label="l")
        self.W_a = self.create_weight(256, 10, label="a")
        self.B_a = self.create_bias(10, label="a")

        self.W = [self.W_g0, self.W_g1, self.W_g2_hg, self.W_g2_hl, self.W_h_g, self.W_h, self.W_a]
        self.B = [self.B_h, self.B_a]
        self.parameters = [self.W_l]

def get_network(model=None, std=0.005, disable_reinforce=False, random_glimpse=False):
    """
    Get baseline model.
    Parameters:
        model - model path
    Returns:
        network
    """
    network = NeuralClassifier(input_dim=28 * 28)
    network.stack_layer(AttentionLayer(std=std, disable_reinforce=disable_reinforce, random_glimpse=random_glimpse))
    if model and os.path.exists(model):
        network.load_params(model)
    return network


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os

import numpy as np
from numpy import linalg as LA
from theano import tensor as T
import theano
from theano.tensor.shared_randomstreams import RandomStreams

from deepy import NeuralClassifier
from deepy.utils import get_activation
from deepy.core.disconnected_grad import disconnected_grad
from deepy.utils.functions import FLOATX
from deepy.layers import NeuralLayer
from examples.attention_models.gaussian_sampler import SampleMultivariateGaussian
import theano.tensor.signal.downsample

class FirstGlimpseLayer(NeuralLayer):

    def __init__(self, activation='tanh', std=0.1, disable_reinforce=False, random_glimpse=False):
        self.disable_reinforce = disable_reinforce
        self.random_glimpse = random_glimpse
        self.gaussian_std = std
        super(FirstGlimpseLayer, self).__init__(10, activation)

    def init(self, config, vars, x, input_n, id="UNKNOWN"):
        self._config = config
        self._vars = vars
        self.input_n = input_n
        self.id = id
        self.x = x
        self._setup_params()
        self._setup_functions()
        self.connected = True

    def _first_glimpse_sensor(self, x_t):
        """
        Compute first glimpse position using down-sampled image.
        """
        downsampled_img = theano.tensor.signal.downsample.max_pool_2d(x_t, (4,4))
        downsampled_img = downsampled_img.flatten()
        first_l = T.dot(downsampled_img, self.W_f)
        if self.disable_reinforce:
            wf_grad = self.W_f
            if self.random_glimpse:
                first_l = self.srng.uniform((2,), low=-1.7, high=1.7)
        else:
            sampled_l_t = self._sample_gaussian(first_l, self.cov)
            sampled_pdf = self._multi_gaussian_pdf(disconnected_grad(sampled_l_t), first_l)
            wf_grad = T.grad(T.log(sampled_pdf), self.W_f)
            first_l = sampled_l_t
        return first_l, wf_grad

    def _refined_glimpse_sensor(self, x_t, l_p):
        """
        Parameters:
            x_t - 28x28 image
            l_p - 2x1 focus vector
        Returns:
            7*14 matrix
        """
        # Turn l_p to the left-top point of rectangle
        l_p = l_p * 6.67 + 14 - 4
        l_p = T.cast(T.round(l_p), "int32")

        l_p = l_p * (l_p >= 0)
        l_p = l_p * (l_p < 21) + (l_p >= 21) * 20
        glimpse_1 = x_t[l_p[0]: l_p[0] + 7][:, l_p[1]: l_p[1] + 7]
        return glimpse_1

    def _multi_gaussian_pdf(self, vec, mean):
        norm2d_var = ((1.0 / T.sqrt((2*np.pi)**2 * self.cov_det_var)) *
                      T.exp(-0.5 * ((vec-mean).T.dot(self.cov_inv_var).dot(vec-mean))))
        return norm2d_var

    def _glimpse_network(self, x_t, l_p):
        """
        """
        sensor_output = self._refined_glimpse_sensor(x_t, l_p)
        sensor_output = T.flatten(sensor_output)
        h_g = self._relu(T.dot(sensor_output, self.W_g0))
        h_l = self._relu(T.dot(l_p, self.W_g1))
        g = self._relu(T.dot(h_g, self.W_g2_hg) + T.dot(h_l, self.W_g2_hl))
        return g

    def _location_network(self, h_t):
        """
        Parameters:
            h_t - 256x1 vector
        Returns:
            2x1 focus vector
        """
        return T.dot(h_t, self.W_l)

    def _action_network(self, h_t):
        """
        Parameters:
            h_t - 256x1 vector
        Returns:
            10x1 vector
        """
        z = self._relu(T.dot(h_t, self.W_a) + self.B_a)
        return self._softmax(z)

    def _core_network(self, l_p, h_p, x_t):
        """
        Parameters:
            x_t - 28x28 image
            l_p - 2x1 focus vector
            h_p - 256x1 vector
        Returns:
            h_t, 256x1 vector
        """
        g_t = self._glimpse_network(x_t, l_p)
        h_t = self._tanh(T.dot(g_t, self.W_h_g) + T.dot(h_p, self.W_h) + self.B_h)
        l_t = self._location_network(h_t)

        if not self.disable_reinforce:
            sampled_l_t = self._sample_gaussian(l_t, self.cov)
            sampled_pdf = self._multi_gaussian_pdf(disconnected_grad(sampled_l_t), l_t)
            wl_grad = T.grad(T.log(sampled_pdf), self.W_l)
        else:
            sampled_l_t = l_t
            wl_grad = self.W_l

        if self.random_glimpse and self.disable_reinforce:
            sampled_l_t = self.srng.uniform((2,), low=-1.7, high=1.7)

        a_t = self._action_network(h_t)

        return sampled_l_t, h_t, a_t, wl_grad

    def _output_func(self):
        self.x = self.x.reshape((28, 28))
        first_l, wf_grad = self._first_glimpse_sensor(self.x)

        [l_ts, h_ts, a_ts, wl_grads], _ = theano.scan(fn=self._core_network,
                         outputs_info=[first_l, self.h0, None, None],
                         non_sequences=[self.x],
                         n_steps=5)

        self.positions = l_ts
        self.last_decision = T.argmax(a_ts[-1])
        wl_grad = T.sum(wl_grads, axis=0) / wl_grads.shape[0]
        self.wl_grad = wl_grad
        self.wf_grad = wf_grad
        return a_ts[-1].reshape((1,10))

    def _setup_functions(self):
        self._relu = get_activation("tanh")
        self._tanh = get_activation("tanh")
        self._softmax = get_activation("softmax")
        self.output_func = self._output_func()

    def _setup_params(self):
        self.srng = RandomStreams(seed=234)
        self.large_cov = np.array([[0.06,0],[0,0.06]], dtype=FLOATX)
        self.small_cov = np.array([[self.gaussian_std,0],[0,self.gaussian_std]], dtype=FLOATX)
        self.cov = theano.shared(np.array(self.small_cov, dtype=FLOATX))
        self.cov_inv_var = theano.shared(np.array(LA.inv(self.small_cov), dtype=FLOATX))
        self.cov_det_var = theano.shared(np.array(LA.det(self.small_cov), dtype=FLOATX))
        self._sample_gaussian = SampleMultivariateGaussian()

        self.W_g0 = self.create_weight(7 * 7, 128, label="g0")
        self.W_g1 = self.create_weight(2, 128, label="g1")
        self.W_g2_hg = self.create_weight(128, 256, label="g2_hg")
        self.W_g2_hl = self.create_weight(128, 256, label="g2_hl")

        self.W_h_g = self.create_weight(256, 256, label="h_g")
        self.W_h = self.create_weight(256, 256, label="h")
        self.B_h = self.create_bias(256, label="h")
        self.h0 = self.create_vector(256, "h0")
        self.l0 = self.create_vector(2, "l0")
        self.l0.set_value(np.array([-1, -1], dtype=FLOATX))

        self.W_l = self.create_weight(256, 2, label="l")
        self.W_l.set_value(self.W_l.get_value() / 10)
        self.B_l = self.create_bias(2, label="l")
        self.W_a = self.create_weight(256, 10, label="a")
        self.B_a = self.create_bias(10, label="a")

        self.W_f = self.create_weight(7 * 7, 2, label="f")

        self.W = [self.W_g0, self.W_g1, self.W_g2_hg, self.W_g2_hl, self.W_h_g, self.W_h, self.W_a]
        self.B = [self.B_h, self.B_a]
        self.parameters = [self.W_l, self.W_f]

def get_network(model=None, std=0.005, disable_reinforce=False, random_glimpse=False):
    """
    Get baseline model.
    Parameters:
        model - model path
    Returns:
        network
    """
    network = NeuralClassifier(input_dim=28 * 28)
    network.stack_layer(FirstGlimpseLayer(std=std, disable_reinforce=disable_reinforce, random_glimpse=random_glimpse))
    if model and os.path.exists(model):
        network.load_params(model)
    return network


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging, os
logging.basicConfig(level=logging.INFO)
import theano.tensor as T

from deepy.dataset import MnistDataset, MiniBatches
from deepy.core.comp_graph import ComputationalGraph
from deepy.layers import Dense, Softmax, NeuralLayer, Chain
from deepy.trainers import MomentumTrainer, LearningRateAnnealer
from deepy.core import AutoEncoderCost, CrossEntropyCost, ErrorRateCost

model_path = os.path.join(os.path.dirname(__file__), "models", "tutorial2.gz")

class MyJointTrainingModel(NeuralLayer):
    """
    A customized model that trains an auto-encoder and MLP classifier simultaneously.
    """

    def __init__(self, internal_layer_size=100):
        super(MyJointTrainingModel, self).__init__("my joint-training model")
        self.internal_layer_size = internal_layer_size

    def prepare(self):
        """
        All codes that create parameters should be put into 'setup' function.
        """
        self.output_dim = 10
        self.encoder = Chain(self.input_dim).stack(Dense(self.internal_layer_size, 'tanh'))
        self.decoder = Chain(self.internal_layer_size).stack(Dense(self.input_dim))
        self.classifier = Chain(self.internal_layer_size).stack(Dense(50, 'tanh'),
                                                      Dense(self.output_dim),
                                                      Softmax())

        self.register_inner_layers(self.encoder, self.decoder, self.classifier)

        self.target_input = T.ivector('target')
        self.register_external_inputs(self.target_input)

    def compute_tensor(self, x):
        """
        Build the computation graph here.
        """
        internal_variable = self.encoder.compute_tensor(x)

        decoding_output = self.decoder.compute_tensor(internal_variable)

        classification_output = self.classifier.compute_tensor(internal_variable)

        auto_encoder_cost = AutoEncoderCost(decoding_output, x).get()

        classification_cost = CrossEntropyCost(classification_output, self.target_input).get()

        final_cost = 0.01 * auto_encoder_cost + classification_cost

        error_rate = ErrorRateCost(classification_output, self.target_input).get()

        self.register_monitors(("err", error_rate),
                               ("encoder_cost", auto_encoder_cost),
                               ("classify_cost", classification_cost))

        return final_cost

if __name__ == '__main__':
    model = graph.compile(input_dim=28 * 28)

    mnist = MiniBatches(MnistDataset(), batch_size=20)

    trainer = MomentumTrainer(model, {'weight_l2': 0.0001})

    trainer.run(mnist, epoch_controllers=[LearningRateAnnealer()])

    model.save_params(model_path)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from . import Dataset
from deepy.core.env import FLOATX
import numpy as np

import logging as loggers
logging = loggers.getLogger(__name__)

class BasicDataset(Dataset):
    """
    Basic dataset class.
    """

    def __init__(self, train, valid=None, test=None):
        self._train_set = train
        self._valid_set = valid
        self._test_set = test

    def train_set(self):
        return self._train_set

    def valid_set(self):
        return self._valid_set

    def test_set(self):
        return self._test_set

    def map(self, func):
        """
        Process all data with given function.
        The scheme of function should be x,y -> x,y.
        """
        if self._train_set:
            self._train_set = map(func, self._train_set)
        if self._valid_set:
            self._valid_set = map(func, self._valid_set)
        if self._test_set:
            self._test_set = map(func, self._test_set)

    def _vectorize_set(self, subset, size):
        newset = []
        for x, y in subset:
            l = [0.] * size
            l[y] = 1.
            new_y = np.array(l, dtype=FLOATX)
            newset.append((x, new_y))
        return newset

    def vectorize_target(self, size):
        """
        Make targets be one-hot vectors.
        """
        if self._train_set:
            self._train_set = self._vectorize_set(self._train_set, size)
        if self._valid_set:
            self._valid_set = self._vectorize_set(self._valid_set, size)
        if self._test_set:
            self._test_set = self._vectorize_set(self._test_set, size)

    def report(self):
        """
        Print dataset statistics.
        """
        logging.info("%s train=%d valid=%d test=%d" % (self.__class__.__name__,
                                                       len(list(self._train_set)) if self._train_set else 0,
                                                       len(list(self._valid_set)) if self._valid_set else 0,
                                                       len(list(self._test_set)) if self._test_set else 0))


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging as loggers
from abc import ABCMeta, abstractmethod
from deepy.trainers.base import NeuralTrainer

logging = loggers.getLogger(__name__)
THEANO_LINKER = 'cvm'

class CustomizeTrainer(NeuralTrainer):
    """
    DEPRECATED !!!
    A customized trainer.
    """
    __metaclass__ = ABCMeta

    def __init__(self, network, config=None):
        """
        Basic neural network trainer.
        :type network: deepy.NeuralNetwork
        :type config: deepy.conf.TrainerConfig
        :return:
        """
        super(CustomizeTrainer, self).__init__(network, config)

    def train(self, train_set, valid_set=None, test_set=None, train_size=None):
        '''We train over mini-batches and evaluate periodically.'''
        iteration = 0
        while True:
            if not iteration % self.config.test_frequency and test_set:
                try:
                    self.test(iteration, test_set)
                except KeyboardInterrupt:
                    logging.info('interrupted!')
                    break

            if not iteration % self.validation_frequency and valid_set:
                try:
                    if not self.evaluate(iteration, valid_set):
                        logging.info('patience elapsed, bailing out')
                        break
                except KeyboardInterrupt:
                    logging.info('interrupted!')
                    break

            train_message = ""
            try:
                train_message = self.train_func(train_set)
            except KeyboardInterrupt:
                logging.info('interrupted!')
                break
            if not iteration % self.config.monitor_frequency:
                logging.info('monitor (iter=%i) %s', iteration + 1, train_message)

            iteration += 1
            if hasattr(self.network, "iteration_callback"):
                self.network.iteration_callback()

            yield train_message

        if valid_set:
            self.set_params(self.best_params)
        if test_set:
            self.test(0, test_set)

    @abstractmethod
    def train_func(self, train_set):
        return ""


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import logging as loggers
logging = loggers.getLogger(__name__)

from deepy.networks import NeuralNetwork
from deepy.layers import OneHotEmbedding
from deepy.utils import EPSILON
from deepy.tensor.onehot import onehot
import theano.tensor as T

from cost import LMCost

class NeuralLM(NeuralNetwork):
    """
    LM Network.
    """

    def __init__(self, input_dim, input_tensor=None):
        self.class_based = class_based
        super(NeuralLM, self).__init__(0, input_tensor=T.imatrix('x'))
        self.stack(OneHotEmbedding(vocab_size))
        self.test_data = test_data
        if test_data:
            self.epoch_callbacks.append(self._report_ppl)

    def setup_variables(self):
        super(NeuralLM, self).setup_variables()
        if not self.class_based:
            self.k = T.imatrix('k')
            self.target_variables.append(self.k)

    def _cost_func(self, y):
        if self.class_based:
            return y
        else:
            return LMCost(y, self.k).get()

    def _error_func(self, y):
        y2 = y.reshape((-1, y.shape[-1]))
        k2 = self.k.reshape((-1,))
        return 100 * T.mean(T.neq(T.argmax(y2[:k2.shape[0]], axis=1), k2))

    def _perplexity_func(self, y):
        return 2**self._cost_func(y)

    @property
    def cost(self):
        return self._cost_func(self.output)

    @property
    def test_cost(self):
        return self._cost_func(self.test_output)

    def predict(self, x):
        return self.compute(x).argmax(axis=1)

    def sample(self, input, steps):
        """
        Sample outputs from LM.
        """
        inputs = [[onehot(self.input_dim, x) for x in input]]
        for _ in range(steps):
            target = self.compute(inputs)[0,-1].argmax()
            input.append(target)
            inputs[0].append(onehot(self.input_dim, target))
        return input

    def prepare_training(self):
        if not self.class_based:
            self.training_monitors.append(("err", self._error_func(self.output)))
            self.testing_monitors.append(("err", self._error_func(self.test_output)))
        self.training_monitors.append(("approx_PPL", self._perplexity_func(self.output)))
        self.testing_monitors.append(("approx_PPL", self._perplexity_func(self.test_output)))
        super(NeuralLM, self).prepare_training()

    def _report_ppl(self):
        if not self.test_data: return
        import numpy as np
        costs = []
        train_set = list(self.test_data)
        for i in range(len(train_set)):
            k2 =train_set[i][1].reshape((-1,))
            y2 = self.compute(train_set[i][0]).reshape((-1, 10001))
            c = -np.log2(y2[np.arange(k2.shape[0]), k2]).mean()
            costs.append(c)
        print np.array(costs).mean()
        # return
        # logp_sum  = 0
        # word_count = 0
        # for x_batch, y_batch in self.test_data:
        #     yhat_batch = self.compute(x_batch)
        #     for x_seq, y_seq, yhat_seq in zip(x_batch, y_batch, yhat_batch):
        #         for x, y, yhat in  zip(x_seq, y_seq, yhat_seq):
        #             logp = - np.log2(yhat[y])
        #             logp_sum += logp
        #             word_count += 1
        #             if y == 0:
        #                 # End of the sentence
        #                 break
        # ppl = 0
        # mean_logp = 0
        # if word_count > 0:
        #     mean_logp = logp_sum / word_count
        #     ppl = 2 ** mean_logp
        # logging.info("Accurate test PPL: %.2f, Entropy: %.2f" % (ppl, mean_logp))


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import theano
import theano.tensor as T
from deepy.layers import NeuralLayer, Softmax3D, Softmax, Dense, Chain
from deepy.core import CrossEntropyCost

from cost import LMCost

class FullOutputLayer(NeuralLayer):

    def __init__(self, vocab_size):
        super(FullOutputLayer, self).__init__("full_output")
        self.vocab_size = vocab_size

    def prepare(self):
        self.core = Chain(self.input_dim).stack(Dense(self.vocab_size),
                                                Softmax3D())
        self.register_inner_layers(self.core)

    def compute_tensor(self, x):
        return self.core.compute_tensor(x)

class ClassOutputLayer(NeuralLayer):

    def __init__(self, output_size, class_size):
        super(ClassOutputLayer, self).__init__("class_output")
        self.output_size = output_size
        self.class_size = class_size

    def prepare(self):
        # Output layers
        self.output_layer = Chain(self.input_dim).stack(Dense(self.output_size * self.class_size))
        self.softmax_layer = Softmax().init(input_dim=self.output_size)

        self.class_layer = Chain(self.input_dim).stack(Dense(self.class_size),
                                                        Softmax3D())
        self.register_inner_layers(self.class_layer, self.output_layer)
        # Target tensor
        self.target_tensor = T.imatrix('target')
        self.register_external_targets(self.target_tensor)
        # arange cache
        self.arange_cache = theano.shared(np.arange(10*64), name="arange_cache")

    def compute_tensor(self, x):
        """
        :param x: (batch, time, vec)
        """
        # Target class
        class_matrix = self.target_tensor // self.output_size
        class_vector = class_matrix.reshape((-1,))
        # Target index
        target_matrix = self.target_tensor % self.output_size
        target_vector = target_matrix.reshape((-1,))
        # Input matrix
        input_matrix = x.reshape((-1, self.input_dim))
        # Output matrix
        output_tensor3d = self.output_layer.compute_tensor(x)
        output_matrix = output_tensor3d.reshape((-1, self.class_size, self.output_size))
        arange_vec = self.arange_cache[:output_matrix.shape[0]]
        sub_output_matrix = output_matrix[arange_vec, class_vector]
        # Softmax
        softmax_output_matrix = self.softmax_layer.compute_tensor(sub_output_matrix)
        # Class prediction
        class_output_matrix = self.class_layer.compute_tensor(x)
        # Costs
        output_cost = LMCost(softmax_output_matrix, target_vector).get()
        class_cost = LMCost(class_output_matrix, class_matrix).get()
        final_cost = output_cost + class_cost

        return final_cost


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from layer import NeuralLayer
import deepy.tensor as T

class Attention(NeuralLayer):

    def __init__(self, hidden_size, input_dim):
        super(Attention, self).__init__("attention")
        self.input_dim = input_dim if input_dim else hidden_size
        self.hidden_size = hidden_size
        self.init(input_dim)

    def prepare(self):
        self.Ua = self.create_weight(self.input_dim, self.hidden_size, "ua")
        self.Wa = self.create_weight(self.hidden_size, self.hidden_size, "wa")
        self.Va = self.create_weight(label="va", shape=(self.hidden_size,))
        self.register_parameters(self.Va, self.Wa, self.Ua)

    def precompute(self, inputs):
        """
        Precompute partial values in the score function.
        """
        return T.dot(inputs, self.Ua)

    def compute_alignments(self, prev_state, precomputed_values, mask=None):
        """
        Compute the alignment weights based on the previous state.
        """

        WaSp = T.dot(prev_state, self.Wa)
        UaH = precomputed_values
        # For test time the UaH will be (time, output_dim)
        if UaH.ndim == 2:
            preact = WaSp[:, None, :] + UaH[None, :, :]
        else:
            preact = WaSp[:, None, :] + UaH
        act = T.activate(preact, 'tanh')
        align_scores = T.dot(act, self.Va)  # ~ (batch, time)
        if mask:
            mask = (1 - mask) * -99.00
            if align_scores.ndim == 3:
                align_scores += mask[None, :]
            else:
                align_scores += mask
        align_weights = T.nnet.softmax(align_scores)
        return align_weights

    def compute_context_vector(self, prev_state, inputs, precomputed_values=None, mask=None):
        """
        Compute the context vector with soft attention.
        """
        precomputed_values = precomputed_values if precomputed_values else self.precompute(inputs)
        align_weights = self.compute_alignments(prev_state, precomputed_values, mask)
        context_vector = T.sum(align_weights[:, :, None] * inputs, axis=1)
        return context_vector


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import time
import os
from collections import OrderedDict
import numpy as np

from deepy.trainers import GeneralNeuralTrainer

import logging

class MultiGPUTrainer(GeneralNeuralTrainer):
    """
    General neural network trainer.
    """

    def __init__(self,
                 network, config=None, method='sgd',
                 server_port=5567,
                 start_halving_at=6, end_at=10, sync_freq=3,
                 valid_freq=1500, learning_rate=None, halving_freq=1,
                 using_easgd=True
                 ):
        super(MultiGPUTrainer, self).__init__(network, method, config)
        self._report_time = False
        self._port = server_port
        self.logger = logging.getLogger('MultiGPUTrainingWorker')
        self.epoch = 0
        self._using_easgd = using_easgd
        if not learning_rate:
            learning_rate = float(self.config.learning_rate.get_value())
        self._schedule_params = {
            'learning_rate': learning_rate,
            'start_halving_at': start_halving_at,
            'end_at': end_at,
            'sync_freq': sync_freq,
            'valid_freq': valid_freq,
            'halving_freq': halving_freq
        }

    def create_param_map(self):
        param_map = OrderedDict()
        for i, param in enumerate(self.training_params()):
            param_map["param_{}".format(i)] = param
        return param_map

    def sync_hyperparams(self, param_map):
        self.logger.info("(proc {}) sync hyperparameters".format(os.getpid()))
        if 'epoch' in param_map:
            self.epoch = param_map['epoch']
        if 'learning_rate' in param_map:
            self.config.learning_rate.set_value(param_map['learning_rate'])

    def fix_costs(self):
        self.last_run_costs = [(a, float(b)) for (a,b) in self.last_run_costs]

    def train(self, train_set, valid_set=None, test_set=None, train_size=None):
        """
        Train the model in multi-GPU environment.
        """
        from platoon.channel import Worker
        from platoon.param_sync import EASGD, ASGD
        server_port = self._port
        param_map = self.create_param_map()
        # Initialize the worker
        worker = Worker(control_port=server_port)
        if self.config.learning_rate:
            worker.send_req({'init_schedule': self._schedule_params})
        self.sync_hyperparams(worker.send_req('sync_hyperparams')['sync_hyperparams'])
        easgd_alpha = worker.send_req('get_easgd_alpha')
        if self._using_easgd:
            self.logger.info("using EASGD with alpha={}".format(easgd_alpha))
        else:
            self.logger.info("using ASGD rule")
        rule = EASGD(easgd_alpha) if self._using_easgd else ASGD()
        worker.init_shared_params(param_map.values(), param_sync_rule=rule)
        worker.send_req({
            "set_names": None,
            "training_names": self.training_names,
            "evaluation_names": self.evaluation_names
        })
        # Load all training batches, consume vast memory here
        self.logger.info("started process {}".format(os.getpid()))
        self.logger.info("(proc {}) load training data".format(os.getpid()))
        train_batches = list(train_set)
        network_callback = bool(self.network.training_callbacks)
        trainer_callback = bool(self._iter_controllers)
        # Start from valid, so the performance when a worked join can be known
        worker.copy_to_local()
        if valid_set:
            self._run_valid(self.epoch, valid_set, dry_run=True)
            self.fix_costs()
        worker.send_req({
            "valid_done": None,
            "valid_costs": self.last_run_costs,
            "auto_save": self.config.auto_save
        })
        worker.copy_to_local()
        # Begin the loop
        while True:
            resp = worker.send_req('next')
            if resp == 'stop':
                break
            elif resp == 'wait':
                time.sleep(1)
            elif resp == 'get_num_batches':
                worker.send_req({'get_num_batches_done': len(train_batches)})
            elif 'eval' in resp:
                self.best_cost = resp['best_valid_cost']
                worker.copy_to_local()
                valid_costs = None
                test_costs = None
                if valid_set:
                    self._run_valid(self.epoch, valid_set)
                    self.fix_costs()
                    valid_costs = self.last_run_costs
                if test_set:
                    self._run_test(self.epoch, test_set)
                    self.fix_costs()
                    test_costs = self.last_run_costs
                worker.send_req({
                    "eval_done": None,
                    "valid_costs": valid_costs,
                    "test_costs": test_costs,
                    "auto_save": self.config.auto_save
                })
            elif 'valid' in resp:
                self.best_cost = resp['best_valid_cost']
                worker.copy_to_local()
                if valid_set:
                    self._run_valid(self.epoch, valid_set, dry_run=True)
                    self.fix_costs()
                worker.send_req({
                    "valid_done": None,
                    "valid_costs": self.last_run_costs,
                    "auto_save": self.config.auto_save
                })
            elif 'train' in resp:
                batch_ids = resp['train']
                batch_costs = [[] for _ in self.training_names]
                for batch_id in batch_ids:
                    x = train_batches[batch_id]
                    cost_x = self.learn(*x)
                    for i, cost in enumerate(cost_x):
                        batch_costs[i].append(cost)
                    self.last_cost = cost_x[0]
                if network_callback:
                    self.network.training_callback()
                if trainer_callback:
                    for func in self._iter_controllers:
                        func(self)
                worker.sync_params(synchronous=True)
                worker.send_req({'train_done': None, 'costs': [float(np.mean(c)) for c in batch_costs]})
            elif 'sync_hyperparams' in resp:
                self.sync_hyperparams(resp['sync_hyperparams'])
        worker.close()
        return []


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import theano.tensor as TT
from theano.ifelse import ifelse as theano_ifelse

from deepy.core.tensor_conversion import neural_computation
from deepy.layers.concat import Concatenate
from wrapper import deepy_tensor

def concat(vars, axis=-1):
    """
    A shortcut for concatenation.
    """
    return concatenate(vars, axis)

@neural_computation
def reverse(tensor, axis=-1):
    ndim = tensor.ndim
    selectors = [slice(None)] * ndim
    selectors[axis] = slice(None, None, -1)
    ret = tensor[selectors]
    if hasattr(tensor.tag, "last_dim"):
        ret.tag.last_dim = tensor.tag.last_dim
    return ret

@neural_computation
def activate(var, method):
    """
    An activation function.
    :param var: input var
    :param method: type of activation, such as `relu`,`tanh`,`sigmoid`
    """
    from activations import get_activation
    return get_activation(method)(var)

def concatenate(vars, axis=-1):
    """
    A utility function of concatenate.
    """
    from deepy.core.neural_var import NeuralVariable
    if isinstance(vars[0], NeuralVariable):
        concat_var = Concatenate(axis=axis).compute(*vars)
        if axis == -1 or axis == vars[0].tensor.ndim - 1:
            concat_var.output_dim = sum([x.output_dim for x in vars], 0)
    else:
        concat_var = TT.concatenate(vars, axis)
    return concat_var

@neural_computation
def ifelse(condition, then_branch, else_branch):
    return theano_ifelse(condition, then_branch, else_branch)

def apply(func, *args, **kwargs):
    from deepy.core.neural_var import NeuralVariable
    dim = kwargs['dim'] if 'dim' in kwargs else args[0].dim()
    return NeuralVariable(func(*[x.tensor for x in args]), dim)

def repeat(*args, **kwargs):
    return deepy_tensor.repeat(*args, **kwargs)

def vars(*tensor_types):
    """
    Create multiple variables without specifying last dimension and shape.
    :rtype: list of deepy.core.neural_var.NeuralVariable
    """
    return map(var, tensor_types)

def var(tensor_type, last_dim=0, test_shape=None):
    """
    Wrap a Theano tensor into the variable for defining neural network.
    :param last_dim: last dimension of tensor, 0 indicates that the last dimension is flexible
    :rtype: deepy.core.neural_var.NeuralVariable
    """
    # Create tensor
    from deepy.core.neural_var import NeuralVariable
    from deepy.core.env import env
    from theano.tensor.var import TensorVariable
    if isinstance(tensor_type, NeuralVariable):
        var = tensor_type
        if last_dim != 0:
            var.output_dim = last_dim
    elif isinstance(tensor_type, TensorVariable):
        var = NeuralVariable(tensor_type, dim=last_dim)
    elif isinstance(tensor_type, str):
        theano_tensor = getattr(TT, tensor_type)()
        var = NeuralVariable(theano_tensor, dim=last_dim)
    else:
        raise Exception("tensor_type shall be a string or a NeuralVariable")
    # Set test value
    if test_shape:
        if type(test_shape) != list and type(test_shape) != tuple:
            # May be it's a value
            var.set_test_value(test_shape)
        else:
            test_val = env.numpy_rand.rand(*test_shape)
            if len(test_shape) > 0:
                test_val = test_val.astype(var.tensor.dtype)
            elif var.tensor.dtype.startswith("int"):
                test_val = 1
            var.set_test_value(test_val)
    else:
        # Create a general test_shape
        dims = [(d + 1) * 3 for d in range(var.tensor.ndim)]
        if var.dim() != 0:
            dims[-1] = var.dim()
        test_val = env.numpy_rand.rand(*dims)
        if len(dims) > 0:
            test_val = test_val.astype(var.tensor.dtype)
        elif var.tensor.dtype.startswith("int"):
            test_val = 1
        var.set_test_value(test_val)
    return var

def is_neural_var(var):
    from deepy.core.neural_var import NeuralVariable
    return isinstance(var, NeuralVariable)

def is_theano_var(var):
    from theano.tensor.var import TensorVariable
    return isinstance(var, TensorVariable)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from . import BasicDataset
from padding import pad_dataset

class SequentialDataset(BasicDataset):
    """
    Dataset class for sequences.
    """

    def __init__(self, train, valid=None, test=None):
        super(SequentialDataset, self).__init__(train, valid, test)

    def _pad(self, side, length):
        """
        Pad sequences to given length in the left or right side.
        """
        if self._train_set:
            self._train_set = pad_dataset(self._train_set, side, length)
        if self._valid_set:
            self._valid_set = pad_dataset(self._valid_set, side, length)
        if self._test_set:
            self._test_set = pad_dataset(self._test_set, side, length)

    def pad_left(self, length):
        """
        Pad sequences to given length in the left side.
        """
        self._pad('left', length)

    def pad_right(self, length):
        """
        Pad sequences to given length in the left side.
        """
        self._pad('right', length)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import theano
import theano.tensor as T
import numpy as np

def rmsprop_core(params, gradients, momentum=0.9, learning_rate=0.01):
    """
    RMSPROP optimization core.
    """
    for param, grad in zip(params, gradients):
            rms_ = theano.shared(np.zeros_like(param.get_value()), name=param.name + '_rms')
            rms = momentum * rms_ + (1 - momentum) * grad * grad
            yield rms_, rms
            yield param, param - learning_rate * grad / T.sqrt(rms + 1e-8)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
from deepy.core.env import FLOATX

def pad_dataset(subset, side="right", length=-1):
    """
    Pad data set to specified length.
    Parameters:
        length - max length, a just to the max length in the batch if length is -1
    """
    assert length == -1 or length > 0
    if type(subset[0][0][0]) in [float, int, np.int64, np.int32, np.float32]:
        return _pad_2d(subset, side, length)
    else:
        return _pad_3d(subset, side, length)

def _pad_2d(subset, side, length):
    new_set = []
    max_len = max([len(x) for x, _ in subset]) if length == -1 else length
    for x, y in subset:
        if len(y) > max_len:
            y = y[:max_len]
        elif len(y) < max_len:
            if side == "left":
                y = [0 for _ in range(max_len - len(y))] + y
            elif side == "right":
                y = y + [0 for _ in range(max_len - len(y))]
        if len(x) > max_len:
            x = x[:max_len]
        elif len(x) < max_len:
            if side == "left":
                x = [0 for _ in range(max_len - len(x))] + x
            elif side == "right":
                x = x + [0 for _ in range(max_len - len(x))]
        new_set.append((x, y))
    return new_set

def _pad_3d(subset, side, length):
    row_size = subset[0][0][0].shape[0]
    new_set = []
    max_len = max([len(x) for x, _ in subset]) if length == -1 else length
    for x, y in subset:
        if type(y) == list:
            # Clip target vector
            if len(y) > max_len:
                y = y[:max_len]
            elif len(y) < max_len:
                if side == "left":
                    y = [0 for _ in range(max_len - len(y))] + y
                elif side == "right":
                    y = y + [0 for _ in range(max_len - len(y))]
        if len(x) > max_len:
            x = x[:max_len]
        elif len(x) < max_len:
            pad_length = max_len - len(x)
            pad_matrix = np.zeros((pad_length,row_size), dtype=FLOATX)
            if side == "left":
                x = np.vstack([pad_matrix, x])
            elif side == "right":
                x = np.vstack([x, pad_matrix])
            else:
                return Exception("Side of padding must be 'left' or 'right'")
        new_set.append((x, y))
    return new_set


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function
import time
import sys, os
import numpy as np
import logging as loggers
from threading import Lock
logging = loggers.getLogger("ScheduledTrainingServer")
loggers.basicConfig(level=loggers.INFO)

from platoon.channel import Controller
from argparse import ArgumentParser

CONTROLLER_PORT = 5567

class ScheduledTrainingServer(Controller):
    """
    This multi-process controller implements patience-based early-stopping SGD
    """

    def __init__(self, port=CONTROLLER_PORT, easgd_alpha=0.5,
                 # Following arguments can be received from workers
                 start_halving_at=6, end_at=10, sync_freq=10, halving_freq=1,
                 valid_freq=1500, learning_rate=0.1, log_path=None):
        """
        Initialize the controller.

        Args:
            port (int): batches in one training step
            easgd_alpha (float)
        """

        Controller.__init__(self, port)
        self.epoch_start_halving = start_halving_at
        self.end_at = end_at
        self.sync_freq = sync_freq
        self.start_time = None
        self.rand = np.random.RandomState(3)
        self.epoch = 0
        self._current_iter = 0
        self._iters_from_last_valid = 0
        self._evaluating = False
        self._valid_freq = valid_freq
        self._halving_freq = halving_freq
        self._done = False
        self._lr = learning_rate
        self._easgd_alpha = easgd_alpha
        self._training_names = []
        self._evaluation_names = []
        self._best_valid_cost = sys.float_info.max
        self._lock = Lock()

        self.num_train_batches = 0
        self.batch_pool = []
        self._train_costs = []
        self._epoch_start_time = None
        self.prepared_worker_pool = set()
        self.log_file = open(log_path, "w") if log_path else None
        if log_path:
            logging.info("write logs into {}".format(log_path))
        logging.info("multi-gpu server is listening port {}".format(port))

    def log(self, msg):
        logging.info(msg)
        if self.log_file:
            self.log_file.write(msg + "\n")

    def prepare_epoch(self):
        """
        Prepare for one epoch.
        Returns:
            bool: False if to stop the training.
        """
        self.epoch += 1
        if self.epoch >= self.epoch_start_halving and ((self.epoch - self.epoch_start_halving) % self._halving_freq == 0):
            self._lr *= 0.5
        self._current_iter = 0
        self._iters_from_last_valid = 0
        self._train_costs = []
        self.prepared_worker_pool.clear()
        self.batch_pool = range(self.num_train_batches)
        self.rand.shuffle(self.batch_pool)
        if self.epoch > self.end_at:
            self.log("Training is done, wait all workers to stop")
            return False
        else:
            self.log("start epoch {} with lr={}".format(self.epoch, self._lr))
            return True

    def feed_batches(self):
        if not self.batch_pool:
            return None
        else:
            batches = self.batch_pool[:self.sync_freq]
            self.batch_pool = self.batch_pool[self.sync_freq:]
            self._current_iter += len(batches)
            self._iters_from_last_valid += len(batches)
            return batches

    def feed_hyperparams(self):
        retval = {
            "epoch": self.epoch,
            "learning_rate": self._lr,
            "easgd_alpha": self._easgd_alpha
        }
        return retval

    def get_monitor_string(self, costs):
        return " ".join(["{}={:.2f}".format(n, c) for (n, c) in costs])

    def handle_control(self, req, worker_id, req_info):
        """
        Handles a control_request received from a worker.
        Returns:
            string or dict: response

            'stop' - the worker should quit
            'wait' - wait for 1 second
            'eval' - evaluate on valid and test set to start a new epoch
            'sync_hyperparams' - set learning rate
            'valid' - evaluate on valid and test set, then save the params
            'train' - train next batches
        """
        if self.start_time is None: self.start_time = time.time()
        response = ""

        if req == 'next':
            if self.num_train_batches == 0:
                response = "get_num_batches"
            elif self._done:
                response = "stop"
                self.worker_is_done(worker_id)
            elif self._evaluating:
                response = 'wait'
            elif not self.batch_pool:
                # End of one iter
                if self._train_costs:
                    with self._lock:
                        sys.stdout.write("\r")
                        sys.stdout.flush()
                        mean_costs = []
                        for i in range(len(self._training_names)):
                            mean_costs.append(np.mean([c[i] for c in self._train_costs]))
                        self.log("train   (epoch={:2d}) {}".format(
                            self.epoch,
                            self.get_monitor_string(zip(self._training_names, mean_costs)))
                        )
                response = {'eval': None, 'best_valid_cost': self._best_valid_cost}
                self._evaluating = True
            else:
                # Continue training
                if worker_id not in self.prepared_worker_pool:
                    response = {"sync_hyperparams": self.feed_hyperparams()}
                    self.prepared_worker_pool.add(worker_id)
                elif self._iters_from_last_valid >= self._valid_freq:
                    response = {'valid': None, 'best_valid_cost': self._best_valid_cost}
                    self._iters_from_last_valid = 0
                else:
                    response = {"train": self.feed_batches()}
        elif 'eval_done' in req:
            with self._lock:
                self._evaluating = False
                sys.stdout.write("\r")
                sys.stdout.flush()
                if 'test_costs' in req and req['test_costs']:
                    self.log("test    (epoch={:2d}) {} (worker {})".format(
                        self.epoch,
                        self.get_monitor_string(req['test_costs']),
                        worker_id)
                    )
                if 'valid_costs' in req and req['test_costs']:
                    valid_J = req['valid_costs'][0][1]
                    if valid_J < self._best_valid_cost:
                        self._best_valid_cost = valid_J
                        star_str = "*"
                    else:
                        star_str = ""
                        self.log("valid   (epoch={:2d}) {} {} (worker {})".format(
                        self.epoch,
                        self.get_monitor_string(req['valid_costs']),
                        star_str,
                        worker_id))
                    # if star_str and 'auto_save' in req and req['auto_save']:
                    #     self.log("(worker {}) save the model to {}".format(
                    #         worker_id,
                    #         req['auto_save']
                    #     ))
                continue_training = self.prepare_epoch()
                self._epoch_start_time = time.time()
                if not continue_training:
                    self._done = True
                    self.log("training time {:.4f}s".format(time.time() - self.start_time))
                    response = "stop"
        elif 'valid_done' in req:
            with self._lock:
                sys.stdout.write("\r")
                sys.stdout.flush()
                if 'valid_costs' in req:
                    valid_J = req['valid_costs'][0][1]
                    if valid_J < self._best_valid_cost:
                        self._best_valid_cost = valid_J
                        star_str = "*"
                    else:
                        star_str = ""
                    self.log("valid   ( dryrun ) {} {} (worker {})".format(
                        self.get_monitor_string(req['valid_costs']),
                        star_str,
                        worker_id
                    ))
                    # if star_str and 'auto_save' in req and req['auto_save']:
                    #     self.log("(worker {}) save the model to {}".format(
                    #         worker_id,
                    #         req['auto_save']
                    #     ))
        elif 'train_done' in req:
            costs = req['costs']
            self._train_costs.append(costs)
            sys.stdout.write("\x1b[2K\r> %d%% | J=%.2f | %.1f batch/s" % (
                self._current_iter * 100 / self.num_train_batches,
                costs[0], float(len(self._train_costs) * self.sync_freq) / (time.time() - self._epoch_start_time)))
            sys.stdout.flush()
        elif 'get_num_batches_done' in req:
            self.num_train_batches = req['get_num_batches_done']
        elif 'get_easgd_alpha' in req:
            response = self._easgd_alpha
        elif 'sync_hyperparams' in req:
            response = {"sync_hyperparams": self.feed_hyperparams()}
        elif 'init_schedule' in req:
            with self._lock:
                sys.stdout.write("\r")
                sys.stdout.flush()
                self.log("worker {} connected".format(worker_id))
                if self.epoch == 0:
                    schedule_params = req['init_schedule']
                    sch_str = " ".join("{}={}".format(a, b) for (a, b) in schedule_params.items())
                    self.log("initialize the schedule with {}".format(sch_str))
                    for key, val in schedule_params.items():
                        if not val: continue
                        if key == 'learning_rate':
                            self._lr = val
                        elif key == 'start_halving_at':
                            self.epoch_start_halving = val
                        elif key == 'halving_freq':
                            self._halving_freq = val
                        elif key == 'end_at':
                            self.end_at = val
                        elif key == 'sync_freq':
                            self.sync_freq = val
                        elif key == 'valid_freq':
                            self._valid_freq = val

        elif 'set_names' in req:
            self._training_names = req['training_names']
            self._evaluation_names = req['evaluation_names']

        return response

if __name__ == '__main__':
    ap = ArgumentParser()
    ap.add_argument("--port", type=int, default=5567)
    ap.add_argument("--easgd_alpha", type=float, default=0.5)
    ap.add_argument("--log", type=str, default=None)
    args = ap.parse_args()

    server = ScheduledTrainingServer(
        port=args.port,
        easgd_alpha=args.easgd_alpha,
        log_path=args.log
    )

    server.serve()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import time

class Timer(object):
    """
    A timer for measure runing time.
    """

    def __init__(self):
        self.start_time = time.time()
        self.end_time = None

    def end(self):
        """
        Stop the timer.
        """
        self.end_time = time.time()

    def report(self):
        """
        Report elapsed time.
        """
        if not self.end_time:
            self.end()
        print ("Time: {} mins".format((self.end_time - self.start_time )/ 60))


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from collections import OrderedDict, defaultdict

class TrainingController(object):
    """
    Abstract class of training controllers.
    """

    def bind(self, trainer):
        """
        :type trainer: deepy.trainers.base.NeuralTrainer
        """
        self._trainer = trainer

    def invoke(self):
        """
        Return True to exit training.
        """
        return False

class TrainingValidator(TrainingController):
    """
    A validator that allows validating the model with another graph.
    """

    def __init__(self, valid_model=None, data_split='valid', freq=1500, save_path=None, criteria='cost',
                 smaller_is_better=True):
        """
        Initialize the training validator.
        """
        self._model = valid_model
        self._data_split = data_split
        self._freq = freq
        self._save_path = save_path
        self._criteria = criteria
        self._smaller_is_better = smaller_is_better
        self._best_criteria = None
        self._counter = 0

    def compare(self, cost_map):
        """
        Compare to previous records and return whether the given cost is a new best.
        :return: True if the given cost is a new best
        """
        cri_val = cost_map[self._criteria]
        if self._best_criteria is None:
            self._best_criteria = cri_val
            return True
        else:
            if self._smaller_is_better and cri_val < self._best_criteria:
                self._best_criteria = cri_val
                return True
            elif not self._smaller_is_better and cri_val > self._best_criteria:
                self._best_criteria = cri_val
                return True
            else:
                return False

    def compute(self, *x):
        """
        Compute with the validation model given data x.
        """
        return self._model.compute(*x)

    def _extract_costs(self, vars):
        ret_map = OrderedDict()
        sub_costs = OrderedDict()
        for k, val in vars.items():
            if val.ndim == 0:
                if k == self._criteria:
                    ret_map[k] = val
                else:
                    sub_costs[k] = val
        ret_map.update(sub_costs)
        return ret_map

    def run(self, data_x):
        """
        Run the model with validation data and return costs.
        """
        output_vars = self.compute(*data_x)
        return self._extract_costs(output_vars)

    def invoke(self):
        """
        This function will be called after each iteration.
        """
        self._counter += 1
        if self._counter % self._freq == 0:
            cnt = 0.
            sum_map = defaultdict(float)
            for x in self._trainer.get_data(self._data_split):
                val_map = self.run(x)
                if not isinstance(val_map, dict):
                    raise Exception("Monitor.run must return a dict.")
                for k, val in val_map.items():
                    sum_map[k] += val
                cnt += 1
            for k in sum_map:
                sum_map[k] /= cnt
            new_best = self.compare(sum_map)
            self._trainer.report(sum_map, self._data_split, new_best=new_best)
            if new_best:
                self._trainer.save_checkpoint(self._save_path)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from deepy.utils import Scanner, MapDict
from dummy_loop_utils import get_dummy_args, finish_scan
import theano

class LoopVars(MapDict):
    """
    Variables inside the loop.
    """

class Loop(object):

    def __init__(self, sequences=None, outputs=None, non_sequences=None, block=None, **kwargs):
        """
        A loop function to support "with" grammar.
        """
        self._sequences = sequences if sequences else {}
        self._outputs = outputs if outputs else {}
        self._non_sequences = non_sequences if non_sequences else {}
        if not isinstance(self._sequences, dict) or not isinstance(self._outputs, dict) != dict or not isinstance(self._non_sequences, dict):
            raise Exception("Arguments of Loop shall be dicts.")
        self._block = block
        self._kwargs = kwargs
        self._loop_vars = None
        self._dummy_nodes = None
        self._scan_local_vars = None
        self._ordered_out_keys = []
        self._scan_outputs = None

    def _build_loop_vars(self):
        """
        Create inner loop variables.
        """
        from theano.tensor.var import TensorVariable
        from deepy.core.neural_var import NeuralVariable
        if not self._loop_vars:
            self._ordered_out_keys = self._outputs.keys()
            seq_keys = self._sequences.keys()
            filled_out_keys = [k for k in self._ordered_out_keys if self._outputs[k]]
            nonseq_keys = self._non_sequences.keys()
            dummy_tensors, self._scan_local_vars = get_dummy_args(
                sequences=[self._sequences[k].tensor for k in seq_keys],
                outputs_info=[self._outputs[k].tensor for k in self._ordered_out_keys],
                non_sequences=[self._non_sequences[k].tensor for k in nonseq_keys],
                **self._kwargs
            )
            dummy_map = dict(zip(seq_keys + filled_out_keys + nonseq_keys, dummy_tensors))
            arg_map = self._sequences.copy()
            arg_map.update(self._outputs)
            arg_map.update(self._non_sequences)
            self._loop_vars = LoopVars()
            for k, dummy_tensor in dummy_map.items():
                dummy_var = NeuralVariable(dummy_tensor, dim=arg_map[k].dim())
                self._loop_vars[k] = dummy_var
            # self._dummy_nodes = dict(self._loop_vars.items()[:])

    def __enter__(self):
        self._build_loop_vars()
        return self._loop_vars

    def __exit__(self, exc_type, exc_val, exc_tb):
        from neural_var import NeuralVariable
        output_tensors = []
        for k in self._ordered_out_keys:
            if k not in self._loop_vars:
                raise Exception("{} can not be found in loop vars.".format(k))
            output_tensors.append(self._loop_vars[k].tensor)

        result_tensors, updates = finish_scan(output_tensors, self._scan_local_vars)
        if self._block and updates:
            if type(updates) == dict:
                updates = updates.items()
            self._block.register_updates(*updates)

        outputs = MapDict()
        for k, tensor in zip(self._ordered_out_keys, result_tensors):
            out_var = NeuralVariable(tensor)
            if self._outputs[k] is not None:
                out_var.output_dim = self._outputs[k].dim()
            outputs[k] = out_var
        self._scan_outputs = outputs

    def _scan_step(self, vars):
        """
        Internal scan with dummy input variables.
        """
        from neural_var import NeuralVariable
        if not self._loop_vars:
            raise Exception("The loop is not initialized. To initialize the loop, use `with loop as vars`")
        replace_map = {}
        for k, var in vars.items():
            if var is not None:
                replace_map[self._dummy_nodes[k].tensor] = var.tensor
        outputs = {}
        for k in self._outputs:
            if k not in self._loop_vars:
                raise Exception("{} can not be found in loop vars.".format(k))
            output_node = theano.clone(self._loop_vars[k].tensor, replace_map)
            outputs[k] = NeuralVariable(output_node, self._loop_vars[k].dim())
        return outputs

    @property
    def outputs(self):
        """
        Get the output of the loop.
        :rtype: MapDict
        """
        if not self._scan_outputs:
            raise Exception("The loop is not executed.")
        else:
            return self._scan_outputs

    def get_outputs(self, *args):
        """
        Get the outputs of the loop.
        Return specific variables by passing the keys to the arguments.
        :rtype: MapDict
        """
        if args:
            output_vars = map(self._scan_outputs.get, args)
            if len(output_vars) == 1:
                return output_vars[0]
            else:
                return output_vars
        else:
            return self._scan_outputs


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import theano
import numpy as np

def momentum_core(params, gradients, momentum=0.9, learning_rate=0.01):
    """
    Momentum SGD optimization core.
    """
    free_parameters = []
    updates = []
    for param, grad in zip(params, gradients):
            delta = learning_rate * grad
            velocity = theano.shared(np.zeros_like(param.get_value()), name=param.name + '_vel')
            updates.append((velocity, momentum * velocity - delta))
            updates.append((param, param + velocity))
            free_parameters.append(velocity)
    return updates, free_parameters


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import theano
import theano.tensor as T
from theano.ifelse import ifelse
from tensor_conversion import neural_computation

class Runtime(object):
    """
    Manage runtime variables in deepy.
    """

    def __init__(self):
        self._training_flag = theano.shared(0, name="training_flag")
        self._is_training = False

    @neural_computation
    def iftrain(self, then_branch, else_branch):
        """
        Execute `then_branch` when training.
        """
        return ifelse(self._training_flag, then_branch, else_branch, name="iftrain")

    def switch_training(self, flag):
        """
        Switch training mode.
        :param flag: switch on training mode when flag is True.
        """
        if self._is_training == flag: return
        self._is_training = flag
        if flag:
            self._training_flag.set_value(1)
        else:
            self._training_flag.set_value(0)

if "runtime" not in globals():
    runtime = Runtime()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import theano
import theano.tensor as T
import numpy as np

def nag_core(params, J, momentum=0.9, learning_rate=0.01):
    """
    Nesterov's Accelerated Gradient (NAG).
    See http://www.cs.toronto.edu/~fritz/absps/momentum.pdf .
    Still unfinished
    """
    # TODO: this requires some refractorings.
    for param in params:
        step = theano.shared(np.zeros_like(param.get_value()), name=param.name + '_step')
        velocity = theano.shared(np.zeros_like(param.get_value()), name=param.name + '_vel')
        yield step, momentum * velocity
        yield param, param + step
        yield velocity, step - learning_rate * T.grad(J, param)
        yield param, param + velocity - step


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import time
import numpy as np
import theano

from ..conf import TrainerConfig
from ..core import env, runtime
from ..utils import Timer
from ..dataset import Dataset
from controllers import TrainingController

from abc import ABCMeta, abstractmethod

from logging import getLogger
logging = getLogger("trainer")

class NeuralTrainer(object):
    """
    A base class for all trainers.
    """
    __metaclass__ = ABCMeta

    def __init__(self, network, config=None, validator=None, annealer=None):
        """
        Basic neural network trainer.
        :type network: deepy.NeuralNetwork
        :type config: deepy.conf.TrainerConfig
        :return:
        """
        super(NeuralTrainer, self).__init__()

        self.config = None
        if isinstance(config, TrainerConfig):
            self.config = config
        elif isinstance(config, dict):
            self.config = TrainerConfig(config)
        else:
            self.config = TrainerConfig()
        if type(self.config.learning_rate) == float:
            self.config.learning_rate = np.array(self.config.learning_rate, dtype=env.FLOATX)
        # Model and network all refer to the computational graph
        self.model = self.network = network

        self.network.prepare_training()
        self._setup_costs()

        self.evaluation_func = None

        self.validation_frequency = self.config.validation_frequency
        self.min_improvement = self.config.min_improvement
        self.patience = self.config.patience

        self._iter_controllers = []
        self._epoch_controllers = []
        if annealer:
            annealer.bind(self)
            self._epoch_controllers.append(annealer)
        if validator:
            validator.bind(self)
            self._iter_controllers.append(validator)

        self.best_cost = 1e100
        self.best_epoch = 0
        self.best_params = self.copy_params()
        self._skip_batches = 0
        self._skip_epochs = 0
        self._progress = 0
        self.last_cost = 0
        self.last_run_costs = None
        self._report_time = True
        self._epoch = 0

        self._current_train_set = None
        self._current_valid_set = None
        self._current_test_set = None
        self._ended = False

    def _compile_evaluation_func(self):
        if not self.evaluation_func:
            logging.info("compile evaluation function")
            self.evaluation_func = theano.function(
                self.network.input_variables + self.network.target_variables,
                self.evaluation_variables,
                updates=self.network.updates,
                allow_input_downcast=True, mode=self.config.get("theano_mode", None))

    def skip(self, n_batches, n_epochs=0):
        """
        Skip N batches in the training.
        """
        logging.info("skip %d epochs and %d batches" % (n_epochs, n_batches))
        self._skip_batches = n_batches
        self._skip_epochs = n_epochs

    def epoch(self):
        """
        Get current epoch.
        """
        return self._epoch

    def _setup_costs(self):
        self.cost = self._add_regularization(self.network.cost)
        self.test_cost = self._add_regularization(self.network.test_cost)
        self.training_variables = [self.cost]
        self.training_names = ['J']
        for name, monitor in self.network.training_monitors:
            self.training_names.append(name)
            self.training_variables.append(monitor)
        logging.info("monitor list: %s" % ",".join(self.training_names))

        self.evaluation_variables = [self.test_cost]
        self.evaluation_names = ['J']
        for name, monitor in self.network.testing_monitors:
            self.evaluation_names.append(name)
            self.evaluation_variables.append(monitor)

    def _add_regularization(self, cost):
        if self.config.weight_l1 > 0:
            logging.info("L1 weight regularization: %f" % self.config.weight_l1)
            cost += self.config.weight_l1 * sum(abs(w).sum() for w in self.network.parameters)
        if self.config.hidden_l1 > 0:
            logging.info("L1 hidden unit regularization: %f" % self.config.hidden_l1)
            cost += self.config.hidden_l1 * sum(abs(h).mean(axis=0).sum() for h in self.network._hidden_outputs)
        if self.config.hidden_l2 > 0:
            logging.info("L2 hidden unit regularization: %f" % self.config.hidden_l2)
            cost += self.config.hidden_l2 * sum((h * h).mean(axis=0).sum() for h in self.network._hidden_outputs)

        return cost

    def set_params(self, targets, free_params=None):
        for param, target in zip(self.network.parameters, targets):
            param.set_value(target)
        if free_params:
            for param, param_value in zip(self.network.free_parameters, free_params):
                param.set_value(param_value)

    def save_params(self, path):
        self.set_params(*self.best_params)
        self.network.save_params(path)

    def load_params(self, path, exclude_free_params=False):
        """
        Load parameters for the training.
        This method can load free parameters and resume the training progress.
        """
        self.network.load_params(path, exclude_free_params=exclude_free_params)
        self.best_params = self.copy_params()
        # Resume the progress
        if self.network.train_logger.progress() > 0 or self.network.train_logger.epoch() > 0:
            self.skip(self.network.train_logger.progress(), self.network.train_logger.epoch() - 1)

    def copy_params(self):
        checkpoint = (map(lambda p: p.get_value().copy(), self.network.parameters),
                      map(lambda p: p.get_value().copy(), self.network.free_parameters))
        return checkpoint

    def exit(self):
        """
        End the training.
        """
        self._ended = True

    def add_iter_controllers(self, *controllers):
        """
        Add iteration callbacks function (receives an argument of the trainer).
        :param controllers: can be a `TrainingController` or a function.
        :type funcs: list of TrainingContoller
        """
        for controller in controllers:
            if isinstance(controller, TrainingController):
                controller.bind(self)
            self._iter_controllers.append(controller)

    def add_epoch_controllers(self, *controllers):
        """
        Add epoch callbacks function.
        :param controllers: can be a `TrainingController` or a function.
        """
        for controller in controllers:
            if isinstance(controller, TrainingController):
                controller.bind(self)
            self._epoch_controllers.append(controller)

    def train(self, train_set, valid_set=None, test_set=None, train_size=None):
        """
        Train the model and return costs.
        """
        self._epoch = 0
        while True:
            if self._skip_epochs > 0:
                logging.info("skipping one epoch ...")
                self._skip_epochs -= 1
                self._epoch += 1
                yield None
                continue
            # Test
            if not self._epoch % self.config.test_frequency and test_set:
                try:
                    self._run_test(self._epoch, test_set)
                except KeyboardInterrupt:
                    logging.info('interrupted!')
                    break
            # Validate
            if not self._epoch % self.validation_frequency and valid_set:
                try:

                    if not self._run_valid(self._epoch, valid_set):
                        logging.info('patience elapsed, bailing out')
                        break
                except KeyboardInterrupt:
                    logging.info('interrupted!')
                    break
            # Train one step

            try:
                costs = self._run_train(self._epoch, train_set, train_size)
            except KeyboardInterrupt:
                logging.info('interrupted!')
                break
            # Check costs
            if np.isnan(costs[0][1]):
                logging.info("NaN detected in costs, rollback to last parameters")
                self.set_params(*self.checkpoint)
            else:
                self._epoch += 1
                self.network.epoch_callback()

            yield dict(costs)

        if valid_set and self.config.get("save_best_parameters", True):
            self.set_params(*self.best_params)
        if test_set:
            self._run_test(-1, test_set)

    @abstractmethod
    def learn(self, *variables):
        """
        Update the parameters and return the cost with given data points.
        :param variables:
        :return:
        """

    def _run_test(self, epoch, test_set):
        """
        Run on test epoch.
        """
        costs = self.test_step(test_set)
        self.report(dict(costs), "test", epoch)
        self.last_run_costs = costs

    def _run_train(self, epoch, train_set, train_size=None):
        """
        Run one training iteration.
        """
        self.network.train_logger.record_epoch(epoch + 1)
        costs = self.train_step(train_set, train_size)
        if not epoch % self.config.monitor_frequency:
            self.report(dict(costs), "train", epoch)
        self.last_run_costs = costs
        return costs

    def _run_valid(self, epoch, valid_set, dry_run=False, save_path=None):
        """
        Run one valid iteration, return true if to continue training.
        """
        costs = self.valid_step(valid_set)
        # this is the same as: (J_i - J_f) / J_i > min improvement
        _, J = costs[0]
        new_best = False
        if self.best_cost - J > self.best_cost * self.min_improvement:
            # save the best cost and parameters
            self.best_params = self.copy_params()
            new_best = True
            if not dry_run:
                self.best_cost = J
                self.best_epoch = epoch
            self.save_checkpoint(save_path)

        self.report(dict(costs), type="valid", epoch=0 if dry_run else epoch, new_best=new_best)
        self.last_run_costs = costs
        return epoch - self.best_epoch < self.patience

    def save_checkpoint(self, save_path=None):
        save_path = save_path if save_path else self.config.auto_save
        self.checkpoint = self.copy_params()
        if save_path and self._skip_batches == 0:
            self.network.train_logger.record_progress(self._progress)
            self.network.save_params(save_path, new_thread=True)

    def report(self, score_map, type="valid", epoch=-1, new_best=False):
        """
        Report the scores and record them in the log.
        """
        type_str = type
        if len(type_str) < 5:
            type_str += " " * (5 - len(type_str))
        info = " ".join("%s=%.2f" % el for el in score_map.items())
        current_epoch = epoch if epoch > 0 else self.current_epoch()
        epoch_str = "epoch={}".format(current_epoch + 1)
        if epoch < 0:
            epoch_str = "dryrun"
            sys.stdout.write("\r")
            sys.stdout.flush()
        marker = " *" if new_best else ""
        message = "{} ({}) {}{}".format(type_str, epoch_str, info, marker)
        self.network.train_logger.record(message)
        logging.info(message)

    def test_step(self, test_set):
        runtime.switch_training(False)
        self._compile_evaluation_func()
        costs = list(zip(
            self.evaluation_names,
            np.mean([self.evaluation_func(*x) for x in test_set], axis=0)))
        return costs

    def valid_step(self, valid_set):
        runtime.switch_training(False)
        self._compile_evaluation_func()
        costs = list(zip(
            self.evaluation_names,
            np.mean([self.evaluation_func(*x) for x in valid_set], axis=0)))
        return costs

    def train_step(self, train_set, train_size=None):
        dirty_trick_times = 0
        network_callback = bool(self.network.training_callbacks)
        trainer_callback = bool(self._iter_controllers)
        cost_matrix = []
        exec_count = 0
        start_time = time.time()
        self._compile_time = 0
        self._progress = 0

        for x in train_set:
            runtime.switch_training(True)
            if self._skip_batches == 0:

                if dirty_trick_times > 0:
                    cost_x = self.learn(*[t[:(t.shape[0]/2)] for t in x])
                    cost_matrix.append(cost_x)
                    cost_x = self.learn(*[t[(t.shape[0]/2):] for t in x])
                    dirty_trick_times -= 1
                else:
                    try:
                        cost_x = self.learn(*x)
                    except MemoryError:
                        logging.info("Memory error was detected, perform dirty trick 30 times")
                        dirty_trick_times = 30
                        # Dirty trick
                        cost_x = self.learn(*[t[:(t.shape[0]/2)] for t in x])
                        cost_matrix.append(cost_x)
                        cost_x = self.learn(*[t[(t.shape[0]/2):] for t in x])
                cost_matrix.append(cost_x)
                self.last_cost = cost_x[0]
                exec_count += 1
                if network_callback:
                    self.network.training_callback()
                if trainer_callback:
                    for func in self._iter_controllers:
                        if isinstance(func, TrainingController):
                            func.invoke()
                        else:
                            func(self)
            else:
                self._skip_batches -= 1
            if train_size:
                self._progress += 1
                spd = float(exec_count) / (time.time() - start_time - self._compile_time)
                sys.stdout.write("\x1b[2K\r> %d%% | J=%.2f | spd=%.2f batch/s" % (self._progress * 100 / train_size, self.last_cost, spd))
                sys.stdout.flush()
        self._progress = 0

        if train_size:
            sys.stdout.write("\r")
            sys.stdout.flush()
        costs = list(zip(self.training_names, np.mean(cost_matrix, axis=0)))
        return costs

    def current_epoch(self):
        """
        Get current epoch.
        """
        return self._epoch

    def get_data(self, data_split="train"):
        """
        Get specified split of data.
        """
        if data_split == 'train':
            return self._current_train_set
        elif data_split == 'valid':
            return self._current_valid_set
        elif data_split == 'test':
            return self._current_test_set
        else:
            return None

    def run(self, train_set, valid_set=None, test_set=None, train_size=None, epoch_controllers=None):
        """
        Run until the end.
        :param epoch_controllers: deprecated
        """
        epoch_controllers = epoch_controllers if epoch_controllers else []
        epoch_controllers += self._epoch_controllers
        if isinstance(train_set, Dataset):
            dataset = train_set
            train_set = dataset.train_set()
            valid_set = dataset.valid_set()
            test_set = dataset.test_set()
            train_size = dataset.train_size()
        self._current_train_set = train_set
        self._current_valid_set = valid_set
        self._current_test_set = test_set
        if epoch_controllers:
            for controller in epoch_controllers:
                controller.bind(self)
        timer = Timer()
        for _ in self.train(train_set, valid_set=valid_set, test_set=test_set, train_size=train_size):
            if epoch_controllers:
                for controller in epoch_controllers:
                    controller.invoke()
            if self._ended:
                break
        if self._report_time:
            timer.report()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging as loggers
logging = loggers.getLogger(__name__)

from . import MiniBatches
import numpy as np
from deepy.utils import FakeGenerator
from deepy.core import env

class BunchSequences(MiniBatches):
    """
    DEPRECATED.
    Arrange sequences in bunch mode.
    See http://mi.eng.cam.ac.uk/~xc257/papers/RNNLMTrain_Interspeech2014.pdf .
    """

    def __init__(self, dataset, batch_size=20, fragment_length=5):
        super(BunchSequences, self).__init__(dataset, batch_size=batch_size)
        self.fragment_length = fragment_length
        if self.fragment_length < 1:
            raise SystemError("fragment_length must be greater than 1")

    def _yield_data(self, subset):
        subset = list(subset)
        env.numpy_rand.shuffle(subset)

        bunch_stack_x = [[] for _ in range(self.size)]
        bunch_stack_y = [[] for _ in range(self.size)]

        for x, y in subset:
            stack_lens = map(len, bunch_stack_x)
            shortest_i = stack_lens.index(min(stack_lens))
            bunch_stack_x[shortest_i].extend(x)
            bunch_stack_y[shortest_i].extend(y)
        self._pad_zeros(bunch_stack_x)
        self._pad_zeros(bunch_stack_y)
        pieces_x = self._cut_to_pieces(bunch_stack_x)
        pieces_y = self._cut_to_pieces(bunch_stack_y)
        logging.info("%d pieces this time" % int(float(len(bunch_stack_x[0])) / self.fragment_length))
        for piece in zip(pieces_x, pieces_y):
            yield piece

    def _train_set(self):
        if not self.origin.train_set():
            return None
        return self._yield_data(self.origin.train_set())

    def train_set(self):
        return FakeGenerator(self, "_train_set")

    def train_size(self):
        size = len([_ for _ in self.train_set()])
        return size

    def _cut_to_pieces(self, bunch_stack):
        """
        :type bunch_stack: list of list of int
        """
        stack_len = len(bunch_stack[0])
        for i in xrange(0, stack_len, self.fragment_length):
            yield np.array(map(lambda stack: stack[i: i + self.fragment_length], bunch_stack))

    def _pad_zeros(self, bunch_stack):
        """
        :type bunch_stack: list of list
        """
        min_len = min(map(len, bunch_stack))
        for i in range(len(bunch_stack)):
            bunch_stack[i] = bunch_stack[i][:min_len]
        # for stack in bunch_stack:
        #     for _ in range(max_len - len(stack)):
        #         stack.append(0)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import theano.tensor as TT

from deepy.core.tensor_conversion import neural_computation, convert_to_theano_var
from deepy.layers.layer import NeuralLayer

class NeuralVariable(NeuralLayer):
    """
    Create a constant layer with tensors.
    """

    def __init__(self, tensor, dim=0):
        """
        Create a tensor layer.
        :type tensor: theano.tensor.var.TensorVariable
        """
        super(NeuralVariable, self).__init__("const")
        self.output_dim = dim
        self.tensor = tensor
        self.init(0)

    def __getattr__(self, name):
        return NeuralVariable(getattr(self.tensor, name), dim=self.dim())

    def apply(self, func, dim=None):
        """
        Apply a function to tensors.
        """
        output_dim = dim if dim else self.output_dim
        return NeuralVariable(func(self.tensor), output_dim)

    def compute_tensor(self, x):
        return self.tensor

    def set_test_value(self, value):
        self.tensor.tag.test_value = value

    def dim(self):
        return self.output_dim

    def _other_tensor(self, other):
        return  other.tensor if isinstance(other, NeuralVariable) else other

    def __eq__(self, other):
        return NeuralVariable(TT.eq(self.tensor, self._other_tensor(other)), dim=self.dim())

    def __ne__(self, other):
        return NeuralVariable(TT.neq(self.tensor, self._other_tensor(other)), dim=self.dim())

    def __add__(self, other):
        return NeuralVariable(self.tensor + self._other_tensor(other), dim=self.dim())

    def __sub__(self, other):
        return NeuralVariable(self.tensor - self._other_tensor(other), dim=self.dim())

    def __mul__(self, other):
        return NeuralVariable(self.tensor * self._other_tensor(other), dim=self.dim())

    def __div__(self, other):
        return NeuralVariable(self.tensor / self._other_tensor(other), dim=self.dim())

    def __neg__(self):
        return NeuralVariable(- self.tensor, dim=self.dim())

    def __radd__(self, other):
        return NeuralVariable(self._other_tensor(other) + self.tensor, dim=self.dim())

    def __rsub__(self, other):
        return NeuralVariable(self._other_tensor(other) - self.tensor, dim=self.dim())

    def __rmul__(self, other):
        return NeuralVariable(self._other_tensor(other) * self.tensor, dim=self.dim())

    def __rdiv__(self, other):
        return NeuralVariable(self._other_tensor(other) / self.tensor, dim=self.dim())

    def __getitem__(self, index):
        @neural_computation
        def getitem_wrapper(t, index):
            if type(index) == list:
                index = tuple(index)
            return t.__getitem__(index)
        ret = getitem_wrapper(self, index)
        if (hasattr(ret.tensor, 'tag') and hasattr(ret.tensor.tag, 'test_value')
            and ret.tensor.tag.test_value is not None and len(ret.tensor.tag.test_value.shape) > 0):
            ret.output_dim = ret.tensor.tag.test_value.shape[-1]
        else:
            ret.output_dim = self.dim()
        return ret

    def __call__(self, *args, **kwargs):
        normal_args, tensor_found_in_args, neural_found_in_args = convert_to_theano_var(args)
        normal_kwargs, tensor_found_in_kwargs, neural_found_in_kwargs = convert_to_theano_var(kwargs)

        tensor_found = tensor_found_in_args or tensor_found_in_kwargs

        if tensor_found:
            raise Exception("Theano tensor variables can not be used together with neural variables.")

        return NeuralVariable(self.tensor(*normal_args, **normal_kwargs), dim=self.dim())

    def debug_monitor(self, name=""):
        from deepy.debug import monitor_var_sum
        self.tensor += monitor_var_sum(self.tensor, name=name)

    @property
    def test_value(self):
        if hasattr(self.tensor.tag, 'test_value'):
            return self.tensor.tag.test_value
        else:
            return None

    @property
    def ndim(self):
        return self.tensor.ndim

    @property
    def tv(self):
        return self.test_value

    @property
    def ts(self):
        if self.test_value is not None:
            return self.test_value.shape
        else:
            return None


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import theano.tensor as T
import theano

def rprop_core(params, gradients, rprop_increase=1.01, rprop_decrease=0.99, rprop_min_step=0, rprop_max_step=100,
               learning_rate=0.01):
    """
    Rprop optimizer.
    See http://sci2s.ugr.es/keel/pdf/algorithm/articulo/2003-Neuro-Igel-IRprop+.pdf.
    """
    for param, grad in zip(params, gradients):
        grad_tm1 = theano.shared(np.zeros_like(param.get_value()), name=param.name + '_grad')
        step_tm1 = theano.shared(np.zeros_like(param.get_value()) + learning_rate, name=param.name+ '_step')

        test = grad * grad_tm1
        same = T.gt(test, 0)
        diff = T.lt(test, 0)
        step = T.minimum(rprop_max_step, T.maximum(rprop_min_step, step_tm1 * (
            T.eq(test, 0) +
            same * rprop_increase +
            diff * rprop_decrease)))
        grad = grad - diff * grad
        yield param, param - T.sgn(grad) * step
        yield grad_tm1, grad
        yield step_tm1, step


#!/usr/bin/env python
# -*- coding: utf-8 -*-

class GeneralConfig(object):

    def __init__(self, logger=None):
        """
        Create a general config
        """
        object.__setattr__(self, "attrs", {})
        object.__setattr__(self, "used_parameters", set())
        object.__setattr__(self, "undefined_parameters", set())
        object.__setattr__(self, "logger", logger)

    def __getattr__(self, key):
        self.used_parameters.add(key)
        if key not in self.attrs:
            self.undefined_parameters.add(key)
            return None
        else:
            return self.attrs[key]

    def __setattr__(self, key, value):
        self.attrs[key] = value

    def get(self, key, default=None):
        key = getattr(self, key)
        if key != None:
            return key
        else:
            return default

    def merge(self, config):
        for k in config.attrs:
            setattr(self, k, getattr(config, k))
        return self

    def report(self):
        """
        Report usage of training parameters.
        """
        if self.logger:
            self.logger.info("accessed parameters:")
            for key in self.used_parameters:
                self.logger.info(" - %s %s" % (key, "(undefined)" if key in self.undefined_parameters else ""))


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import numpy as np
import pickle
import gzip
from inspect import getargspec
from env import env
import theano
import theano.tensor as TT
import logging as loggers
from tensor_conversion import neural_computation
from disconnected_grad import disconnected_grad
from deepy.utils import Scanner
logging = loggers.getLogger(__name__)

class GraphBuilder(object):
    """
    Tool for creating computational graph in deepy.
    """

    def __init__(self):
        self._default_block = self.new_block("default_block")

    def default_block(self):
        """
        Return the default block.
        """
        return self._default_block

    def collect_parameters(self):
        """
        Return the default block, as all parameters will be registered to the default one.
        """
        return self._default_block

    def new_block(self, *layers, **kwargs):
        """
        Create a parameters block.
        :param layers: register some layers in the block
        :param name: specify the name of this block
        """
        from deepy.layers.block import Block
        block = Block(*layers, **kwargs)
        return block

    def var(self, tensor_type, last_dim=0, test_shape=None):
        """
        An alias of deepy.tensor.var.
        """
        from deepy.tensor import var
        return var(tensor_type, last_dim=last_dim, test_shape=test_shape)

    def create_vars_from_data(self, dataset, split="train"):
        """
        Create vars given a dataset and set test values.
        Useful when dataset is already defined.
        """
        from deepy.core.neural_var import NeuralVariable
        vars = []
        if split == "valid":
            data_split = dataset.valid_set()
        elif split == "test":
            data_split = dataset.test_set()
        else:
            data_split = dataset.train_set()
        first_data_piece = list(data_split)[0]
        for i, numpy_tensor in enumerate(first_data_piece):
            if numpy_tensor.dtype == "int64":
                numpy_tensor = numpy_tensor.astype("int32")
            if numpy_tensor.dtype == "float64":
                numpy_tensor = numpy_tensor.astype(env.FLOATX)
            type_map = {
                0: "scalar",
                1: "vector",
                2: "matrix",
                3: "tensor3",
                4: "tensor4",
                5: "tensor5",
            }
            tensor_type = type_map[numpy_tensor.ndim] if numpy_tensor.ndim in type_map else type_map[0]
            if numpy_tensor.dtype.kind == "i":
                tensor_type = "i" + tensor_type
            theano_tensor = getattr(TT, tensor_type)("input_{}_{}".format(i + 1, tensor_type))
            last_dim = numpy_tensor.shape[-1]
            var = NeuralVariable(theano_tensor, dim=last_dim)
            var.set_test_value(numpy_tensor)
            vars.append(var)
        return vars

    @neural_computation
    def scan(self, func, sequences=None, outputs=None, non_sequences=None, block=None, **kwargs):
        """
        A loop function, the usage is identical with the theano one.
        :type block: deepy.layers.Block
        """
        results, updates = Scanner(func, sequences, outputs, non_sequences, neural_computation=True, **kwargs).compute()
        if block and updates:
            if type(updates) == dict:
                updates = updates.items()
            block.register_updates(*updates)
        return results

    def loop(self, sequences=None, outputs=None, non_sequences=None, block=None, **kwargs):
        """
        Start a loop.
        Usage:
        ```
        with deepy.graph.loop(sequences={"x": x}, outputs={"o": None}) as vars:
            vars.o = vars.x + 1
        loop_outputs = deepy.graph.loop_outputs()
        result = loop_outputs.o
        ```
        """
        from loop import Loop
        return Loop(sequences, outputs, non_sequences, block, **kwargs)

    def get_trainer(self, model,  method='sgd', config=None, annealer=None, validator=None):
        """
        Get a trainer to optimize given model.
        :rtype: deepy.trainers.GeneralNeuralTrainer
        """
        from deepy.trainers import GeneralNeuralTrainer
        return GeneralNeuralTrainer(model, method=method, config=config, annealer=annealer, validator=validator)

    @neural_computation
    def shared(self, value, name=None):
        """
        Create a shared theano scalar value.
        """
        if type(value) == int:
            final_value = np.array(value, dtype="int32")
        elif type(value) == float:
            final_value = np.array(value, dtype=env.FLOATX)
        else:
            final_value = value

        return theano.shared(final_value, name=name)

    @neural_computation
    def disconnect(self, x):
        """
        Disconnect a variable from backpropagation.
        """
        return disconnected_grad(x)

    def compile(self, input_dim=0, model=None, input_tensor=None, monitors=None,
                 cost=None, output=None, outputs=None, blocks=None, input_vars=None, target_vars=None):
        from comp_graph import ComputationalGraph
        # Pass the arguments to `ComputationalGraph`
        args = [arg for arg in getargspec(GraphBuilder.compile).args if arg != "self"]
        arg_vals = [locals()[k] for k in args]
        kwargs = dict(zip(args, arg_vals))
        return ComputationalGraph(**kwargs)

    def fill_parameters(self, path, blocks, exclude_free_params=False, check_parameters=False):
        """
        Load parameters from file to fill all blocks sequentially.
        :type blocks: list of deepy.layers.Block
        """
        if not os.path.exists(path):
            raise Exception("model {} does not exist".format(path))
        # Decide which parameters to load
        normal_params = sum([nn.parameters for nn in blocks], [])
        all_params = sum([nn.all_parameters for nn in blocks], [])
        # Load parameters
        if path.endswith(".gz"):
            opener = gzip.open if path.lower().endswith('.gz') else open
            handle = opener(path, 'rb')
            saved_params = pickle.load(handle)
            handle.close()
            # Write parameters
            if len(all_params) != len(saved_params):
                logging.warning(
                    "parameters in the network: {}, parameters in the dumped model: {}".format(len(all_params),
                                                                                len(saved_params)))
            for target, source in zip(all_params, saved_params):
                if not exclude_free_params or target not in normal_params:
                    target.set_value(source)
        elif path.endswith(".npz"):
            arrs = np.load(path)
            # Write parameters
            if len(all_params) != len(arrs.keys()):
                logging.warning(
                    "parameters in the network: {}, parameters in the dumped model: {}".format(len(all_params),
                                                                                len(arrs.keys())))
            for target, idx in zip(all_params, range(len(arrs.keys()))):
                if not exclude_free_params or target not in normal_params:
                    source = arrs['arr_%d' % idx]
                    target.set_value(source)
        else:
            raise Exception("File format of %s is not supported, use '.gz' or '.npz' or '.uncompressed.gz'" % path)

if "graph" not in globals():
    graph = GraphBuilder()