def grad(*dy): """Gradient function for build_rnn.""" dy = nest.pack_sequence_as(ret, dy) def _continue(unused_theta, unused_dy, unused_dstate1, unused_dtheta, unused_dinput, i): return _should_continue(i, True) dstate1, dtheta, dinput = tf.while_loop(_continue, _cell_grad_fn, [ theta, dy, { k: tf.zeros_like(state0[k]) for k in state0 if k not in skipped_state }, { k: tf.zeros_like(theta[k]) for k in theta if k not in skipped_theta }, {k: tf.zeros_like(inputs[k]) for k in inputs}, tf.zeros([], tf.int32) if reverse else max_length - 1, ])[2:5] dtheta, dinput = _cell_grad_fn_with_state0( state0, theta, dy, dstate1, dtheta, dinput, max_length - 1 if reverse else tf.zeros([], dtype=tf.int32))[3:5] state0_h = tf.reshape(acc_state["h"], [-1, theta["kernel"].shape[0]]) state0_atten = tf.reshape(acc_state["attention"], [ -1, theta["attention_kernel"].shape[0] ]) if "attention_kernel" in theta else None grad = tf.reshape(dinput["rnn"], [-1, theta["kernel"].shape[1]]) if reverse: state0_h = tf.split(state0_h, [batch_size, -1])[1] grad = tf.split(grad, [-1, batch_size])[0] else: if state0_atten is not None: state0_atten = tf.split(state0_atten, [-1, batch_size])[0] state0_h = tf.split(state0_h, [-1, batch_size])[0] grad = tf.split(grad, [batch_size, -1])[1] if state0_atten is not None: dtheta["attention_kernel"] = tf.matmul( tf.transpose(state0_atten), grad) dtheta["kernel"] = tf.matmul(tf.transpose(state0_h), grad) if "memory_kernel" in orig_theta: dtheta["memory_kernel"] = tf.zeros_like( orig_theta["memory_kernel"]) dtheta["seq_mask"] = tf.zeros_like(orig_theta["seq_mask"]) return dinput, dtheta
def tpu_eval_step(): """Generate the TPU graph.""" values = self.eval_infeed_queue[0].generate_dequeue_op( tpu_device=0) unflattened_inputs = data_nest.pack_sequence_as( self.eval_feature_structure, values) features = unflattened_inputs["features"] estimator_spec = model_fn(features, None, tf.estimator.ModeKeys.PREDICT, params) for k, v in six.iteritems(estimator_spec.predictions): self.outfeed_names.append(k) self.outfeed_tensors.append(v) with tf.device( device_for_tpu_core(get_host(self.resolver, self.hparams))): outfeed_enqueue_ops = tpu_ops.outfeed_enqueue_tuple( self.outfeed_tensors) with tf.control_dependencies([outfeed_enqueue_ops]): return tf.no_op()
def tpu_train_step(loss): """Generate the TPU graph.""" del loss values = self.infeed_queue[0].generate_dequeue_op(tpu_device=0) unflattened_inputs = data_nest.pack_sequence_as( self.feature_structure, values) features = unflattened_inputs["features"] core_id = unflattened_inputs["core_id"] new_features = {} for k in features: s = features[k].shape.as_list() s = [self.hparams.num_shards, s[0] // self.hparams.num_shards ] + s[1:] new_features[k] = tf.squeeze( tf.gather( tf.reshape(tpu_ops.cross_replica_sum(features[k]), s), core_id), [0]) estimator_spec = model_fn(new_features, None, tf.estimator.ModeKeys.TRAIN, params) loss, train_op = estimator_spec.loss, estimator_spec.train_op with tf.control_dependencies([train_op]): return tf.identity(loss)
def attention_cell_grad(theta, state0, unused_inputs, extras, dstate1): """Gradient function for attention_cell.""" new_lstm_state = lstm_cell_split(extras, state0["c"], None) new_states = attention(theta, new_lstm_state) del new_states["alignments"] y = nest.flatten(new_states) x = [extras, state0["c"]] + nest.flatten(theta) dy = nest.flatten(dstate1) g = tf.gradients(y, x, dy) dtheta = nest.pack_sequence_as(theta, g[2:]) grad, dstate_c = g[:2] dtheta["bias"] = tf.reduce_sum(grad, 0) datten = tf.matmul(grad, tf.transpose(theta["attention_kernel"])) dstate_h = tf.matmul(grad, tf.transpose(theta["kernel"])) dstate = { "h": dstate_h, "c": dstate_c, "attention": datten, } return dtheta, dstate, {"rnn": grad}
def _rnn(*inp): """Function that drives RNN with early stop.""" inputs = nest.pack_sequence_as(orig_inputs, inp[0:len(orig_inputs)]) theta = nest.pack_sequence_as(orig_theta, inp[len(orig_inputs):]) def _cell_fn(theta, state0, acc_state, acc_gate, i): """RNN cell function.""" input_slice = {k: tf.gather(inputs[k], i) for k in inputs} state1, gate = cell_fn(theta, state0, input_slice) for k in state0: if k not in skipped_state: acc_state[k] = tf.stop_gradient( inplace_ops.alias_inplace_update(acc_state[k], i, state1[k])) acc_gate = tf.stop_gradient( inplace_ops.alias_inplace_update(acc_gate, i, gate)) return theta, state1, acc_state, acc_gate, i - 1 if reverse else i + 1 def _should_continue(i, is_backward=False): if is_backward: return i < max_length - 1 if reverse else i > 0 else: return i >= 0 if reverse else i < max_length acc_state = { k: tf.zeros([max_time, batch_size, state0["c"].shape[-1]], state0["c"].dtype) for k in state0 if k not in skipped_state } acc_state, acc_gate = tf.while_loop( lambda theta, state0, acc_state, acc_gate, i: _should_continue(i), _cell_fn, [ theta, state0, acc_state, tf.zeros_like(inputs["rnn"]), max_length - 1 if reverse else tf.zeros([], tf.int32) ])[2:4] ret = {"h": acc_state["h"]} if "attention" in acc_state: ret["attention"] = acc_state["attention"] def _cell_grad_fn_with_state0(state0, theta, dy, dstate1, dtheta, dinput, i): """Gradient cell function.""" state0 = { k: tf.stop_gradient(state0[k]) for k in state0 if k not in skipped_state } theta = {k: tf.stop_gradient(theta[k]) for k in theta} if "padding" in inputs: inputs_slice = {"padding": tf.gather(inputs["padding"], i)} else: inputs_slice = None gate = tf.gather(acc_gate, i) for k in dy: dstate1[k] = dstate1[k] + tf.gather(dy[k], i) dt, dstate, di = cell_grad(theta, state0, inputs_slice, gate, dstate1) dtheta = {k: dtheta[k] + dt[k] for k in dtheta if k not in skipped_theta} dinput = { k: inplace_ops.alias_inplace_update(dinput[k], i, di[k]) for k in di } return theta, dy, dstate, dtheta, dinput, i + 1 if reverse else i - 1 def _cell_grad_fn(theta, dy, dstate1, dtheta, dinput, i): """Gradient cell function wrapper.""" return _cell_grad_fn_with_state0( { k: tf.gather(acc_state[k], i + 1 if reverse else i - 1) for k in acc_state }, theta, dy, dstate1, dtheta, dinput, i) def grad(*dy): """Gradient function for build_rnn.""" dy = nest.pack_sequence_as(ret, dy) def _continue(unused_theta, unused_dy, unused_dstate1, unused_dtheta, unused_dinput, i): return _should_continue(i, True) dstate1, dtheta, dinput = tf.while_loop(_continue, _cell_grad_fn, [ theta, dy, { k: tf.zeros_like(state0[k]) for k in state0 if k not in skipped_state }, {k: tf.zeros_like(theta[k]) for k in theta if k not in skipped_theta}, {k: tf.zeros_like(inputs[k]) for k in inputs}, tf.zeros([], tf.int32) if reverse else max_length - 1, ])[2:5] dtheta, dinput = _cell_grad_fn_with_state0( state0, theta, dy, dstate1, dtheta, dinput, max_length - 1 if reverse else tf.zeros([], dtype=tf.int32))[3:5] state0_h = tf.reshape(acc_state["h"], [-1, theta["kernel"].shape[0]]) state0_atten = tf.reshape(acc_state["attention"], [-1, theta["attention_kernel"].shape[0] ]) if "attention_kernel" in theta else None grad = tf.reshape(dinput["rnn"], [-1, theta["kernel"].shape[1]]) if reverse: state0_h = tf.split(state0_h, [batch_size, -1])[1] grad = tf.split(grad, [-1, batch_size])[0] else: if state0_atten is not None: state0_atten = tf.split(state0_atten, [-1, batch_size])[0] state0_h = tf.split(state0_h, [-1, batch_size])[0] grad = tf.split(grad, [batch_size, -1])[1] if state0_atten is not None: dtheta["attention_kernel"] = tf.matmul(tf.transpose(state0_atten), grad) dtheta["kernel"] = tf.matmul(tf.transpose(state0_h), grad) if "memory_kernel" in orig_theta: dtheta["memory_kernel"] = tf.zeros_like(orig_theta["memory_kernel"]) dtheta["seq_mask"] = tf.zeros_like(orig_theta["seq_mask"]) return dinput, dtheta return ret, grad