def diagonal_lstm(inputs, scope='diagonal_lstm'): with tf.compat.v1.variable_scope(scope): skewed_inputs = skew(inputs, scope="skewed_i") input_to_state = conv2d(skewed_inputs, 64, [1, 1], mask_type="b", scope="i_to_s") column_wise_inputs = tf.transpose(input_to_state, perm=[0, 2, 1, 3]) batch, width, height, channel = column_wise_inputs.get_shape().as_list( ) rnn_inputs = tf.reshape(column_wise_inputs, [-1, width, height * channel]) cell = DiagonalLSTMCell(16, height, channel) outputs, states = tf.compat.v1.nn.dynamic_rnn(cell, inputs=rnn_inputs, dtype=tf.float32) width_first_outputs = tf.reshape(outputs, [-1, width, height, 16]) skewed_outputs = tf.transpose(width_first_outputs, perm=[0, 2, 1, 3]) outputs = unskew(skewed_outputs) return outputs
def __call__(self, i_to_s, state, scope="DiagonalBiLSTMCell"): c_prev = tf.slice(state, [0, 0], [-1, self._num_units]) h_prev = tf.slice(state, [0, self._num_units], [-1, self._num_units]) with tf.compat.v1.variable_scope(scope): conv1d_inputs = tf.reshape( h_prev, [-1, self._height, 1, self._hidden_dims], name='conv1d_inputs') conv_s_to_s = conv1d(conv1d_inputs, 4 * self._hidden_dims, 2, scope='s_to_s') s_to_s = tf.reshape(conv_s_to_s, [-1, self._height * self._hidden_dims * 4]) lstm_matrix = tf.sigmoid(s_to_s + i_to_s) i, g, f, o = tf.split(lstm_matrix, 4, 1) c = f * c_prev + i * g h = tf.multiply(o, tf.tanh(c), name='hid') new_state = tf.concat([c, h], 1) return h, new_state
def skew(inputs, scope="skew"): with tf.compat.v1.name_scope(scope): batch, height, width, channel = inputs.get_shape().as_list() rows = tf.split(inputs, height, 1) new_width = width + height - 1 new_rows = [] for idx, row in enumerate(rows): transposed_row = tf.transpose(tf.squeeze(row, [1]), perm=[0, 2, 1]) squeezed_row = tf.reshape(transposed_row, [-1, width]) padded_row = tf.pad(squeezed_row, paddings=((0, 0), (idx, height - 1 - idx))) unsqueezed_row = tf.reshape(padded_row, [-1, channel, new_width]) untransposed_row = tf.transpose(unsqueezed_row, perm=[0, 2, 1]) new_rows.append(untransposed_row) outputs = tf.stack(new_rows, axis=1, name="output") return outputs