def eval_loop(): """Docs.""" def _cond(step, *args): # pylint: disable=unused-argument return tf.less(step, tf.cast(num_eval_steps, step.dtype)) def _body(step, *args): outs = training_utils.eval_step_fn(params, model) new_args = [step + 1 ] + [a.write(step, o) for a, o in zip(args, outs)] return tuple(new_args) batch_size = batch_size = params.eval_batch_size // params.num_replicas num_classes = params.num_classes logits = tf.TensorArray(dtype=tf.float32, size=num_eval_steps, element_shape=[batch_size, num_classes]) labels = tf.TensorArray(dtype=tf.int32, size=num_eval_steps, element_shape=[batch_size, 1]) masks = tf.TensorArray(dtype=tf.float32, size=num_eval_steps, element_shape=[batch_size]) loop_inps = [0, logits, labels, masks] loop_outs = tf.while_loop(_cond, _body, loop_inps, parallel_iterations=1, name='eval') return [o.concat() for o in loop_outs[1:]]
def _scan_initial_state(self): """Create TensorArrays and indices to track bin assignment. availability: TensorArray[queue_size, num_sequences] This represents the number of tokens available in the ith bin. See implementation note below. contents: TensorArray[queue_size, num_sequences * 2] This holds the actual contents of the packed strings as well as a bit mask indicating where sequences begin. It is stored in a flat vector and is accessed in offsets of packed_length. top_index: scalar [0, queue_size) Integer tensor indicating which index is the "top" bin. See implementation note below. IMPLEMENTATION_NOTE: The FFD algorithm periodically pops the topmost queue and pushes a new one to replace it. In order to replicate those semantics with a fixed size TensorArray, indexing operations are shifted by top_index. For example, instead of: `queue_available.read(i)` a read is instead performed as: `queue_available.read((i - top_index) % queue_size)` to account for the fact that the "ith" logical FFD queue is stored at position j. This means that the pop / push update can be performed by simply incrementing top_index. (And zeroing the old top_index position.) Returns: The state for the binning scan. """ all_available = ( tf.ones((self._queue_size, self._num_sequences), dtype=INDEX_DTYPE) * self._packed_length ) total_size = self._packed_length * self._queue_size total_size_range = tf.range(total_size, dtype=INDEX_DTYPE) empty = tf.zeros((total_size, self._num_sequences * 2), dtype=self._token_dtype) availability = tf.TensorArray( dtype=INDEX_DTYPE, size=self._queue_size, dynamic_size=False, clear_after_read=False, element_shape=(self._num_sequences,), ).scatter(tf.range(self._queue_size, dtype=INDEX_DTYPE), all_available) contents = tf.TensorArray( dtype=self._token_dtype, size=total_size, dynamic_size=False, clear_after_read=False, element_shape=(self._num_sequences * 2,), ).scatter(total_size_range, empty) # Which index should be considered the "top" bucket for the purpose of # the first-fit descending algorithm. top_index = tf.zeros((), dtype=INDEX_DTYPE) return availability, contents, top_index
def body_outer(i, rows): """Calculates the i-th row of the triangle (not strictly). The shape of the row should be [T-1, state_size] """ assert isinstance(rows, tf.TensorArray) # Initialize grad grad_init = dlds[i] assert isinstance(grad_init, tf.Tensor) and len( grad_init.shape) == 1 grad_init = tf.reshape(grad_init, [-1, 1]) def body_inner(j, row, grad): assert isinstance(row, tf.TensorArray) grad_to_write = tf.transpose(grad, [1, 0]) # Write grad to row row = row.write( j - 1, tf.cond(tf.less_equal(j, i), lambda: grad_to_write, lambda: tf.zeros_like(grad_to_write))) # Update grad and return grad = tf.matmul(dsds[j - 1], grad) return j - 1, row, grad # Use tf.while_loop to generate row empty_row = tf.TensorArray(tf.float32, size=T - 1, dynamic_size=False) _, row, _ = tf.while_loop(lambda j, *_: tf.greater(j, 0), body_inner, (T - 1, empty_row, grad_init), back_prop=False) # each entry in row has a shape [1, state_size] assert isinstance(row, tf.TensorArray) return i + 1, rows.write(i - 1, row.concat(name='row'))
def align(encoder_states, decoder_hidden_state,scope="attention"): with tf.variable_scope(scope,reuse=tf.AUTO_REUSE): Wp = tf.get_variable("Wp", shape=[2*hidden_size,128],# pt = S·sigmoid(vp tanh(Wp.ht)) #ht shape [32,600] dtype=tf.float32, trainable=True, initializer=tf.glorot_uniform_initializer()) Vp = tf.get_variable("Vp", shape=[128,1], dtype=tf.float32, trainable=True, initializer=tf.glorot_uniform_initializer()) positions = tf.cast(S-window_len,dtype=tf.float32) # Maximum valid attention window starting position # Predict attention window starting position sigmoid_values = tf.nn.sigmoid(tf.matmul(tf.tanh(tf.matmul(decoder_hidden_state,Wp)),Vp)) ps = positions*sigmoid_values #tf.nn.sigmoid(tf.matmul(tf.tanh(tf.matmul(decoder_hidden_state,Wp)),Vp)) # ps = (soft-)predicted starting position of attention window pt = ps+D # pt = center of attention window where the whole window length is 2*D+1 pt = tf.reshape(pt,[N]) # pt is the point of interest i.e is center next step is to create a gausian dis # tribution centered around pt whose size is sentence length so point near to pt will be given more priority i = 0 gaussian_position_based_scores = tf.TensorArray(size=S,dtype=tf.float32) sigma = tf.constant(D/2,dtype=tf.float32)
def _sandwich_bottom(self, dlds, dsds): assert isinstance(dlds, tf.Tensor) and isinstance(dsds, tf.Tensor) assert len(dlds.shape) == 2 and len(dsds.shape) == 3 Ts = hub.error_injection_step dlds, dsds = dlds[:Ts], dsds[:Ts] Ts = tf.shape(dlds)[0] def body(tau, dLtdS, grad): assert isinstance(dLtdS, tf.TensorArray) grad_to_write = tf.transpose(grad, [1, 0]) # Write gradient dLtdS = dLtdS.write(tau, grad_to_write) # Update grad grad = tf.matmul(dsds[tau], grad) return tau - 1, dLtdS, grad # dLt/dS will be put into tensor array in while_loop ta = tf.TensorArray(tf.float32, size=Ts - 1) grad_src = tf.reshape(dlds[-1], [-1, 1]) _, dLtdS, _ = tf.while_loop(lambda tau, *_: tf.greater_equal(tau, 0), body, (Ts - 2, ta, grad_src), back_prop=False) assert isinstance(dLtdS, tf.TensorArray) # each entry in row has a shape [1, state_size] dLtdS = dLtdS.concat('dLt_dS') return dLtdS
def batch_image_files_decode(image_files): raw_images = tf.TensorArray(tf.uint8, size=0, dynamic_size=True) for i in tf.range(tf.shape(image_files)[0]): image = tf.io.decode_image(image_files[i]) image.set_shape([None, None, None]) raw_images = raw_images.write(i, image) return raw_images.stack()
def sample(self, num_samples=1): sample_ta = tf.TensorArray(dtype=tf.float32, size=self.data_dim, dynamic_size=False, clear_after_read=False, element_shape=[num_samples]).unstack( tf.zeros([self.data_dim, num_samples])) def while_body(i, sample_ta): contexts = self.arnn_net( tf.transpose(sample_ta.stack()) - self.data_mean) # [num_samples, context_dim] context_i = contexts[:, i, :] logits_i = tf.reshape(self.enn_net(context_i), [num_samples]) p = tfd.Bernoulli(logits=logits_i) sample = tf.to_float(p.sample()) sample_ta = sample_ta.write(i, sample) return i + 1, sample_ta def while_cond(i, unused_ta): return i < self.data_dim outs = tf.while_loop(while_cond, while_body, (0, sample_ta), back_prop=False) return tf.transpose(outs[1].stack())
def my_dynamic_rnn(rnn_cell, input_data, initial_state, sequence_length=None): """A handwritten reimplementation of dynamic_rnn.""" input_data = tf.transpose(input_data, [1, 0, 2]) outputs = tf.TensorArray(tf.float32, input_data.shape[0]) if sequence_length is None: max_seq_len = input_data.shape[0] else: max_seq_len = tf.reduce_max(sequence_length) def while_body(i, state, outputs): new_output, new_state = rnn_cell(input_data[i], state) output = tf.where(i < sequence_length, new_output, tf.zeros(new_output.shape)) state = tf.where(i < sequence_length, new_state, state) outputs = outputs.write(i, output) return i + 1, state, outputs def while_cond(i, unused_state, unused_outputs): return i < max_seq_len _, state, outputs = tf.while_loop(while_cond, while_body, loop_vars=(tf.constant(0), initial_state, outputs)) return tf.transpose(outputs.stack(), [1, 0, 2]), state
def est_hess_trace_while2(y, x, num_v=1): batch_size, data_dim = x.get_shape().as_list() v_dist = tfd.Normal(loc=tf.zeros([data_dim]), scale=tf.ones([data_dim])) # [num_v, data_dim] v = v_dist.sample(num_v) # [batch_size, data_dim] score = tf.gradients(y, x)[0] # [batch_size, num_v] vt_score = tf.linalg.matmul(score, v, transpose_b=True) grad_ta = tf.TensorArray(tf.float32, size=num_v, dynamic_size=False, element_shape=x.shape) def while_body(n, grad_ta): vt_score_i = vt_score[:, n] grad = tf.gradients(vt_score_i, x)[0] grad_ta = grad_ta.write(n, grad) return n + 1, grad_ta def while_predicate(n, unused_grad_ta): return n < num_v while_outs = tf.while_loop(while_predicate, while_body, (0, grad_ta), parallel_iterations=num_v) # [num_v, batch_size, data_dim] grads = while_outs[1].stack() estimate = tf.linalg.matmul(grads, v[:, :, tf.newaxis]) estimate = tf.reduce_mean(estimate, axis=0) return estimate
def for_compile_fn(*infeed_args): def _iterate_compile(i, sample, *args): if isinstance(sample, (list, tuple)): arqs = list(args) + list(sample) out = body(*arqs) else: out = body(*args, **sample) if isinstance(out, (list, tuple)): return (tf.add(i, 1), ) + tuple(out) else: return tf.add(i, 1), out batch_size = infeed_args[0].get_shape()[0] // n superargs_rs = [ tf.reshape(t, [n, batch_size] + t.get_shape().as_list()[1:]) for t in infeed_args ] superargs_arr = [ tf.TensorArray(t.dtype, n).unstack(t) for t in superargs_rs ] def iterate_compile(i, *subargs): infeed_batches = [ta.read(i) for ta in superargs_arr] return _iterate_compile(i, infeed_batches, *subargs) return tf.while_loop( cond=lambda i, *args, **kwargs: tf.less(i, n), body=iterate_compile, loop_vars=[tf.constant(0, dtype=tf.int32)] + list(inputs), parallel_iterations=1, back_prop=True)[1:]
def nms_tf(dets, thresh): """Non-maximum suppression with tf graph mode.""" x1 = dets[:, 0] y1 = dets[:, 1] x2 = dets[:, 2] y2 = dets[:, 3] scores = dets[:, 4] areas = (x2 - x1 + 1) * (y2 - y1 + 1) order = tf.argsort(scores, direction='DESCENDING') keep = tf.TensorArray(tf.int32, size=0, dynamic_size=True) index = 0 while tf.shape(order)[0] > 0: i = order[0] keep = keep.write(index, i) xx1 = tf.maximum(x1[i], tf.gather(x1, order[1:])) yy1 = tf.maximum(y1[i], tf.gather(y1, order[1:])) xx2 = tf.minimum(x2[i], tf.gather(x2, order[1:])) yy2 = tf.minimum(y2[i], tf.gather(y2, order[1:])) w = tf.maximum(0.0, xx2 - xx1 + 1) h = tf.maximum(0.0, yy2 - yy1 + 1) intersection = w * h overlap = intersection / ( areas[i] + tf.gather(areas, order[1:]) - intersection) inds = tf.where_v2(overlap <= thresh) order = tf.concat(tf.gather(order, inds + 1), axis=1) order = tf.squeeze(order, axis=-1) index += 1 return keep.stack()
def _lstm(x, prev_c, prev_h, w_lstm, layer_masks): """Multi-layer LSTM. Args: x: [batch_size, num_steps, hidden_size]. prev_c: [[batch_size, hidden_size] * num_layers]. prev_h: [[batch_size, hidden_size] * num_layers]. w_lstm: [[2 * hidden_size, 4 * hidden_size] * num_layers]. layer_masks: [([hidden_size, hidden_size] or None)* num_layers]. Returns: next_c: [[batch_size, hidden_size] * num_layers]. next_h: [[batch_size, hidden_size] * num_layers]. all_h: [batch_size, num_steps, hidden_size]. """ _, num_steps, _ = tf.unstack(tf.shape(x)) num_layers = len(w_lstm) all_h = [ tf.TensorArray(dtype=tf.float32, size=num_steps, infer_shape=False) for _ in range(num_layers) ] def _condition(step, *unused_args): return tf.less(step, num_steps) def _body(step, pprev_c, pprev_h, all_h): """Apply LSTM at each step.""" next_c, next_h = [], [] for layer_id, (p_c, p_h, w, m) in enumerate( zip(pprev_c, pprev_h, w_lstm, layer_masks)): inp = x[:, step, :] if layer_id == 0 else next_h[-1] if m is not None: inp *= m ifog = tf.matmul(tf.concat([inp, p_h], axis=1), w) i, f, o, g = tf.split(ifog, 4, axis=1) i = tf.sigmoid(i) f = tf.sigmoid(f) o = tf.sigmoid(o) g = tf.tanh(g) c = i * g + f * p_c h = o * tf.tanh(c) all_h[layer_id] = all_h[layer_id].write(step, h) next_c.append(c) next_h.append(h) return step + 1, next_c, next_h, all_h loop_inps = [tf.constant(0, dtype=tf.int32), prev_c, prev_h, all_h] _, next_c, next_h, all_h = tf.while_loop(_condition, _body, loop_inps, parallel_iterations=1) all_h = [tf.transpose(h.stack(), [1, 0, 2]) for h in all_h] return next_c, next_h, all_h
def langevin(energy_fn, init_X, step_size=1.0, burn_in=100, num_samples=1000, thinning_steps=1, max_steps=None): samples = tf.TensorArray(init_X.dtype, size=num_samples * thinning_steps, dynamic_size=False, name='samples_ta') #init_X = init_X[tf.newaxis,:] X_shape = tf.shape(init_X) if max_steps == None: max_steps = 1000 * num_samples * thinning_steps def langevin_step(i, num_accepted, X, samples): # Compute initial kinetic and potential energies. grad_X = tf.gradients(energy_fn(X), X)[0] new_X = X + step_size * grad_X + tf.math.sqrt( 2 * step_size) * tf.random.normal(X_shape) new_X = tf.debugging.check_numerics(new_X, "Nans in X.") accept = tf.squeeze(i > burn_in) samples = tf.cond(accept, lambda: samples.write(num_accepted, new_X), lambda: samples) new_X = tf.cond(accept, lambda: new_X, lambda: X) return i + 1, num_accepted + tf.to_int32(accept), new_X, samples def langevin_predicate(i, num_accepted, unused_X, unused_samples): # We're not done as long as we haven't gone over the max number of steps # and we haven't accepted enough samples. return tf.logical_and( tf.less(i, burn_in + max_steps), tf.less(num_accepted, num_samples * thinning_steps)) results = tf.while_loop(langevin_predicate, langevin_step, (0, 0, init_X, samples), back_prop=False) #[num_samples, data_dim] samples = results[-1].stack() samples = tf.reshape(samples, [num_samples, thinning_steps, -1]) samples = samples[:, -1, :] num_steps = results[0] num_accepted = results[1] accept_ratio = num_accepted / (num_steps - burn_in) #tf.summary.scalar("acceptance_ratio", accept_ratio) #tf.summary.scalar("num_langevin_steps", num_steps - burn_in) return samples
def run_hmc(self, z_sample_init, unnorm_log_prob_target): # # Standard normal prior # p_z = tfd.MultivariateNormalDiag(tf.zeros_like(z_sample_init), tf.ones_like(z_sample_init)) # with tf.variable_scope('hmc', reuse=tf.AUTO_REUSE, use_resource=True): hmc_step_size_tr = tf.get_variable( name='hmc_step_size_tr', shape=(), initializer=tf.constant_initializer(self.hmc_step_size_init), dtype=self.dtype, use_resource=True, trainable=False) if self.testing: if self.init_test_mcmc_train_step_size: hmc_step_size = hmc_step_size_tr else: hmc_step_size = self.hmc_step_size_init else: hmc_step_size = hmc_step_size_tr self.target_log_prob = unnorm_log_prob_target if self.keep_samples: z_samples = tf.TensorArray(dtype=self.dtype, size=self.n_burn_in_steps + self.n_post_burn_steps, name='hmc_keep_samples_arr') else: z_samples = tf.zeros((), dtype=self.dtype) _, z_t, z_samples_arr, new_step_size = \ tf.while_loop(cond=self.hmc_cond, body=self.hmc_step, loop_vars=[tf.constant(0), z_sample_init, z_samples, hmc_step_size], back_prop=False, maximum_iterations=self.n_burn_in_steps + self.n_post_burn_steps, name='hmc_while_loop') hmc_step_size_op = tf.no_op() if self.testing else tf.assign( hmc_step_size, new_step_size) if self.keep_samples: z_samples = tf.reshape( z_samples_arr.stack(), (self.n_burn_in_steps + self.n_post_burn_steps, tf.shape(z_sample_init)[0], tf.shape(z_sample_init)[-1]))[-self.n_post_burn_steps:] return z_samples, hmc_step_size_op else: return z_t, hmc_step_size_op
def rollout(simulator, features, num_steps): """Rolls out a trajectory by applying the model in sequence.""" initial_positions = features['position'][:, 0:INPUT_SEQUENCE_LENGTH] ground_truth_positions = features['position'][:, INPUT_SEQUENCE_LENGTH:] global_context = features.get('step_context') def step_fn(step, current_positions, predictions): if global_context is None: global_context_step = None else: global_context_step = global_context[step + INPUT_SEQUENCE_LENGTH - 1][tf.newaxis] next_position = simulator( current_positions, n_particles_per_example=features['n_particles_per_example'], particle_types=features['particle_type'], global_context=global_context_step) # Update kinematic particles from prescribed trajectory. kinematic_mask = get_kinematic_mask(features['particle_type']) next_position_ground_truth = ground_truth_positions[:, step] next_position = tf.where(kinematic_mask, next_position_ground_truth, next_position) updated_predictions = predictions.write(step, next_position) # Shift `current_positions`, removing the oldest position in the sequence # and appending the next position at the end. next_positions = tf.concat( [current_positions[:, 1:], next_position[:, tf.newaxis]], axis=1) return (step + 1, next_positions, updated_predictions) predictions = tf.TensorArray(size=num_steps, dtype=tf.float32) _, _, predictions = tf.while_loop( cond=lambda step, state, prediction: tf.less(step, num_steps), body=step_fn, loop_vars=(0, initial_positions, predictions), back_prop=False, parallel_iterations=1) output_dict = { 'initial_positions': tf.transpose(initial_positions, [1, 0, 2]), 'predicted_rollout': predictions.stack(), 'ground_truth_rollout': tf.transpose(ground_truth_positions, [1, 0, 2]), 'particle_types': features['particle_type'], } if global_context is not None: output_dict['global_context'] = global_context return output_dict
def _sample2(self): sample_ta0 = tf.TensorArray(self.data_mean.dtype, size=self.data_dim, dynamic_size=False, element_shape=[1], clear_after_read=False, name="samples_ta") def sample_body(i, sample_ta): # [1, i] prev_pixels = tf.transpose(sample_ta.gather(tf.range(0, i))) # [1, data_dim] arnn_input = tf.pad(prev_pixels, [[0, 0], [0, self.data_dim - i]], "CONSTANT") arnn_input = arnn_input - self.data_mean[tf.newaxis, :] # [1, data_dim, context_dim] contexts = self.arnn_net(arnn_input) # [1, context_dim] context = contexts[:, i, :] # [num_x_samples, context_dim] tiled_context = tf.tile(context, [self.num_x_samples, 1]) # [num_x_samples, 1] xs = tf.range(0, 1.0, 1. / self.num_x_samples)[:, tf.newaxis] centered_xs = xs - self.data_mean[i] # [num_x_samples, context_dim+1] enn_input = tf.concat([centered_xs, tiled_context], axis=-1) # [num_x_samples] log_energies = self.enn_net(enn_input) # [1] log_Z_hat = (tf.math.reduce_logsumexp(log_energies, axis=0) - tf.log(tf.to_float(self.num_x_samples))) # [num_x_samples] weights = log_energies - log_Z_hat # [1] ind = tfd.Categorical(probs=weights).sample() new_sample = tf.reshape(tf.gather(tf.squeeze(xs), ind), [1]) sample_ta = sample_ta.write(i, new_sample) return i + 1, sample_ta def sample_predicate(i, unused_sample_ta): return tf.less(i, self.data_dim) results = tf.while_loop(sample_predicate, sample_body, (0, sample_ta0), back_prop=False) sample = tf.reshape(results[1].stack(), [self.data_dim]) return sample
def backtrack_alignment_tf(bt_ta, input_lengths, label_lengths, blank_index): """Computes the alignment from the backtracking matrix. :param tf.TensorArray bt_ta: [T+U] * (B, U, 2) :param tf.Tensor input_lengths: (B,) :param tf.Tensor label_lengths: (B,) :param int blank_index: """ max_path = bt_ta.size() n_batch = tf.shape(input_lengths)[0] alignments_ta = tf.TensorArray( dtype=tf.int32, size=max_path, dynamic_size=False, infer_shape=False, element_shape=(None, ), # [V] name="alignments") initial_idx = tf.zeros((n_batch, 2), dtype=tf.int32) def body(s, alignments, idx): """Runs over s=[max_path-1,..., 1] (along the alignment) :param int s: path index :param tf.TensorArray alignments: [T+U] * (V,) :param tf.Tensor idx: (B, 2) -> [from_u, symbol/blank] """ backtrack = bt_ta.read(s) # (B, U, 2) init_u = tf.where( tf.greater_equal(s, input_lengths + label_lengths - 1), label_lengths, # if we are at the end of some path (or behind) idx[:, 0] # if we are within a path, continue backtracking. ) backtrack_indices = tf.stack([tf.range(n_batch), init_u], axis=-1) # (B, 2) idx = tf.gather_nd(backtrack, backtrack_indices) align_write = tf.where( tf.less_equal(s, input_lengths + label_lengths), idx[:, 1], # within alignment tf.ones((n_batch, ), dtype=tf.int32) * blank_index) # outside, assume blank alignments = alignments.write(s, align_write) return s - 1, alignments, idx init_s = max_path - 1 final_s, final_alignments_ta, final_idx = tf.while_loop( lambda s, *args: tf.greater_equal(s, 1), body, (init_s, alignments_ta, initial_idx)) final_alignments_ta = final_alignments_ta.write( 0, tf.tile([blank_index], [n_batch])) return tf.transpose(final_alignments_ta.stack())
def _sample(self): sample_ta0 = tf.TensorArray(self.data_mean.dtype, size=self.data_dim, dynamic_size=False, element_shape=[1], clear_after_read=False, name="samples_ta") def sample_body(i, sample_ta): # [1, i] prev_pixels = tf.transpose(sample_ta.gather(tf.range(0, i))) # [1, data_dim] arnn_input = tf.pad(prev_pixels, [[0, 0], [0, self.data_dim - i]], "CONSTANT") arnn_input = arnn_input - self.data_mean[tf.newaxis, :] # [1, data_dim, context_dim] contexts = self.arnn_net(arnn_input) # [1, context_dim] context = contexts[:, i, :] def le(x): centered_x = x - self.data_mean[i] # [1, context_dim+1] enn_input = tf.concat([centered_x, context], axis=-1) # [1] energy = self.enn_net(enn_input) return -energy # [1, 1] new_samples = utils.langevin(le, tf.reshape(self.data_mean[i], [1]), step_size=0.1, num_samples=1, burn_in=100, thinning_steps=100) new_samples = tf.reshape(new_samples, [1]) sample_ta = sample_ta.write(i, new_samples) return i + 1, sample_ta def sample_predicate(i, unused_sample_ta): return tf.less(i, self.data_dim) results = tf.while_loop(sample_predicate, sample_body, (0, sample_ta0), back_prop=False) sample = tf.reshape(results[1].stack(), [self.data_dim]) return sample
def batching_loop_initial_values(self): """Returns a list of initial values for the batching tf.while_loop. The return value here defines the shepherd-specific loop vars for a tf.while_loop that does dynamic batching. Updated values for these loop vars will get passed in to `batching_loop_body` and `batching_loop_results`. Returns: A possibly nested list of tensors. """ source_offset = tf.constant(0, dtype=tf.int64) dest_offset = tf.constant(0, dtype=tf.int64) source_indices_ta = tf.TensorArray(tf.int64, size=_TENSORARRAY_INITIAL_SIZE, dynamic_size=True, element_shape=None, infer_shape=False) dest_indices_ta = tf.TensorArray(tf.int64, size=_TENSORARRAY_INITIAL_SIZE, dynamic_size=True, element_shape=None, infer_shape=False) return _SparseTensorLoopVars(source_offset, dest_offset, source_indices_ta, dest_indices_ta)
def mask_joint_logits(input_mask, start_end_logits): """Masks logits based on input mask and valid start/end combinations.""" _, _, length = modeling.get_shape_list(start_end_logits, expected_rank=3) mask = tf.TensorArray(input_mask.dtype, size=length, dynamic_size=False) for i in range(length): mask = mask.write(i, input_mask) # The permitted span length is determined by the existing mask combined # with its being shifted up by one. input_mask = input_mask * tf.pad(input_mask[:, 1:], [[0, 0], [0, 1]]) mask = mask.stack() mask = tf.transpose(mask, [1, 2, 0]) mask.shape.assert_is_compatible_with(start_end_logits.shape) start_end_logits -= 1e6 * tf.cast(1 - mask, tf.float32) return start_end_logits
def batching_loop_initial_values(self): """Initial values for all loop variables this shepherd is responsible for. This is method 1 of 3 that shepherds need to implement in order to be able to dynamically batch data. It is called before the start of the tf.while_loop that implements the batching. Returns: A list of TensorFlow Variables defining the tf.while_loop state managed by the shepherd. """ tensor_array = tf.TensorArray(self._dtype, size=0, dynamic_size=True, element_shape=None, infer_shape=False) return [tensor_array]
def batching_loop_initial_values(self): """Initial values for all loop variables this shepherd is responsible for. This is method 1 of 3 that shepherds need to implement in order to be able to dynamically batch data. It is called before the start of the tf.while_loop that implements the batching. Returns: A list of TensorFlow Variables defining the tf.while_loop state managed by the shepherd. """ # pylint: disable=protected-access tensor_array = tf.TensorArray(tf.string, size=_TENSORARRAY_INITIAL_SIZE, dynamic_size=True, element_shape=None, infer_shape=False) return [tensor_array]
def sample(self, num_samples=1, num_importance_samples=20): sample_ta = tf.TensorArray(dtype=tf.float32, size=self.data_dim, dynamic_size=False, clear_after_read=False, element_shape=[num_samples]).unstack( tf.zeros([self.data_dim, num_samples])) def while_body(i, sample_ta): contexts, q = self.arnn(tf.transpose(sample_ta.stack())) # [num_samples, context_dim] context_i = contexts[:, i, :] # [num_importance_samples, num_samples] sample = q.sample(num_importance_samples)[:, :, i] centered_sample = sample - self.data_mean[i] # [num_importance_samples, num_samples, context_dim] tiled_contexts = tf.tile(context_i[tf.newaxis, :, :], [num_importance_samples, 1, 1]) # [num_importance_samples, num_samples, context_dim+1] enn_input = tf.concat( [centered_sample[:, :, tf.newaxis], tiled_contexts], axis=-1) # [num_importance_samples, num_samples] log_energies = self.enn_net(enn_input) # [num_samples] log_Z_hat = (tf.math.reduce_logsumexp(log_energies, axis=0) - tf.log(tf.to_float(num_importance_samples))) # [num_importance_samples, num_samples] weights = log_energies - log_Z_hat[tf.newaxis, :] inds = tfd.Categorical(probs=tf.transpose(weights)).sample() # [num_samples] outs = tf.reshape( tf.gather(tf.transpose(sample), inds[:, tf.newaxis], batch_dims=1), [num_samples]) sample_ta = sample_ta.write(i, outs) return i + 1, sample_ta def while_cond(i, unused_ta): return i < self.data_dim outs = tf.while_loop(while_cond, while_body, (0, sample_ta), back_prop=False) return tf.transpose(outs[1].stack())
def _rollout(model, initial_state, num_steps): """Rolls out a model trajectory.""" node_type = initial_state['node_type'][:, 0] mask = tf.logical_or(tf.equal(node_type, NodeType.NORMAL), tf.equal(node_type, NodeType.OUTFLOW)) def step_fn(step, velocity, trajectory): prediction = model({**initial_state, 'velocity': velocity}) # don't update boundary nodes next_velocity = tf.where(mask, prediction, velocity) trajectory = trajectory.write(step, velocity) return step + 1, next_velocity, trajectory _, _, output = tf.while_loop( cond=lambda step, cur, traj: tf.less(step, num_steps), body=step_fn, loop_vars=(0, initial_state['velocity'], tf.TensorArray(tf.float32, num_steps)), parallel_iterations=1) return output.stack()
def fill_batches(size, state_fn): """Fill a tensor array with batches of data.""" dummy_batch = state_fn() tas = nest.map_structure( lambda b: tf.TensorArray( b.dtype, size=size, clear_after_read=False), dummy_batch) cond = lambda i, ta: tf.less(i, size) def body(i, tas): batch = state_fn() out_tas = [] for ta, b in py_utils.eqzip(nest.flatten(tas), nest.flatten(batch)): out_tas.append(ta.write(i, b)) return (i + 1, nest.pack_sequence_as(dummy_batch, out_tas)) _, batches = tf.while_loop(cond, body, [tf.constant(0, dtype=tf.int32), tas]) return batches
def batching_loop_initial_values(self): """Returns a list of initial values for the batching tf.while_loop. The return value here defines the shepherd-specific loop vars for a tf.while_loop that does dynamic batching. Updated values for these loop vars will get passed in to `batching_loop_body` and `batching_loop_results`. Returns: A possibly nested list of tensors. """ sparse_tensor_vars = (super(SparseTensorWithValuesShepherd, self).batching_loop_initial_values()) values_ta = tf.TensorArray(self.feature_types()[self.values_key()], size=_TENSORARRAY_INITIAL_SIZE, dynamic_size=True, element_shape=None, infer_shape=False) return _SparseTensorWithValuesLoopVars( sparse_tensor_vars.source_offset, sparse_tensor_vars.dest_offset, sparse_tensor_vars.source_indices_ta, sparse_tensor_vars.dest_indices_ta, values_ta)
def ag_dynamic_rnn(rnn_cell, input_data, initial_state, sequence_length=None): """An autograph-able reimplementation of subset of dynamic_rnn.""" # [batch, time, features] -> [time, batch, features] input_data = tf.transpose(input_data, [1, 0, 2]) if sequence_length is None: max_seq_len = input_data.shape[0] else: max_seq_len = tf.reduce_max(sequence_length) outputs = tf.TensorArray(tf.float32, size=max_seq_len) state = initial_state for i in tf.range(max_seq_len): new_output, new_state = rnn_cell(input_data[i], state) output = tf.where(i < sequence_length, new_output, tf.zeros(new_output.shape)) state = tf.where(i < sequence_length, new_state, state) outputs = outputs.write(i, output) return tf.transpose(outputs.stack(), [1, 0, 2]), state
def _rollout(model, initial_state, num_steps): """Rolls out a model trajectory.""" mask = tf.equal(initial_state['node_type'][:, 0], NodeType.NORMAL) def step_fn(step, prev_pos, cur_pos, trajectory): prediction = model({ **initial_state, 'prev|world_pos': prev_pos, 'world_pos': cur_pos }) # don't update kinematic nodes next_pos = tf.where(mask, prediction, cur_pos) trajectory = trajectory.write(step, cur_pos) return step + 1, cur_pos, next_pos, trajectory _, _, _, output = tf.while_loop( cond=lambda step, last, cur, traj: tf.less(step, num_steps), body=step_fn, loop_vars=(0, initial_state['prev|world_pos'], initial_state['world_pos'], tf.TensorArray(tf.float32, num_steps)), parallel_iterations=1) return output.stack()
def dynamic_rnn(input_data, cell, loop_state_fn, initial_loop_state): inputs_shape_g = tf.shape(input_data) input_shape_l = input_data.get_shape().as_list() pad_input = tf.zeros([ inputs_shape_g[0], ] + input_shape_l[2:]) seq_lengths = inputs_shape_g[1] # raw_rnn uses TensorArray for the input and outputs, in which Tensor must be in [time, batch_size, input_depth] shape. inputs_ta = tf.TensorArray(size=inputs_shape_g[1], dtype=tf.float32).unstack( _transpose_batch_time(input_data), 'TBD_Input') initial_state = cell.zero_state(inputs_shape_g[0], None) def loop_fn(time, previous_output, previous_state, previous_loop_state): # this operation produces boolean tensor of [batch_size] defining if corresponding sequence has ended # all False at the initial step (time == 0) finished = time >= seq_lengths if previous_state is None: # time == 0 return (finished, inputs_ta.read(time), initial_state, previous_output, initial_loop_state) else: step_input = tf.cond(tf.reduce_all(finished), lambda: pad_input, lambda: inputs_ta.read(time)) previous_loop_state = loop_state_fn(time, previous_loop_state, previous_state) return (finished, step_input, previous_state, previous_output, previous_loop_state) outputs_ta, final_state, final_loop_state = tf.nn.raw_rnn(cell, loop_fn) output = _transpose_batch_time(outputs_ta.stack()) return output, final_state, final_loop_state
def self_all_attention(facts, ATTENTION_SIZE, mask, stag='null'): if len(facts.get_shape().as_list()) == 2: facts = tf.expand_dims(facts, 1) def cond(batch, output, i): return tf.less(i, tf.shape(batch)[1]) def body(batch, output, i): self_attention_tmp = din_fcn_attention(batch[:, i, :], batch, ATTENTION_SIZE, mask, softmax_stag=1, stag=stag, mode='LIST') self_attention_tmp = tf.reduce_sum(self_attention_tmp, 1) output = output.write(i, self_attention_tmp) return batch, output, i + 1 output_ta = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True, element_shape=(facts[:, 0, :].get_shape())) _, output_op, _ = tf.while_loop(cond, body, [facts, output_ta, 0], maximum_iterations=MAX_ITERS) self_attention = output_op.stack() self_attention = tf.transpose(self_attention, perm = [1, 0, 2]) return self_attention