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
