layers.py 文件源码

def _compute_states(self):

    _inputs = tf.transpose(self.inputs, [1, 0, 2])
    x_ta = tf.TensorArray(tf.float32, size=self.length).unstack(_inputs)

    h_ta_size = self.num_initial_states + self.length
    initial_states = tf.transpose(self.initial_states, [1, 0, 2])

    # infer_shapes=True is buggy and says that shape (?, num_hidden_units) is incompatible with
    # shape (?, num_hidden_units). I've verified that they both have shape
    # (batch_size, num_hidden_units). To avoid this,  we'll set infer_shape=False and
    # skip the consistency check entirely.
    h_ta = tf.TensorArray(tf.float32, size=h_ta_size, clear_after_read=False, infer_shape=False)
    h_ta = h_ta.unstack(initial_states)

    def cond(t, h_ta):
      return tf.less(t, self.length)

    def body(t, h_ta):

      h = h_ta.read(self.num_initial_states + t - 1)
      x = x_ta.read(t)
      num_units, input_size = self.num_hidden_units, self.input_size

      with tf.variable_scope('pre_act'):

        # Shape [batch_size, pre_act_mixture_delays.size, num_units]
        h_history = tf.transpose(h_ta.gather(self.num_initial_states + t - self.pre_act_mixture_delays), [1, 0, 2])

        # Shape [batch_size, pre_act_mixture_delays.size, 1]
        coefs = tf.expand_dims(self._linear(h, x, self.pre_act_mixture_delays.size, scope='coefs'), 2)
        coefs = tf.nn.softmax(coefs, dim=1)

        # Shape [batch_size, num_units]
        h_pre_act = tf.reduce_sum(coefs * h_history, axis=[1])

        r = tf.nn.sigmoid(self._linear(h, x, num_units, scope='r'))
        h_pre_act = r * h_pre_act

      h_tilde = self.activation(self._linear(h_pre_act, x, num_units, scope='mist'))

      h_ta_new = h_ta.write(self.num_initial_states + t, h_tilde)
      return t + 1, h_ta_new

    t = tf.constant(0)
    _, h_ta = tf.while_loop(cond, body, [t, h_ta])

    all_states = h_ta.stack()
    states = tf.transpose(all_states[self.num_initial_states:], [1, 0, 2], name='states')
    outputs = tf.identity(states, name='outputs')
    return outputs, states