Completed
Push — master ( 48255b...bf2b0c )
by Raphael
01:13
created

RecurrentLayer.compute_step()   A

Complexity

Conditions 2

Size

Total Lines 15

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 2
c 0
b 0
f 0
dl 0
loc 15
rs 9.4285
1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
4
from . import NeuralLayer
5
from neural_var import NeuralVariable
6
from deepy.utils import build_activation, FLOATX, XavierGlorotInitializer, OrthogonalInitializer, Scanner, neural_computation
7
import numpy as np
8
import theano.tensor as T
9
from abc import ABCMeta, abstractmethod
10
11
OUTPUT_TYPES = ["sequence", "one"]
12
INPUT_TYPES = ["sequence", "one"]
13
14
15
16
class RecurrentLayer(NeuralLayer):
17
    __metaclass__ = ABCMeta
18
19
    def __init__(self, name, state_names, hidden_size=100, input_type="sequence", output_type="sequence",
20
                 inner_init=None, outer_init=None,
21
                 gate_activation='sigmoid', activation='tanh',
22
                 steps=None, backward=False, mask=None,
23
                 additional_input_dims=None):
24
        super(RecurrentLayer, self).__init__(name)
25
        self.state_names = state_names
26
        self.main_state = state_names[0]
27
        self.hidden_size = hidden_size
28
        self._gate_activation = gate_activation
29
        self._activation = activation
30
        self.gate_activate = build_activation(self._gate_activation)
31
        self.activate = build_activation(self._activation)
32
        self._input_type = input_type
33
        self._output_type = output_type
34
        self.inner_init = inner_init if inner_init else OrthogonalInitializer()
35
        self.outer_init = outer_init if outer_init else XavierGlorotInitializer()
36
        self._steps = steps
37
        self._mask = mask.tensor if type(mask) == NeuralVariable else mask
38
        self._go_backwards = backward
39
        self.additional_input_dims = additional_input_dims if additional_input_dims else []
40
41
        if input_type not in INPUT_TYPES:
42
            raise Exception("Input type of {} is wrong: {}".format(name, input_type))
43
        if output_type not in OUTPUT_TYPES:
44
            raise Exception("Output type of {} is wrong: {}".format(name, output_type))
45
46
    @neural_computation
47
    def step(self, step_inputs):
48
        new_states = self.compute_new_state(step_inputs)
49
50
        # apply mask for each step if `output_type` is 'one'
51
        if step_inputs.get("mask"):
52
            mask = step_inputs["mask"].dimshuffle(0, 'x')
53
            for state_name in new_states:
54
                new_states[state_name] = new_states[state_name] * mask + step_inputs[state_name] * (1 - mask)
55
56
        return new_states
57
58
    @abstractmethod
59
    def compute_new_state(self, step_inputs):
60
        """
61
        :type step_inputs: dict
62
        :rtype: dict
63
        """
64
65
    @abstractmethod
66
    def merge_inputs(self, input_var, additional_inputs=None):
67
        """
68
        Merge inputs and return a map, which will be passed to core_step.
69
        :type input_var: T.var
70
        :param additional_inputs: list
71
        :rtype: dict
72
        """
73
74
    @abstractmethod
75
    def prepare(self):
76
        pass
77
78
    @neural_computation
79
    def compute_step(self, state, lstm_cell=None, input=None, additional_inputs=None):
80
        """
81
        Compute one step in the RNN.
82
        :return: one variable for RNN and GRU, multiple variables for LSTM
83
        """
84
        input_map = self.merge_inputs(input, additional_inputs=additional_inputs)
85
        input_map.update({"state": state, "lstm_cell": lstm_cell})
86
        output_map = self.compute_new_state(input_map)
87
        outputs = [output_map.pop("state")]
88
        outputs += output_map.values()
89
        if len(outputs) == 1:
90
            return outputs[0]
91
        else:
92
            return outputs
93
94
    @neural_computation
95
    def get_initial_states(self, input_var):
96
        """
97
        :type input_var: T.var
98
        :rtype: dict
99
        """
100
        initial_states = {}
101
        for state in self.state_names:
102
            if self._input_type == 'sequence' and input_var.ndim == 2:
103
                init_state = T.alloc(np.cast[FLOATX](0.), self.hidden_size)
104
            else:
105
                init_state = T.alloc(np.cast[FLOATX](0.), input_var.shape[0], self.hidden_size)
106
            initial_states[state] = init_state
107
        return initial_states
108
109
    @neural_computation
110
    def get_step_inputs(self, input_var, states=None, mask=None, additional_inputs=None):
111
        """
112
        :type input_var: T.var
113
        :rtype: dict
114
        """
115
        step_inputs = {}
116
        if self._input_type == "sequence":
117
            if not additional_inputs:
118
                additional_inputs = []
119
            if mask:
120
                step_inputs['mask'] = mask.dimshuffle(1, 0)
121
            step_inputs.update(self.merge_inputs(input_var, additional_inputs=additional_inputs))
122
        else:
123
            # step_inputs["mask"] = mask.dimshuffle((1,0)) if mask else None
124
            if additional_inputs:
125
                step_inputs.update(self.merge_inputs(None, additional_inputs=additional_inputs))
126
        if states:
127
            for name in self.state_names:
128
                step_inputs[name] = states[name]
129
130
        return step_inputs
131
132
    def compute(self, input_var, mask=None, additional_inputs=None, steps=None, backward=False, init_states=None, return_all_states=False):
133
        if additional_inputs and not self.additional_input_dims:
134
            self.additional_input_dims = map(lambda var: var.dim(), additional_inputs)
135
        result_var = super(RecurrentLayer, self).compute(input_var,
136
                                                   mask=mask, additional_inputs=additional_inputs, steps=steps, backward=backward, init_states=init_states, return_all_states=return_all_states)
137
        if return_all_states:
138
            state_map = {}
139
            for k in result_var.tensor:
140
                state_map[k] = NeuralVariable(result_var.tensor[k], result_var.test_tensor[k], self.output_dim)
141
            return state_map
142
        else:
143
            return result_var
144
145
    def compute_tensor(self, input_var, mask=None, additional_inputs=None, steps=None, backward=False, init_states=None, return_all_states=False):
146
        # prepare parameters
147
        backward = backward if backward else self._go_backwards
148
        steps = steps if steps else self._steps
149
        mask = mask if mask else self._mask
150
        if mask and self._input_type == "one":
151
            raise Exception("Mask only works with sequence input")
152
        # get initial states
153
        init_state_map = self.get_initial_states(input_var)
154
        if init_states:
155
            for name, val in init_states.items():
156
                if name in init_state_map:
157
                    init_state_map[name] = val
158
        # get input sequence map
159
        if self._input_type == "sequence":
160
            # Move middle dimension to left-most position
161
            # (sequence, batch, value)
162
            if input_var.ndim == 3:
163
                input_var = input_var.dimshuffle((1,0,2))
164
165
            seq_map = self.get_step_inputs(input_var, mask=mask, additional_inputs=additional_inputs)
166
        else:
167
            init_state_map[self.main_state] = input_var
168
            seq_map = self.get_step_inputs(None, mask=mask, additional_inputs=additional_inputs)
169
        # scan
170
        retval_map, _ = Scanner(
171
            self.step,
172
            sequences=seq_map,
173
            outputs_info=init_state_map,
174
            n_steps=steps,
175
            go_backwards=backward
176
        ).compute()
177
        # return main states
178
        main_states = retval_map[self.main_state]
179
        if self._output_type == "one":
180
            if return_all_states:
181
                return_map = {}
182
                for name, val in retval_map.items():
183
                    return_map[name] = val[-1]
184
                return return_map
185
            else:
186
                return main_states[-1]
187
        elif self._output_type == "sequence":
188
            if return_all_states:
189
                return_map = {}
190
                for name, val in retval_map.items():
191
                    return_map[name] = val.dimshuffle((1,0,2))
192
                return return_map
193
            else:
194
                main_states = main_states.dimshuffle((1,0,2)) # ~ batch, time, size
195
                # if mask: # ~ batch, time
196
                #     main_states *= mask.dimshuffle((0, 1, 'x'))
197
                return main_states
198
199
200
class RNN(RecurrentLayer):
201
202
    def  __init__(self, hidden_size, **kwargs):
203
        kwargs["hidden_size"] = hidden_size
204
        super(RNN, self).__init__("RNN", ["state"], **kwargs)
205
206
    @neural_computation
207
    def compute_new_state(self, step_inputs):
208
        xh_t, h_tm1 = map(step_inputs.get, ["xh_t", "state"])
209
        if not xh_t:
210
            xh_t = 0
211
212
        h_t = self.activate(xh_t + T.dot(h_tm1, self.W_h) + self.b_h)
213
214
        return {"state": h_t}
215
216
    @neural_computation
217
    def merge_inputs(self, input_var, additional_inputs=None):
218
        if not additional_inputs:
219
            additional_inputs = []
220
        all_inputs = ([input_var] if input_var else []) + additional_inputs
221
        h_inputs = []
222
        for x, weights in zip(all_inputs, self.input_weights):
223
            wi, = weights
224
            h_inputs.append(T.dot(x, wi))
225
        merged_inputs = {
226
            "xh_t": sum(h_inputs)
227
        }
228
        return merged_inputs
229
230
    def prepare(self):
231
        self.output_dim = self.hidden_size
232
233
        self.W_h = self.create_weight(self.hidden_size, self.hidden_size, "h", initializer=self.outer_init)
234
        self.b_h = self.create_bias(self.hidden_size, "h")
235
236
        self.register_parameters(self.W_h, self.b_h)
237
238
        self.input_weights = []
239
        if self._input_type == "sequence":
240
            normal_input_dims = [self.input_dim]
241
        else:
242
            normal_input_dims = []
243
244
        all_input_dims = normal_input_dims + self.additional_input_dims
245
        for i, input_dim in enumerate(all_input_dims):
246
            wi = self.create_weight(input_dim, self.hidden_size, "wi_{}".format(i+1), initializer=self.outer_init)
247
            weights = [wi]
248
            self.input_weights.append(weights)
249
            self.register_parameters(*weights)