def _most_recent_obs_value(obs_values, indicator, delta_time,
attribution_max_delta_time):
"""Returns the most recent lab result for each test within a time frame.
The eligible lab values fall into a time window until time of prediction -
attribution_max_delta_time. Among those we select their most recent value
or zero if there are none.
Args:
obs_values: A dense representation of the observation_values at the position
of their obs_code_ids. A padded Tensor of shape [batch_size,
max_sequence_length, vocab_size] of type float32 where obs_values[b, t,
id] = observation_values[b, t, 0] and id = observation_code_ids[b, t, 0]
and obs_values[b, t, x] = 0 for all other x != id. If t is greater than
the sequence_length of batch entry b then the result is 0 as well.
indicator: A one-hot encoding of whether a value in obs_values comes from
observation_values or is just filled in to be 0. A Tensor of shape
[batch_size, max_sequence_length, vocab_size] and type float32.
delta_time: A Tensor of shape [batch_size, max_sequence_length] describing
the time to prediction.
attribution_max_delta_time: Time threshold so that we return the most recent
lab values among those that are at least attribution_max_delta_time
seconds old at time of prediction.
Returns:
A Tensor of shape [batch_size, 1, vocab_size] of the most recent lab results
for all lab tests that are at least attribution_max_delta_time old at time
of prediction.
"""
batch_size = tf.shape(indicator)[0]
seq_len = tf.shape(indicator)[1]
num_obs = indicator.shape[2]
# Prepend a dummy so that for lab tests for which we have no eligible lab
# values we will select 0.
obs_values = tf.concat(
[tf.zeros([batch_size, 1, num_obs]), obs_values], axis=1)
indicator = tf.concat([tf.ones([batch_size, 1, num_obs]), indicator], axis=1)
delta_time = tf.to_int32(delta_time)
delta_time = tf.concat(
[
tf.zeros([batch_size, 1, 1], dtype=tf.int32) +
attribution_max_delta_time, delta_time
],
axis=1)
# First we figure out what the eligible lab values are that are at least
# attribution_max_delta_time old.
indicator = tf.to_int32(indicator)
indicator *= tf.to_int32(delta_time >= attribution_max_delta_time)
range_val = tf.expand_dims(tf.range(seq_len + 1), axis=0)
range_val = tf.tile(range_val, multiples=[tf.shape(indicator)[0], 1])
# [[[0], [1], ..., [max_sequence_length]],
# [[0], [1], ..., [max_sequence_length]],
# ...]
range_val = tf.expand_dims(range_val, axis=2)
# [batch_size, max_sequence_length, vocab_size] with 1 non-zero number per
# time-step equal to that time-step.
seq_indicator = indicator * range_val
# [batch_size, vocab_size] with the time-step of the last lab value.
last_val_indicator = tf.reduce_max(seq_indicator, axis=1, keepdims=True)
last_val_indicator = tf.tile(
last_val_indicator, multiples=[1, tf.shape(indicator)[1], 1])
# eq indicates which lab values are the most recent ones.
eq = tf.logical_and(
tf.equal(last_val_indicator, seq_indicator), indicator > 0)
most_recent_obs_value_indicator = tf.where(eq)
# Collect the lab values associated with those indices.
res = tf.gather_nd(obs_values, most_recent_obs_value_indicator)
# Reorder the values by batch and then by lab test.
res_sorted = tf.sparse_reorder(
tf.sparse_transpose(
tf.SparseTensor(
indices=most_recent_obs_value_indicator,
values=res,
dense_shape=tf.to_int64(
tf.stack([batch_size, seq_len + 1, num_obs]))),
perm=[0, 2, 1])).values
return tf.reshape(res_sorted, [batch_size, 1, num_obs])
def cudnn_lstm_layer(inputs,
batch_size,
num_units,
lengths=None,
stack_size=1,
rnn_dropout_drop_amt=0,
is_training=True,
bidirectional=True):
"""Create a LSTM layer that uses cudnn."""
inputs_t = tf.transpose(inputs, [1, 0, 2])
if lengths is not None:
all_outputs = [inputs_t]
for i in range(stack_size):
with tf.variable_scope('stack_' + str(i)):
with tf.variable_scope('forward'):
lstm_fw = contrib_cudnn_rnn.CudnnLSTM(
num_layers=1,
num_units=num_units,
direction='unidirectional',
dropout=rnn_dropout_drop_amt,
kernel_initializer=contrib_layers.
variance_scaling_initializer(),
bias_initializer=tf.zeros_initializer(),
)
c_fw = tf.zeros([1, batch_size, num_units], tf.float32)
h_fw = tf.zeros([1, batch_size, num_units], tf.float32)
outputs_fw, _ = lstm_fw(all_outputs[-1], (h_fw, c_fw),
training=is_training)
combined_outputs = outputs_fw
if bidirectional:
with tf.variable_scope('backward'):
lstm_bw = contrib_cudnn_rnn.CudnnLSTM(
num_layers=1,
num_units=num_units,
direction='unidirectional',
dropout=rnn_dropout_drop_amt,
kernel_initializer=contrib_layers.
variance_scaling_initializer(),
bias_initializer=tf.zeros_initializer(),
)
c_bw = tf.zeros([1, batch_size, num_units], tf.float32)
h_bw = tf.zeros([1, batch_size, num_units], tf.float32)
inputs_reversed = tf.reverse_sequence(all_outputs[-1],
lengths,
seq_axis=0,
batch_axis=1)
outputs_bw, _ = lstm_bw(inputs_reversed, (h_bw, c_bw),
training=is_training)
outputs_bw = tf.reverse_sequence(outputs_bw,
lengths,
seq_axis=0,
batch_axis=1)
combined_outputs = tf.concat([outputs_fw, outputs_bw],
axis=2)
all_outputs.append(combined_outputs)
# for consistency with cudnn, here we just return the top of the stack,
# although this can easily be altered to do other things, including be
# more resnet like
return tf.transpose(all_outputs[-1], [1, 0, 2])
else:
lstm = contrib_cudnn_rnn.CudnnLSTM(
num_layers=stack_size,
num_units=num_units,
direction='bidirectional' if bidirectional else 'unidirectional',
dropout=rnn_dropout_drop_amt,
kernel_initializer=contrib_layers.variance_scaling_initializer(),
bias_initializer=tf.zeros_initializer(),
)
stack_multiplier = 2 if bidirectional else 1
c = tf.zeros([stack_multiplier * stack_size, batch_size, num_units],
tf.float32)
h = tf.zeros([stack_multiplier * stack_size, batch_size, num_units],
tf.float32)
outputs, _ = lstm(inputs_t, (h, c), training=is_training)
outputs = tf.transpose(outputs, [1, 0, 2])
return outputs
def unsupervised_wiener(image,
psf,
reg=None,
user_params=None,
is_real=True,
clip=True):
"""Unsupervised Wiener-Hunt deconvolution.
Return the deconvolution with a Wiener-Hunt approach, where the
hyperparameters are automatically estimated. The algorithm is a
stochastic iterative process (Gibbs sampler) described in the
reference below. See also ``wiener`` function.
Parameters
----------
image : (M, N) ndarray
The input degraded image.
psf : ndarray
The impulse response (input image's space) or the transfer
function (Fourier space). Both are accepted. The transfer
function is automatically recognized as being complex
(``np.iscomplexobj(psf)``).
reg : ndarray, optional
The regularisation operator. The Laplacian by default. It can
be an impulse response or a transfer function, as for the psf.
user_params : dict, optional
Dictionary of parameters for the Gibbs sampler.
clip : boolean, optional
True by default. If true, pixel values of the result above 1 or
under -1 are thresholded for skimage pipeline compatibility.
Returns
-------
x_postmean : (M, N) ndarray
The deconvolved image (the posterior mean).
chains : dict
The keys ``noise`` and ``prior`` contain the chain list of
noise and prior precision respectively.
Other parameters
----------------
The keys of ``user_params`` are:
threshold : float
The stopping criterion: the norm of the difference between to
successive approximated solution (empirical mean of object
samples, see Notes section). 1e-4 by default.
burnin : int
The number of sample to ignore to start computation of the
mean. 15 by default.
min_iter : int
The minimum number of iterations. 30 by default.
max_iter : int
The maximum number of iterations if ``threshold`` is not
satisfied. 200 by default.
callback : callable (None by default)
A user provided callable to which is passed, if the function
exists, the current image sample for whatever purpose. The user
can store the sample, or compute other moments than the
mean. It has no influence on the algorithm execution and is
only for inspection.
References
----------
.. [1] François Orieux, Jean-François Giovannelli, and Thomas
Rodet, "Bayesian estimation of regularization and point
spread function parameters for Wiener-Hunt deconvolution",
J. Opt. Soc. Am. A 27, 1593-1607 (2010)
https://www.osapublishing.org/josaa/abstract.cfm?URI=josaa-27-7-1593
http://research.orieux.fr/files/papers/OGR-JOSA10.pdf
"""
sess = tf.InteractiveSession()
params = {
'threshold': 1e-4,
'max_iter': 200,
'min_iter': 30,
'burnin': 15,
'callback': None
}
params.update(user_params or {})
if reg is None:
reg, _ = laplacian(image.ndim, image.shape, sess, is_real=is_real)
if reg.dtype != tf.complex64:
reg = ir2tf(reg, image.shape, sess, is_real=is_real)
if psf.shape != reg.shape:
trans_fct = ir2tf(psf, image.shape, sess, is_real=is_real)
else:
trans_fct = psf
# The mean of the object
x_postmean = tf.Variable(tf.cast(tf.zeros(trans_fct.shape), tf.complex64))
# The previous computed mean in the iterative loop
prev_x_postmean = tf.Variable(
tf.cast(tf.zeros(trans_fct.shape), tf.complex64))
# Difference between two successive mean
delta = tf.Variable(1e-8)
# Initial state of the chain
gn_chain, gx_chain = [tf.constant(1.0)], [tf.constant(1.0)]
# The correlation of the object in Fourier space (if size is big,
# this can reduce computation time in the loop)
areg2 = tf.abs(reg)**2
atf2 = tf.abs(trans_fct)**2
# The Fourier transform may change the image.size attribute, so we
# store it.
if is_real:
data_spectrum = tf.signal.rfft2d(tf.cast(image, tf.float32))
else:
data_spectrum = tf.fft2d(tf.cast(image, tf.float32))
# Gibbs sampling
update_prev_x_postmean_op = tf.assign(prev_x_postmean, x_postmean)
bool_op = tf.cond(delta < params['threshold'], lambda: True, lambda: False)
# weighting (correlation in direct space)
precision = tf.Variable(gn_chain[-1] * atf2 +
gx_chain[-1] * areg2) # Eq. 29
excursion = tf.Variable(
tf.cast(tf.sqrt(0.5) / tf.sqrt(precision), tf.complex64) *
(tf.cast(tf.random.normal(data_spectrum.shape), tf.complex64) +
1j * tf.cast(tf.random.normal(data_spectrum.shape), tf.complex64)))
# mean Eq. 30 (RLS for fixed gn, gamma0 and gamma1 ...)
wiener_filter = tf.Variable(
tf.cast(gn_chain[-1], tf.complex64) * tf.math.conj(trans_fct) /
tf.cast(precision, tf.complex64))
# sample of X in Fourier space
x_sample = tf.Variable(wiener_filter * data_spectrum + excursion)
# Define oprator for updating variables
update_precision_op = tf.assign(precision,
gn_chain[-1] * atf2 + gx_chain[-1] * areg2)
update_excursion_op = tf.assign(excursion, tf.cast(tf.sqrt(0.5) / tf.sqrt(precision), tf.complex64) * (tf.cast(
tf.random.normal(data_spectrum.shape), tf.complex64) + \
1j * tf.cast(tf.random.normal(data_spectrum.shape), tf.complex64)))
update_wiener_filter_op = tf.assign(
wiener_filter,
tf.cast(gn_chain[-1], tf.complex64) * tf.math.conj(trans_fct) /
tf.cast(precision, tf.complex64))
update_x_sample_op = tf.assign(x_sample,
wiener_filter * data_spectrum + excursion)
update_x_postmean_op = tf.assign(x_postmean, prev_x_postmean + x_sample)
update_group = tf.group([
update_precision_op, update_excursion_op, update_wiener_filter_op,
update_x_sample_op
])
# Initialize variables
variable_set_1 = [x_postmean, prev_x_postmean, delta, precision]
sess.run(tf.variables_initializer(variable_set_1))
variable_set_2 = [excursion, wiener_filter, x_sample]
sess.run(tf.variables_initializer(variable_set_2))
for iteration in range(params['max_iter']):
# Sample of Eq. 27 p(circX^k | gn^k-1, gx^k-1, y).
sess.run(update_group)
# sample of X in Fourier space
if params['callback']:
params['callback'](x_sample)
# sample of Eq. 31 p(gn | x^k, gx^k, y)
gn_chain.append(
tf.random.gamma(shape=[1],
alpha=[int(image.size / 2)],
beta=image_quad_norm(data_spectrum -
x_sample * trans_fct)))
# sample of Eq. 31 p(gx | x^k, gn^k-1, y)
gx_chain.append(
tf.random.gamma(shape=[1],
alpha=[int(image.size / 2)],
beta=image_quad_norm(x_sample * reg)))
# current empirical average
if iteration > params['burnin']:
sess.run(update_x_postmean_op)
if iteration > (params['burnin'] + 1):
current = x_postmean / (iteration - params['burnin'])
previous = prev_x_postmean / (iteration - params['burnin'] - 1)
update_delta_op = tf.assign(
delta, (tf.reduce_sum(tf.abs(current - previous)) /
tf.reduce_sum(tf.abs(x_postmean)) /
(iteration - params['burnin'])))
sess.run(update_delta_op)
sess.run(update_prev_x_postmean_op)
result_is_found = sess.run(bool_op)
# stop of the algorithm
if (iteration > params['min_iter']) and result_is_found:
break
# Empirical average \approx POSTMEAN Eq. 44
x_postmean = x_postmean / (iteration - params['burnin'])
if is_real:
x_postmean = tf.signal.irfft2d(x_postmean)
else:
x_postmean = tf.signal.ifft2d(x_postmean)
x_postmean = x_postmean.eval()
sess.close()
return (x_postmean, {'noise': gn_chain, 'prior': gx_chain})
def __init__(self,
game,
num_hidden_units,
num_hidden_layers=1,
num_hidden_factors=0,
hidden_activation=tf.nn.relu,
use_skip_connections=False,
regularizer=None,
autoencode=False):
"""Creates a new `DeepNeurdModel.
Args:
game: The OpenSpiel game being solved.
num_hidden_units: The number of units in each hidden layer.
num_hidden_layers: The number of hidden layers. Defaults to 1.
num_hidden_factors: The number of hidden factors or the matrix rank of the
layer. If greater than zero, hidden layers will be split into two
separate linear transformations, the first with
`num_hidden_factors`-columns and the second with
`num_hidden_units`-columns. The result is that the logical hidden layer
is a rank-`num_hidden_units` matrix instead of a rank-`num_hidden_units`
matrix. When `num_hidden_units < num_hidden_units`, this is effectively
implements weight sharing. Defaults to 0.
hidden_activation: The activation function to apply over hidden layers.
Defaults to `tf.nn.relu`.
use_skip_connections: Whether or not to apply skip connections (layer
output = layer(x) + x) on hidden layers. Zero padding or truncation is
used to match the number of columns on layer inputs and outputs.
regularizer: A regularizer to apply to each layer. Defaults to `None`.
autoencode: Whether or not to output a reconstruction of the inputs upon
being called. Defaults to `False`.
"""
self._autoencode = autoencode
self._use_skip_connections = use_skip_connections
self._hidden_are_factored = num_hidden_factors > 0
self.layers = []
for _ in range(num_hidden_layers):
if self._hidden_are_factored:
self.layers.append(
tf.keras.layers.Dense(num_hidden_factors,
use_bias=True,
kernel_regularizer=regularizer))
self.layers.append(
tf.keras.layers.Dense(num_hidden_units,
use_bias=True,
activation=hidden_activation,
kernel_regularizer=regularizer))
self.layers.append(
tf.keras.layers.Dense(1 +
self._autoencode * rcfr.num_features(game),
use_bias=True,
kernel_regularizer=regularizer))
# Construct variables for all layers by exercising the network.
x = tf.zeros([1, rcfr.num_features(game)])
for layer in self.layers:
x = layer(x)
self.trainable_variables = sum(
[layer.trainable_variables for layer in self.layers], [])
self.losses = sum([layer.losses for layer in self.layers], [])
def get_policy(observations, hparams, action_space,
distributional_size=1, epoch=-1):
"""Get a policy network.
Args:
observations: observations
hparams: parameters
action_space: action space
distributional_size: optional number of buckets for distributional RL
epoch: optional epoch number
Returns:
Tuple (action logits, value).
"""
if not isinstance(action_space, gym.spaces.Discrete):
raise ValueError("Expecting discrete action space.")
obs_shape = common_layers.shape_list(observations)
(frame_height, frame_width) = obs_shape[2:4]
# TODO(afrozm): We have these dummy problems mainly for hparams, so cleanup
# when possible and do this properly.
if hparams.policy_problem_name == "dummy_policy_problem_ttt":
tf.logging.info("Using DummyPolicyProblemTTT for the policy.")
policy_problem = tic_tac_toe_env.DummyPolicyProblemTTT()
else:
tf.logging.info("Using DummyPolicyProblem for the policy.")
policy_problem = DummyPolicyProblem(action_space, frame_height, frame_width)
trainer_lib.add_problem_hparams(hparams, policy_problem)
hparams.force_full_predict = True
model = registry.model(hparams.policy_network)(
hparams, tf.estimator.ModeKeys.TRAIN
)
try:
num_target_frames = hparams.video_num_target_frames
except AttributeError:
num_target_frames = 1
target_value_shape_suffix = [num_target_frames]
if distributional_size > 1:
target_value_shape_suffix = [num_target_frames, distributional_size]
features = {
"inputs": observations,
"epoch": tf.constant(epoch + 1),
"input_action": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
"input_reward": tf.zeros(obs_shape[:2] + [1], dtype=tf.int32),
"targets": tf.zeros(obs_shape[:1] + [num_target_frames] + obs_shape[2:]),
"target_action": tf.zeros(
obs_shape[:1] + [num_target_frames, 1], dtype=tf.int32),
"target_reward": tf.zeros(
obs_shape[:1] + [num_target_frames, 1], dtype=tf.int32),
"target_policy": tf.zeros(
obs_shape[:1] + [num_target_frames] + [action_space.n]),
"target_value": tf.zeros(
obs_shape[:1] + target_value_shape_suffix)
}
model.distributional_value_size = max(distributional_size, 1)
model.use_epochs = hparams.use_epochs
with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
t2t_model.create_dummy_vars()
(targets, _) = model(features)
target_values = targets["target_value"][:, 0]
if distributional_size > 1:
target_values = targets["target_value"][:, :]
return (targets["target_policy"][:, 0, :], target_values)
def __init__(self,
config,
is_training,
input_ids,
input_mask=None,
token_type_ids=None,
use_one_hot_embeddings=False,
use_einsum=True,
scope=None):
"""Constructor for AlbertModel.
Args:
config: `AlbertConfig` instance.
is_training: bool. true for training model, false for eval model. Controls
whether dropout will be applied.
input_ids: int32 Tensor of shape [batch_size, seq_length].
input_mask: (optional) int32 Tensor of shape [batch_size, seq_length].
token_type_ids: (optional) int32 Tensor of shape [batch_size, seq_length].
use_one_hot_embeddings: (optional) bool. Whether to use one-hot word
embeddings or tf.embedding_lookup() for the word embeddings.
use_einsum: (optional) bool. Whether to use einsum or reshape+matmul for
dense layers
scope: (optional) variable scope. Defaults to "bert".
Raises:
ValueError: The config is invalid or one of the input tensor shapes
is invalid.
"""
config = copy.deepcopy(config)
if not is_training:
config.hidden_dropout_prob = 0.0
config.attention_probs_dropout_prob = 0.0
input_shape = get_shape_list(input_ids, expected_rank=2)
batch_size = input_shape[0]
seq_length = input_shape[1]
if input_mask is None:
input_mask = tf.ones(shape=[batch_size, seq_length],
dtype=tf.int32)
if token_type_ids is None:
token_type_ids = tf.zeros(shape=[batch_size, seq_length],
dtype=tf.int32)
with tf.variable_scope(scope, default_name="bert"):
with tf.variable_scope("embeddings"):
# Perform embedding lookup on the word ids.
(self.word_embedding_output,
self.output_embedding_table) = embedding_lookup(
input_ids=input_ids,
vocab_size=config.vocab_size,
embedding_size=config.embedding_size,
initializer_range=config.initializer_range,
word_embedding_name="word_embeddings",
use_one_hot_embeddings=use_one_hot_embeddings)
# Add positional embeddings and token type embeddings, then layer
# normalize and perform dropout.
self.embedding_output = embedding_postprocessor(
input_tensor=self.word_embedding_output,
use_token_type=True,
token_type_ids=token_type_ids,
token_type_vocab_size=config.type_vocab_size,
token_type_embedding_name="token_type_embeddings",
use_position_embeddings=True,
position_embedding_name="position_embeddings",
initializer_range=config.initializer_range,
max_position_embeddings=config.max_position_embeddings,
dropout_prob=config.hidden_dropout_prob,
use_one_hot_embeddings=use_one_hot_embeddings)
with tf.variable_scope("encoder"):
# Run the stacked transformer.
# `sequence_output` shape = [batch_size, seq_length, hidden_size].
self.all_encoder_layers = transformer_model(
input_tensor=self.embedding_output,
attention_mask=input_mask,
hidden_size=config.hidden_size,
num_hidden_layers=config.num_hidden_layers,
num_hidden_groups=config.num_hidden_groups,
num_attention_heads=config.num_attention_heads,
intermediate_size=config.intermediate_size,
inner_group_num=config.inner_group_num,
intermediate_act_fn=get_activation(config.hidden_act),
hidden_dropout_prob=config.hidden_dropout_prob,
attention_probs_dropout_prob=config.
attention_probs_dropout_prob,
initializer_range=config.initializer_range,
do_return_all_layers=True,
use_einsum=use_einsum)
self.sequence_output = self.all_encoder_layers[-1]
# The "pooler" converts the encoded sequence tensor of shape
# [batch_size, seq_length, hidden_size] to a tensor of shape
# [batch_size, hidden_size]. This is necessary for segment-level
# (or segment-pair-level) classification tasks where we need a fixed
# dimensional representation of the segment.
with tf.variable_scope("pooler"):
# We "pool" the model by simply taking the hidden state corresponding
# to the first token. We assume that this has been pre-trained
first_token_tensor = tf.squeeze(self.sequence_output[:,
0:1, :],
axis=1)
self.pooled_output = tf.layers.dense(
first_token_tensor,
config.hidden_size,
activation=tf.tanh,
kernel_initializer=create_initializer(
config.initializer_range))
def predict(self, features):
batch_size = tf.shape(features)[0]
return tf.zeros((batch_size, 1, self._num_classes, DEFAULT_MASK_SIZE,
DEFAULT_MASK_SIZE),
dtype=tf.float32)
def infer(self, features, **kwargs):
with tf.variable_scope("sparse_transformer", reuse=tf.AUTO_REUSE):
features = self.bottom(features)
decode_length = self.hparams.max_target_length
cache = {}
decoding_stats = {}
targets_old = features.get("targets")
start_step = 0
initial_output = tf.zeros((self.batch_size, decode_length, 1, 1),
dtype=tf.int32)
initial_logits = tf.zeros(
(self.batch_size, decode_length, self.vocab_size))
# call body once to initialize cache with representations of input frames.
features["targets"] = initial_output
# Set shape of inputs
if "inputs" in features:
features["inputs"].set_shape([
self.batch_size, self.hparams.max_length, 1,
self.hparams.hidden_size
])
with tf.variable_scope("sparse_transformer/body", reuse=tf.AUTO_REUSE):
self.body(features,
decode_step=None,
cache=cache,
decoding_stats=decoding_stats)
def infer_step(i, recent_output, recent_logits, cache, decoding_stats):
"""Inference step."""
features_copy = features.copy()
features_copy["targets"] = recent_output
cur_sample, cur_logit = self.sample(features_copy,
decode_step=i,
cache=cache,
decoding_stats=decoding_stats)
pos = i
samples = recent_output + tf.scatter_nd(
indices=[[b, pos, 0, 0] for b in range(self.batch_size)],
updates=cur_sample,
shape=utils.shape_list(recent_output))
logits = recent_logits + tf.scatter_nd(
indices=[[b, pos] for b in range(self.batch_size)],
updates=cur_logit,
shape=utils.shape_list(recent_logits))
return i + 1, samples, logits, cache, decoding_stats
def while_exit_cond(i, result, logits, cache, decoding_stats): # pylint: disable=unused-argument
"""Exit the loop if it reaches decode_length."""
not_overflow = i < decode_length
return not_overflow
_, final_result, final_logits, _, decoding_stats = tf.while_loop(
while_exit_cond,
infer_step, [
start_step, initial_output, initial_logits, cache,
decoding_stats
],
back_prop=False,
parallel_iterations=1)
original_shape = [decode_length]
blocks_per_dim = [
s // q for s, q in zip(original_shape, self.hparams.query_shape)
]
final_result_shape = utils.shape_list(final_result)
final_result = tf.reshape(
final_result,
[final_result_shape[0], -1,
np.prod(self.hparams.query_shape), 1])
final_logits_shape = utils.shape_list(final_logits)
final_logits = tf.reshape(final_logits, [
final_logits_shape[0], -1,
np.prod(self.hparams.query_shape), final_logits_shape[-1]
])
final_result = utils.unflatten_blocks_nd(final_result, blocks_per_dim)
final_result = utils.put_back_blocks_nd(final_result,
self.hparams.query_shape)
final_logits = utils.unflatten_blocks_nd(final_logits, blocks_per_dim)
final_logits = utils.put_back_blocks_nd(final_logits,
self.hparams.query_shape)
for name, value in decoding_stats.items():
tf.summary.scalar("decodes/%s" % name, value / decode_length)
# Reassign targets back to the previous value.
if targets_old is not None:
features["targets"] = targets_old
return {
"outputs": final_result,
"scores": None,
"logits": final_logits,
"losses": None,
}
def construct_model(input_tensors, encoder_w0, prefix=None):
"""Construct model."""
facto = tf.placeholder_with_default(1.0, ())
context_xs = input_tensors['inputa']
context_ys = input_tensors['labela']
target_xs = input_tensors['inputb']
target_ys = input_tensors['labelb']
# sample ws ~ w|(x_all,a), rs = T(ws, ys), r = mean(rs), z = T(r)
# x_all = tf.concat([context_xs, target_xs], axis=1) #n_task * 20 * (128*128)
# y_all = tf.concat([context_ys, target_ys], axis=1)
x_all = context_xs
y_all = context_ys
# n_task * [n_im] * d_z
if 'train' in prefix:
z_samples, mu_w_all, sigma_w_all = xy_to_z(x_all, y_all, encoder_w0)
z_samples = z_samples * facto
else:
z_samples, _, _ = xy_to_z(context_xs, context_ys, encoder_w0)
z_samples = z_samples * facto
target_ws, _, _ = encoder_w(target_xs, encoder_w0)
input_zxs = tf.concat([z_samples, target_ws], axis=-1)
# sample y_hat ~ y|(w,z)
target_yhat_mu = decoder_g(input_zxs) # n_task * n_im * dim_y
# when var of p(y | x,z) is fixed, neg-loglik <=> MSE
mse_loss = mse(target_yhat_mu, target_ys)
tf.summary.scalar(prefix + 'mse', mse_loss)
optimizer1 = tf.train.AdamOptimizer(FLAGS.update_lr)
optimizer2 = tf.train.AdamOptimizer(0.001)
if 'train' in prefix:
# mu_w_all is n_task * n_im * dim_w
# target_yhat_mu is n_task * n_im * dim_w
kl_ib = kl_qp_gaussian(mu_w_all, sigma_w_all,
tf.zeros(tf.shape(mu_w_all)),
tf.ones(tf.shape(mu_w_all)))
THETA = ( # pylint: disable=invalid-name
tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='decoder') +
tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope='encoder_w'))
all_var = tf.trainable_variables()
PHI = [v for v in all_var if v not in THETA] # pylint: disable=invalid-name
loss = mse_loss + FLAGS.beta * kl_ib
gvs_theta = optimizer1.compute_gradients(loss, THETA)
train_theta_op = optimizer1.apply_gradients(gvs_theta)
gvs_phi = optimizer2.compute_gradients(loss, PHI)
train_phi_op = optimizer2.apply_gradients(gvs_phi)
with tf.control_dependencies([train_theta_op, train_phi_op]):
train_op = tf.no_op()
tf.summary.scalar(prefix + 'kl', kl_ib)
tf.summary.scalar(prefix + 'full_loss', loss)
return mse_loss, kl_ib, train_op, facto
else:
return mse_loss
def train(train_dir,
config,
dataset_fn,
checkpoints_to_keep=5,
keep_checkpoint_every_n_hours=1,
num_steps=None,
master='',
num_sync_workers=0,
num_ps_tasks=0,
task=0):
"""Train loop."""
tf.gfile.MakeDirs(train_dir)
is_chief = (task == 0)
if is_chief:
_trial_summary(config.hparams, config.train_examples_path
or config.tfds_name, train_dir)
with tf.Graph().as_default():
with tf.device(
tf.train.replica_device_setter(num_ps_tasks,
merge_devices=True)):
model = config.model
model.build(config.hparams,
config.data_converter.output_depth,
is_training=True)
optimizer = model.train(**_get_input_tensors(dataset_fn(), config))
hooks = []
if num_sync_workers:
optimizer = tf.train.SyncReplicasOptimizer(
optimizer, num_sync_workers)
hooks.append(optimizer.make_session_run_hook(is_chief))
grads, var_list = list(
zip(*optimizer.compute_gradients(model.loss)))
global_norm = tf.global_norm(grads)
tf.summary.scalar('global_norm', global_norm)
if config.hparams.clip_mode == 'value':
g = config.hparams.grad_clip
clipped_grads = [
tf.clip_by_value(grad, -g, g) for grad in grads
]
elif config.hparams.clip_mode == 'global_norm':
clipped_grads = tf.cond(
global_norm < config.hparams.grad_norm_clip_to_zero,
lambda: tf.clip_by_global_norm( # pylint:disable=g-long-lambda
grads,
config.hparams.grad_clip,
use_norm=global_norm)[0],
lambda: [tf.zeros(tf.shape(g)) for g in grads])
else:
raise ValueError('Unknown clip_mode: {}'.format(
config.hparams.clip_mode))
train_op = optimizer.apply_gradients(list(
zip(clipped_grads, var_list)),
global_step=model.global_step,
name='train_step')
logging_dict = {
'global_step': model.global_step,
'loss': model.loss
}
hooks.append(
tf.train.LoggingTensorHook(logging_dict, every_n_iter=100))
if num_steps:
hooks.append(tf.train.StopAtStepHook(last_step=num_steps))
scaffold = tf.train.Scaffold(saver=tf.train.Saver(
max_to_keep=checkpoints_to_keep,
keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours))
tf_slim.training.train(train_op=train_op,
logdir=train_dir,
scaffold=scaffold,
hooks=hooks,
save_checkpoint_secs=60,
master=master,
is_chief=is_chief)
def network(input):
#Xavier initializer
initializer = tf.initializers.glorot_uniform()
filter0 = tf.Variable(initializer([3,3,4,32]))
b0 = tf.Variable(tf.zeros([32]))
conv1 = tf.nn.conv2d(input, filter0,[1,1,1,1],padding='SAME')
conv1 = tf.nn.bias_add(conv1,b0)
conv1 = lrelu(conv1)
filter1 = tf.Variable(initializer([3,3,32,32]))
b1 = tf.Variable(tf.zeros([32]))
conv1 = tf.nn.conv2d(conv1, filter1,[1,1,1,1],padding='SAME')
conv1 = tf.nn.bias_add(conv1,b1)
conv1 = lrelu(conv1)
pool1=tf.nn.max_pool(conv1,[1,2,2,1],[1,2,2,1],padding='SAME')
filter2 = tf.Variable(initializer([3,3,32,64]))
b2 = tf.Variable(tf.zeros([64]))
conv2 = tf.nn.conv2d(pool1, filter2,[1,1,1,1],padding='SAME')
conv2 = tf.nn.bias_add(conv2,b2)
conv2 = lrelu(conv2)
filter3 = tf.Variable(initializer([3,3,64,64]))
b3 = tf.Variable(tf.zeros([64]))
conv2 = tf.nn.conv2d(conv2, filter3,[1,1,1,1],padding='SAME')
conv2 = tf.nn.bias_add(conv2,b3)
conv2 = lrelu(conv2)
pool2=tf.nn.max_pool(conv2,[1,2,2,1],[1,2,2,1],padding='SAME')
filter4 = tf.Variable(initializer([3,3,64,128]))
b4 = tf.Variable(tf.zeros([128]))
conv3 = tf.nn.conv2d(pool2, filter4,[1,1,1,1],padding='SAME')
conv3 = tf.nn.bias_add(conv3,b4)
conv3 = lrelu(conv3)
filter5 = tf.Variable(initializer([3,3,128,128]))
b5 = tf.Variable(tf.zeros([128]))
conv3 = tf.nn.conv2d(conv3, filter5,[1,1,1,1],padding='SAME')
conv3 = tf.nn.bias_add(conv3,b5)
conv3 = lrelu(conv3)
pool3=tf.nn.max_pool(conv3,[1,2,2,1],[1,2,2,1],padding='SAME')
filter6 = tf.Variable(initializer([3,3,128,256]))
b6 = tf.Variable(tf.zeros([256]))
conv4 = tf.nn.conv2d(pool3, filter6,[1,1,1,1],padding='SAME')
conv4 = tf.nn.bias_add(conv4,b6)
conv4 = lrelu(conv4)
filter7 = tf.Variable(initializer([3,3,256,256]))
b7 = tf.Variable(tf.zeros([256]))
conv4 = tf.nn.conv2d(conv4, filter7,[1,1,1,1],padding='SAME')
conv4 = tf.nn.bias_add(conv4,b7)
conv4 = lrelu(conv4)
pool4=tf.nn.max_pool(conv4,[1,2,2,1],[1,2,2,1],padding='SAME')
filter8 = tf.Variable(initializer([3,3,256,512]))
b8 = tf.Variable(tf.zeros([512]))
conv5 = tf.nn.conv2d(pool4, filter8,[1,1,1,1],padding='SAME')
conv5 = tf.nn.bias_add(conv5,b8)
conv5 = lrelu(conv5)
filter9 = tf.Variable(initializer([3,3,512,512]))
b9 = tf.Variable(tf.zeros([512]))
conv5 = tf.nn.conv2d(conv5, filter9,[1,1,1,1],padding='SAME')
conv5 = tf.nn.bias_add(conv5,b9)
conv5 = lrelu(conv5)
up6 = upsample_and_concat( conv5, conv4, 256, 512 ) #todo:bug
filter10 = tf.Variable(initializer([3,3,512,256]))
b10 = tf.Variable(tf.zeros([256]))
conv6 = tf.nn.conv2d(up6, filter10,[1,1,1,1],padding='SAME')
conv6 = tf.nn.bias_add(conv6,b10)
conv6 = lrelu(conv6)
filter11 = tf.Variable(initializer([3,3,256,256]))
b11 = tf.Variable(tf.zeros([256]))
conv6 = tf.nn.conv2d(conv6, filter11,[1,1,1,1],padding='SAME')
conv6 = tf.nn.bias_add(conv6,b11)
conv6 = lrelu(conv6)
up7 = upsample_and_concat( conv6, conv3, 128, 256 )
filter12 = tf.Variable(initializer([3,3,256,128]))
b12 = tf.Variable(tf.zeros([128]))
conv7 = tf.nn.conv2d(up7, filter12,[1,1,1,1],padding='SAME')
conv7 = tf.nn.bias_add(conv7,b12)
conv7 = lrelu(conv7)
filter13 = tf.Variable(initializer([3,3,128,128]))
b13 = tf.Variable(tf.zeros([128]))
conv7 = tf.nn.conv2d(conv7, filter13,[1,1,1,1],padding='SAME')
conv7 = tf.nn.bias_add(conv7,b13)
conv7 = lrelu(conv7)
up8 = upsample_and_concat( conv7, conv2, 64, 128 )
filter14 = tf.Variable(initializer([3,3,128,64]))
b14 = tf.Variable(tf.zeros([64]))
conv8 = tf.nn.conv2d(up8, filter14,[1,1,1,1],padding='SAME')
conv8 = tf.nn.bias_add(conv8,b14)
conv8 = lrelu(conv8)
filter15 = tf.Variable(initializer([3,3,64,64]))
b15 = tf.Variable(tf.zeros([64]))
conv8 = tf.nn.conv2d(conv8, filter15,[1,1,1,1],padding='SAME')
conv8 = tf.nn.bias_add(conv8,b15)
conv8 = lrelu(conv8)
up9 = upsample_and_concat( conv8, conv1, 32, 64 )
filter16 = tf.Variable(initializer([3,3,64,32]))
b16 = tf.Variable(tf.zeros([32]))
conv9 = tf.nn.conv2d(up9, filter16,[1,1,1,1],padding='SAME')
conv9 = tf.nn.bias_add(conv9,b16)
conv9 = lrelu(conv9)
filter17 = tf.Variable(initializer([3,3,32,32]))
b17 = tf.Variable(tf.zeros([32]))
conv9 = tf.nn.conv2d(conv9, filter17,[1,1,1,1],padding='SAME')
conv9 = tf.nn.bias_add(conv9,b17)
conv9 = lrelu(conv9)
filter18 = tf.Variable(initializer([1,1,32,12]))
b18 = tf.Variable(tf.zeros([12]))
conv10 = tf.nn.conv2d(conv9, filter18,[1,1,1,1],padding='SAME')
conv10 = tf.nn.bias_add(conv10,b18)
out = tf.depth_to_space(conv10,2)
return out
def _GetWavData(self): sample_data = tf.zeros([32000, 2]) wav_encoder = tf.audio.encode_wav(sample_data, 16000) wav_data = self.evaluate(wav_encoder) return wav_data
def model(x):
with tf.variable_scope("matmul"):
W = tf.get_variable("W", initializer=tf.ones(shape=(2,2)))
b = tf.get_variable("b", initializer=tf.zeros(shape=(2)))
return x * W + b
data_index += 1
data_index = (data_index + len(data) - span) % len(
data
) # Backtrack a little to avoid skipping words in the end of a batch
return batch, labels
with tf.device(device):
embedding = tf.Variable(
tf.random.normal([vocabulary_size, embedding_size])
) # Create the embedding variable (each row represent a word embedding vector)
nce_weights = tf.Variable(
tf.random.normal([vocabulary_size, embedding_size
])) # Construct the weight variable for NCE loss
nce_biases = tf.Variable(tf.zeros([vocabulary_size]))
def get_embedding(input_batch):
with tf.device(device):
input_embedded = tf.nn.embedding_lookup(
embedding, input_batch
) # Lookup the corresponding embedding vectors for each sample in the input
return input_embedded
# Compute the average NCE loss for the batch
def nce_loss(input_embedded, target):
with tf.device(device):
target = tf.cast(target, tf.int64)
loss = tf.reduce_mean(
def preprocess(self, _):
return tf.zeros((1, 128, 128, 3)), tf.constant([[128, 128, 3]])
def ssd_crop(image, boxes, classes):
"""IoU biassed random crop.
Reference: https://github.com/chauhan-utk/ssd.DomainAdaptation
"""
num_boxes = tf.shape(boxes)[0]
def no_crop_check():
return (tf.random_uniform(shape=(), minval=0, maxval=1, dtype=tf.float32)
< constants.P_NO_CROP_PER_PASS)
def no_crop_proposal():
return (
tf.ones((), tf.bool),
tf.convert_to_tensor([0, 0, 1, 1], dtype=tf.float32),
tf.ones((num_boxes,), tf.bool),
)
def crop_proposal():
rand_vec = lambda minval, maxval: tf.random_uniform(
shape=(constants.NUM_CROP_PASSES, 1), minval=minval, maxval=maxval,
dtype=tf.float32)
width, height = rand_vec(0.3, 1), rand_vec(0.3, 1)
left, top = rand_vec(0, 1-width), rand_vec(0, 1-height)
right = left + width
bottom = top + height
ltrb = tf.concat([left, top, right, bottom], axis=1)
min_iou = tf.random_shuffle(constants.CROP_MIN_IOU_CHOICES)[0]
ious = calc_iou_tensor(ltrb, boxes)
# discard any bboxes whose center not in the cropped image
xc, yc = [tf.tile(0.5 * (boxes[:, i + 0] + boxes[:, i + 2])[tf.newaxis, :],
(constants.NUM_CROP_PASSES, 1)) for i in range(2)]
masks = tf.reduce_all(tf.stack([
tf.greater(xc, tf.tile(left, (1, num_boxes))),
tf.less(xc, tf.tile(right, (1, num_boxes))),
tf.greater(yc, tf.tile(top, (1, num_boxes))),
tf.less(yc, tf.tile(bottom, (1, num_boxes))),
], axis=2), axis=2)
# Checks of whether a crop is valid.
valid_aspect = tf.logical_and(tf.less(height/width, 2),
tf.less(width/height, 2))
valid_ious = tf.reduce_all(tf.greater(ious, min_iou), axis=1, keepdims=True)
valid_masks = tf.reduce_any(masks, axis=1, keepdims=True)
valid_all = tf.cast(tf.reduce_all(tf.concat(
[valid_aspect, valid_ious, valid_masks], axis=1), axis=1), tf.int32)
# One indexed, as zero is needed for the case of no matches.
index = tf.range(1, 1 + constants.NUM_CROP_PASSES, dtype=tf.int32)
# Either one-hot, or zeros if there is no valid crop.
selection = tf.equal(tf.reduce_max(index * valid_all), index)
use_crop = tf.reduce_any(selection)
output_ltrb = tf.reduce_sum(tf.multiply(ltrb, tf.tile(tf.cast(
selection, tf.float32)[:, tf.newaxis], (1, 4))), axis=0)
output_masks = tf.reduce_any(tf.logical_and(masks, tf.tile(
selection[:, tf.newaxis], (1, num_boxes))), axis=0)
return use_crop, output_ltrb, output_masks
def proposal(*args):
return tf.cond(
pred=no_crop_check(),
true_fn=no_crop_proposal,
false_fn=crop_proposal,
)
_, crop_bounds, box_masks = tf.while_loop(
cond=lambda x, *_: tf.logical_not(x),
body=proposal,
loop_vars=[tf.zeros((), tf.bool), tf.zeros((4,), tf.float32), tf.zeros((num_boxes,), tf.bool)],
)
filtered_boxes = tf.boolean_mask(boxes, box_masks, axis=0)
# Clip boxes to the cropped region.
filtered_boxes = tf.stack([
tf.maximum(filtered_boxes[:, 0], crop_bounds[0]),
tf.maximum(filtered_boxes[:, 1], crop_bounds[1]),
tf.minimum(filtered_boxes[:, 2], crop_bounds[2]),
tf.minimum(filtered_boxes[:, 3], crop_bounds[3]),
], axis=1)
left = crop_bounds[0]
top = crop_bounds[1]
width = crop_bounds[2] - left
height = crop_bounds[3] - top
cropped_boxes = tf.stack([
(filtered_boxes[:, 0] - left) / width,
(filtered_boxes[:, 1] - top) / height,
(filtered_boxes[:, 2] - left) / width,
(filtered_boxes[:, 3] - top) / height,
], axis=1)
cropped_image = tf.image.crop_and_resize(
image=image[tf.newaxis, :, :, :],
boxes=crop_bounds[tf.newaxis, :],
box_indices=tf.zeros((1,), tf.int32),
crop_size=(constants.IMAGE_SIZE, constants.IMAGE_SIZE),
)[0, :, :, :]
cropped_classes = tf.boolean_mask(classes, box_masks, axis=0)
return cropped_image, cropped_boxes, cropped_classes
def __init__(self, variable, mesh_impl):
"""Create a LaidOutVariable.
Args:
variable: a Variable (Operation)
mesh_impl: a MeshImpl
"""
self._variable = variable
self._mesh_impl = mesh_impl
shape = variable.outputs[0].shape
slice_shape = mesh_impl.slice_shape(shape)
base_name = variable.name
slices = []
slices_with_master_dtype = []
with tf.device(
variable.master_device), utils.outside_all_rewrites():
zero_tensor = tf.zeros(slice_shape, dtype=variable.slice_dtype)
# pylint: disable=protected-access
init_device_stack = tf.get_default_graph()._device_function_stack
if not mesh_impl.graph_device_function_stacks:
for pnum in xrange(mesh_impl.size):
tpu_device = mesh_impl.device_assignment.tpu_device(
replica=pnum)
with tf.device(tpu_device):
mesh_impl.graph_device_function_stacks.append(
tf.get_default_graph()._device_function_stack.copy(
))
for physical_pnum in xrange(mesh_impl.size):
slice_var_name = base_name + "_slice_%d" % physical_pnum
# Use tf.Variable instead of tf.get_variable since latter adds lots of
# useless operations to the TF graph. Use tf.get_variable only if
# in a AUTO_REUSE scope.
# Note: Repeatedly 'with tf.device():' slows down the graph
# construction. Therefore we directly use the cached device_stack here.
tf.get_default_graph()._device_function_stack = (
mesh_impl.graph_device_function_stacks[physical_pnum])
if tf.get_variable_scope().reuse == tf.AUTO_REUSE:
slice_var = tf.get_variable(
initializer=zero_tensor,
trainable=self._variable.trainable,
collections=["TPU_VAR"],
dtype=variable.slice_dtype,
name=slice_var_name)
else:
slice_var = tf.Variable(initial_value=zero_tensor,
trainable=self._variable.trainable,
collections=["TPU_VAR"],
dtype=variable.slice_dtype,
name=slice_var_name,
expected_shape=slice_shape)
slices.append(slice_var)
# Restore the initial stack
tf.get_default_graph()._device_function_stack = init_device_stack
# pylint: enable=protected-access
self._laid_out_tensor = mesh_impl.LaidOutTensor(
[tpu_variables.ReplicatedVariable(base_name, slices)])
with tf.device(
variable.master_device), utils.outside_all_rewrites():
if os.environ.get("MTF_SEQUENCE_MODE", "") == "1":
if mesh_impl.copy_master_to_slice_ops:
with tf.control_dependencies(
[mesh_impl.copy_master_to_slice_ops[-1]]):
self._copy_master_to_slices = self._gen_copy_master_to_slices_op(
variable.get_master(), shape, slices,
slice_shape)
else:
self._copy_master_to_slices = self._gen_copy_master_to_slices_op(
variable.get_master(), shape, slices, slice_shape)
mesh_impl.copy_master_to_slice_ops.append(
self._copy_master_to_slices)
else:
self._copy_master_to_slices = self._gen_copy_master_to_slices_op(
variable.get_master(), shape, slices, slice_shape)
slices_with_master_dtype = [
tf.cast(s, variable.master_dtype) for s in slices
]
slices_with_master_dtype = [
slices_with_master_dtype[mesh_impl.l2p(logical_pnum)]
for logical_pnum in range(mesh_impl.size)
]
self._copy_slices_to_master = variable.assign_to_master(
mesh_impl.combine_slices(slices_with_master_dtype,
shape,
device=variable.master_device))
def _compute_model_loss(
self, input_sequence, output_sequence, sequence_length, control_sequence):
"""Builds a model with loss for train/eval."""
hparams = self.hparams
batch_size = hparams.batch_size
input_sequence = tf.to_float(input_sequence)
output_sequence = tf.to_float(output_sequence)
max_seq_len = tf.minimum(tf.shape(output_sequence)[1], hparams.max_seq_len)
input_sequence = input_sequence[:, :max_seq_len]
if control_sequence is not None:
control_depth = control_sequence.shape[-1]
control_sequence = tf.to_float(control_sequence)
control_sequence = control_sequence[:, :max_seq_len]
# Shouldn't be necessary, but the slice loses shape information when
# control depth is zero.
control_sequence.set_shape([batch_size, None, control_depth])
# The target/expected outputs.
x_target = output_sequence[:, :max_seq_len]
# Inputs to be fed to decoder, including zero padding for the initial input.
x_input = tf.pad(output_sequence[:, :max_seq_len - 1],
[(0, 0), (1, 0), (0, 0)])
x_length = tf.minimum(sequence_length, max_seq_len)
# Either encode to get `z`, or do unconditional, decoder-only.
if hparams.z_size: # vae mode:
q_z = self.encode(input_sequence, x_length, control_sequence)
z = q_z.sample()
# Prior distribution.
p_z = ds.MultivariateNormalDiag(
loc=[0.] * hparams.z_size, scale_diag=[1.] * hparams.z_size)
# KL Divergence (nats)
kl_div = ds.kl_divergence(q_z, p_z)
# Concatenate the Z vectors to the inputs at each time step.
else: # unconditional, decoder-only generation
kl_div = tf.zeros([batch_size, 1], dtype=tf.float32)
z = None
r_loss, metric_map = self.decoder.reconstruction_loss(
x_input, x_target, x_length, z, control_sequence)[0:2]
free_nats = hparams.free_bits * tf.math.log(2.0)
kl_cost = tf.maximum(kl_div - free_nats, 0)
beta = ((1.0 - tf.pow(hparams.beta_rate, tf.to_float(self.global_step)))
* hparams.max_beta)
self.loss = tf.reduce_mean(r_loss) + beta * tf.reduce_mean(kl_cost)
scalars_to_summarize = {
'loss': self.loss,
'losses/r_loss': r_loss,
'losses/kl_loss': kl_cost,
'losses/kl_bits': kl_div / tf.math.log(2.0),
'losses/kl_beta': beta,
}
return metric_map, scalars_to_summarize
def _generate(self, feature_map_shape_list):
num_anchors = sum(
[shape[0] * shape[1] for shape in feature_map_shape_list])
return box_list.BoxList(tf.zeros((num_anchors, 4), dtype=tf.float32))
caps1_output_expanded = tf.expand_dims(caps1_output, -1,
name="caps1_output_expanded")
caps1_output_tile = tf.expand_dims(caps1_output_expanded, 2,
name="caps1_output_tile")
caps1_output_tiled = tf.tile(caps1_output_tile, [1, 1, caps2_n_caps, 1, 1],
name="caps1_output_tiled")
caps2_predicted = tf.matmul(W_tiled, caps1_output_tiled,
name="caps2_predicted")
"""## Routing by agreement
First let's initialize the raw routing weights $b_{i,j}$ to zero:
"""
raw_weights = tf.zeros([batch_size, caps1_n_caps, caps2_n_caps, 1, 1],
dtype=np.float32, name="raw_weights")
"""### Round 1"""
routing_weights = tf.nn.softmax(raw_weights, axis=2, name="routing_weights")
weighted_predictions = tf.multiply(routing_weights, caps2_predicted,
name="weighted_predictions")
weighted_sum = tf.reduce_sum(weighted_predictions, axis=1, keep_dims=True,
name="weighted_sum")
caps2_output_round_1 = squash(weighted_sum, axis=-2,
name="caps2_output_round_1")
"""### Round 2"""
def compute_prior_log_prob(self,
z_q,
targets_mask,
decoder_self_attention_bias,
check_invertibility=False,
**kwargs):
hparams = self._hparams
batch_size, targets_max_length = (
common_layers.shape_list(targets_mask)[:2])
prior_shape = [batch_size, targets_max_length, hparams.latent_size]
log_abs_det = tf.zeros([batch_size])
if hparams.prior_type == "standard_normal":
log_p_z_base = gops.standard_normal_density(z_q, targets_mask)
elif hparams.prior_type == "diagonal_normal":
diag_prior_params = ops.cond_prior("diag_prior", hparams,
tf.zeros(prior_shape),
targets_mask,
hparams.latent_size * 2,
decoder_self_attention_bias,
**kwargs)
p_dist = gops.diagonal_normal(diag_prior_params, "diag_prior")
log_p_z_base = p_dist.log_prob(z_q) # [B, L, C]
log_p_z_base = gops.reduce_sum_over_lc(log_p_z_base,
targets_mask) # [B]
elif hparams.prior_type in ["affine", "additive", "rq"]:
if self.is_evaluating:
disable_dropout = True
init = False
elif self.is_training:
disable_dropout = False
init = tf.equal(hparams.kl_startup_steps,
tf.cast(tf.train.get_global_step(), tf.int32))
else:
raise ValueError(
"compute_prior shouldn't be used in decoding.")
z_inv, log_abs_det, log_p_z_base, zs = glow.glow(
"glow",
z_q,
targets_mask,
decoder_self_attention_bias,
inverse=False,
init=init,
hparams=self._fparams,
disable_dropout=disable_dropout,
**kwargs)
if self.is_evaluating and check_invertibility:
z_inv_inv, _, _, _ = glow.glow("glow",
z_inv,
targets_mask,
decoder_self_attention_bias,
inverse=True,
split_zs=zs,
init=False,
hparams=self._fparams,
disable_dropout=True,
**kwargs)
z_diff = z_q - z_inv_inv
tf.summary.scalar("flow_recon_forward",
tf.reduce_max(tf.abs(z_diff)))
return log_p_z_base, log_abs_det
def test_autoaugment(self):
"""Smoke test to be sure no syntax errors."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
aug_image = autoaugment.distort_image_with_autoaugment(image, 'v0')
self.assertEqual((224, 224, 3), aug_image.shape)
def build_model(self, hps):
"""Define model architecture."""
if hps.is_training:
self.global_step = tf.Variable(0, name='global_step', trainable=False)
if hps.dec_model == 'lstm':
cell_fn = rnn.LSTMCell
elif hps.dec_model == 'layer_norm':
cell_fn = rnn.LayerNormLSTMCell
elif hps.dec_model == 'hyper':
cell_fn = rnn.HyperLSTMCell
else:
assert False, 'please choose a respectable cell'
if hps.enc_model == 'lstm':
enc_cell_fn = rnn.LSTMCell
elif hps.enc_model == 'layer_norm':
enc_cell_fn = rnn.LayerNormLSTMCell
elif hps.enc_model == 'hyper':
enc_cell_fn = rnn.HyperLSTMCell
else:
assert False, 'please choose a respectable cell'
use_recurrent_dropout = self.hps.use_recurrent_dropout
use_input_dropout = self.hps.use_input_dropout
use_output_dropout = self.hps.use_output_dropout
cell = cell_fn(
hps.dec_rnn_size,
use_recurrent_dropout=use_recurrent_dropout,
dropout_keep_prob=self.hps.recurrent_dropout_prob)
if hps.conditional: # vae mode:
if hps.enc_model == 'hyper':
self.enc_cell_fw = enc_cell_fn(
hps.enc_rnn_size,
use_recurrent_dropout=use_recurrent_dropout,
dropout_keep_prob=self.hps.recurrent_dropout_prob)
self.enc_cell_bw = enc_cell_fn(
hps.enc_rnn_size,
use_recurrent_dropout=use_recurrent_dropout,
dropout_keep_prob=self.hps.recurrent_dropout_prob)
else:
self.enc_cell_fw = enc_cell_fn(
hps.enc_rnn_size,
use_recurrent_dropout=use_recurrent_dropout,
dropout_keep_prob=self.hps.recurrent_dropout_prob)
self.enc_cell_bw = enc_cell_fn(
hps.enc_rnn_size,
use_recurrent_dropout=use_recurrent_dropout,
dropout_keep_prob=self.hps.recurrent_dropout_prob)
# dropout:
tf.logging.info('Input dropout mode = %s.', use_input_dropout)
tf.logging.info('Output dropout mode = %s.', use_output_dropout)
tf.logging.info('Recurrent dropout mode = %s.', use_recurrent_dropout)
if use_input_dropout:
tf.logging.info('Dropout to input w/ keep_prob = %4.4f.',
self.hps.input_dropout_prob)
cell = contrib_rnn.DropoutWrapper(
cell, input_keep_prob=self.hps.input_dropout_prob)
if use_output_dropout:
tf.logging.info('Dropout to output w/ keep_prob = %4.4f.',
self.hps.output_dropout_prob)
cell = contrib_rnn.DropoutWrapper(
cell, output_keep_prob=self.hps.output_dropout_prob)
self.cell = cell
self.sequence_lengths = tf.placeholder(
dtype=tf.int32, shape=[self.hps.batch_size])
self.input_data = tf.placeholder(
dtype=tf.float32,
shape=[self.hps.batch_size, self.hps.max_seq_len + 1, 5])
# The target/expected vectors of strokes
self.output_x = self.input_data[:, 1:self.hps.max_seq_len + 1, :]
# vectors of strokes to be fed to decoder (same as above, but lagged behind
# one step to include initial dummy value of (0, 0, 1, 0, 0))
self.input_x = self.input_data[:, :self.hps.max_seq_len, :]
# either do vae-bit and get z, or do unconditional, decoder-only
if hps.conditional: # vae mode:
self.mean, self.presig = self.encoder(self.output_x,
self.sequence_lengths)
self.sigma = tf.exp(self.presig / 2.0) # sigma > 0. div 2.0 -> sqrt.
eps = tf.random_normal(
(self.hps.batch_size, self.hps.z_size), 0.0, 1.0, dtype=tf.float32)
self.batch_z = self.mean + tf.multiply(self.sigma, eps)
# KL cost
self.kl_cost = -0.5 * tf.reduce_mean(
(1 + self.presig - tf.square(self.mean) - tf.exp(self.presig)))
self.kl_cost = tf.maximum(self.kl_cost, self.hps.kl_tolerance)
pre_tile_y = tf.reshape(self.batch_z,
[self.hps.batch_size, 1, self.hps.z_size])
overlay_x = tf.tile(pre_tile_y, [1, self.hps.max_seq_len, 1])
actual_input_x = tf.concat([self.input_x, overlay_x], 2)
self.initial_state = tf.nn.tanh(
rnn.super_linear(
self.batch_z,
cell.state_size,
init_w='gaussian',
weight_start=0.001,
input_size=self.hps.z_size))
else: # unconditional, decoder-only generation
self.batch_z = tf.zeros(
(self.hps.batch_size, self.hps.z_size), dtype=tf.float32)
self.kl_cost = tf.zeros([], dtype=tf.float32)
actual_input_x = self.input_x
self.initial_state = cell.zero_state(
batch_size=hps.batch_size, dtype=tf.float32)
self.num_mixture = hps.num_mixture
# TODO(deck): Better understand this comment.
# Number of outputs is 3 (one logit per pen state) plus 6 per mixture
# component: mean_x, stdev_x, mean_y, stdev_y, correlation_xy, and the
# mixture weight/probability (Pi_k)
n_out = (3 + self.num_mixture * 6)
with tf.variable_scope('RNN'):
output_w = tf.get_variable('output_w', [self.hps.dec_rnn_size, n_out])
output_b = tf.get_variable('output_b', [n_out])
# decoder module of sketch-rnn is below
output, last_state = tf.nn.dynamic_rnn(
cell,
actual_input_x,
initial_state=self.initial_state,
time_major=False,
swap_memory=True,
dtype=tf.float32,
scope='RNN')
output = tf.reshape(output, [-1, hps.dec_rnn_size])
output = tf.nn.xw_plus_b(output, output_w, output_b)
self.final_state = last_state
# NB: the below are inner functions, not methods of Model
def tf_2d_normal(x1, x2, mu1, mu2, s1, s2, rho):
"""Returns result of eq # 24 of http://arxiv.org/abs/1308.0850."""
norm1 = tf.subtract(x1, mu1)
norm2 = tf.subtract(x2, mu2)
s1s2 = tf.multiply(s1, s2)
# eq 25
z = (tf.square(tf.div(norm1, s1)) + tf.square(tf.div(norm2, s2)) -
2 * tf.div(tf.multiply(rho, tf.multiply(norm1, norm2)), s1s2))
neg_rho = 1 - tf.square(rho)
result = tf.exp(tf.div(-z, 2 * neg_rho))
denom = 2 * np.pi * tf.multiply(s1s2, tf.sqrt(neg_rho))
result = tf.div(result, denom)
return result
def get_lossfunc(z_pi, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr,
z_pen_logits, x1_data, x2_data, pen_data):
"""Returns a loss fn based on eq #26 of http://arxiv.org/abs/1308.0850."""
# This represents the L_R only (i.e. does not include the KL loss term).
result0 = tf_2d_normal(x1_data, x2_data, z_mu1, z_mu2, z_sigma1, z_sigma2,
z_corr)
epsilon = 1e-6
# result1 is the loss wrt pen offset (L_s in equation 9 of
# https://arxiv.org/pdf/1704.03477.pdf)
result1 = tf.multiply(result0, z_pi)
result1 = tf.reduce_sum(result1, 1, keep_dims=True)
result1 = -tf.log(result1 + epsilon) # avoid log(0)
fs = 1.0 - pen_data[:, 2] # use training data for this
fs = tf.reshape(fs, [-1, 1])
# Zero out loss terms beyond N_s, the last actual stroke
result1 = tf.multiply(result1, fs)
# result2: loss wrt pen state, (L_p in equation 9)
result2 = tf.nn.softmax_cross_entropy_with_logits(
labels=pen_data, logits=z_pen_logits)
result2 = tf.reshape(result2, [-1, 1])
if not self.hps.is_training: # eval mode, mask eos columns
result2 = tf.multiply(result2, fs)
result = result1 + result2
return result
# below is where we need to do MDN (Mixture Density Network) splitting of
# distribution params
def get_mixture_coef(output):
"""Returns the tf slices containing mdn dist params."""
# This uses eqns 18 -> 23 of http://arxiv.org/abs/1308.0850.
z = output
z_pen_logits = z[:, 0:3] # pen states
z_pi, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr = tf.split(z[:, 3:], 6, 1)
# process output z's into MDN parameters
# softmax all the pi's and pen states:
z_pi = tf.nn.softmax(z_pi)
z_pen = tf.nn.softmax(z_pen_logits)
# exponentiate the sigmas and also make corr between -1 and 1.
z_sigma1 = tf.exp(z_sigma1)
z_sigma2 = tf.exp(z_sigma2)
z_corr = tf.tanh(z_corr)
r = [z_pi, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr, z_pen, z_pen_logits]
return r
out = get_mixture_coef(output)
[o_pi, o_mu1, o_mu2, o_sigma1, o_sigma2, o_corr, o_pen, o_pen_logits] = out
self.pi = o_pi
self.mu1 = o_mu1
self.mu2 = o_mu2
self.sigma1 = o_sigma1
self.sigma2 = o_sigma2
self.corr = o_corr
self.pen_logits = o_pen_logits
# pen state probabilities (result of applying softmax to self.pen_logits)
self.pen = o_pen
# reshape target data so that it is compatible with prediction shape
target = tf.reshape(self.output_x, [-1, 5])
[x1_data, x2_data, eos_data, eoc_data, cont_data] = tf.split(target, 5, 1)
pen_data = tf.concat([eos_data, eoc_data, cont_data], 1)
lossfunc = get_lossfunc(o_pi, o_mu1, o_mu2, o_sigma1, o_sigma2, o_corr,
o_pen_logits, x1_data, x2_data, pen_data)
self.r_cost = tf.reduce_mean(lossfunc)
if self.hps.is_training:
self.lr = tf.Variable(self.hps.learning_rate, trainable=False)
optimizer = tf.train.AdamOptimizer(self.lr)
self.kl_weight = tf.Variable(self.hps.kl_weight_start, trainable=False)
self.cost = self.r_cost + self.kl_cost * self.kl_weight
gvs = optimizer.compute_gradients(self.cost)
g = self.hps.grad_clip
capped_gvs = [(tf.clip_by_value(grad, -g, g), var) for grad, var in gvs]
self.train_op = optimizer.apply_gradients(
capped_gvs, global_step=self.global_step, name='train_step')
def parser(record):
"""function used to parse tfrecord."""
record_spec = {
"input": tf.FixedLenFeature([seq_len], tf.int64),
"target": tf.FixedLenFeature([seq_len], tf.int64),
"seg_id": tf.FixedLenFeature([seq_len], tf.int64),
"label": tf.FixedLenFeature([1], tf.int64),
"is_masked": tf.FixedLenFeature([seq_len], tf.int64),
}
# retrieve serialized example
example = tf.parse_single_example(serialized=record,
features=record_spec)
inputs = example.pop("input")
target = example.pop("target")
is_masked = tf.cast(example.pop("is_masked"), tf.bool)
non_reuse_len = seq_len - reuse_len
assert perm_size <= reuse_len and perm_size <= non_reuse_len
perm_mask_0, target_0, target_mask_0, input_k_0, input_q_0 = _local_perm(
inputs[:reuse_len], target[:reuse_len], is_masked[:reuse_len],
perm_size, reuse_len)
perm_mask_1, target_1, target_mask_1, input_k_1, input_q_1 = _local_perm(
inputs[reuse_len:], target[reuse_len:], is_masked[reuse_len:],
perm_size, non_reuse_len)
perm_mask_0 = tf.concat(
[perm_mask_0, tf.ones([reuse_len, non_reuse_len])], axis=1)
perm_mask_1 = tf.concat(
[tf.zeros([non_reuse_len, reuse_len]), perm_mask_1], axis=1)
perm_mask = tf.concat([perm_mask_0, perm_mask_1], axis=0)
target = tf.concat([target_0, target_1], axis=0)
target_mask = tf.concat([target_mask_0, target_mask_1], axis=0)
input_k = tf.concat([input_k_0, input_k_1], axis=0)
input_q = tf.concat([input_q_0, input_q_1], axis=0)
if num_predict is not None:
indices = tf.range(seq_len, dtype=tf.int64)
bool_target_mask = tf.cast(target_mask, tf.bool)
indices = tf.boolean_mask(indices, bool_target_mask)
##### extra padding due to CLS/SEP introduced after prepro
actual_num_predict = tf.shape(indices)[0]
pad_len = num_predict - actual_num_predict
##### target_mapping
target_mapping = tf.one_hot(indices, seq_len, dtype=tf.float32)
paddings = tf.zeros([pad_len, seq_len], dtype=target_mapping.dtype)
target_mapping = tf.concat([target_mapping, paddings], axis=0)
example["target_mapping"] = tf.reshape(target_mapping,
[num_predict, seq_len])
##### target
target = tf.boolean_mask(target, bool_target_mask)
paddings = tf.zeros([pad_len], dtype=target.dtype)
target = tf.concat([target, paddings], axis=0)
example["target"] = tf.reshape(target, [num_predict])
##### target mask
target_mask = tf.concat([
tf.ones([actual_num_predict], dtype=tf.float32),
tf.zeros([pad_len], dtype=tf.float32)
],
axis=0)
example["target_mask"] = tf.reshape(target_mask, [num_predict])
else:
example["target"] = tf.reshape(target, [seq_len])
example["target_mask"] = tf.reshape(target_mask, [seq_len])
# reshape back to fixed shape
example["perm_mask"] = tf.reshape(perm_mask, [seq_len, seq_len])
example["input_k"] = tf.reshape(input_k, [seq_len])
example["input_q"] = tf.reshape(input_q, [seq_len])
_convert_example(example, use_bfloat16)
for k, v in example.items():
logging.info("%s: %s", k, v)
return example
def model_fn(features, labels, mode, params, config):
"""Builds the acoustic model."""
del config
hparams = params
length = features.length
spec = features.spec
is_training = mode == tf.estimator.ModeKeys.TRAIN
if is_training:
onset_labels = labels.onsets
offset_labels = labels.offsets
velocity_labels = labels.velocities
frame_labels = labels.labels
frame_label_weights = labels.label_weights
if hparams.stop_activation_gradient and not hparams.activation_loss:
raise ValueError(
'If stop_activation_gradient is true, activation_loss must be true.'
)
losses = {}
with slim.arg_scope([slim.batch_norm, slim.dropout],
is_training=is_training):
with tf.variable_scope('onsets'):
onset_outputs = acoustic_model(spec,
hparams,
lstm_units=hparams.onset_lstm_units,
lengths=length,
is_training=is_training)
onset_probs = slim.fully_connected(onset_outputs,
constants.MIDI_PITCHES,
activation_fn=tf.sigmoid,
scope='onset_probs')
# onset_probs_flat is used during inference.
onset_probs_flat = flatten_maybe_padded_sequences(
onset_probs, length)
if is_training:
onset_labels_flat = flatten_maybe_padded_sequences(
onset_labels, length)
onset_losses = tf_utils.log_loss(onset_labels_flat,
onset_probs_flat)
tf.losses.add_loss(tf.reduce_mean(onset_losses))
losses['onset'] = onset_losses
with tf.variable_scope('offsets'):
offset_outputs = acoustic_model(
spec,
hparams,
lstm_units=hparams.offset_lstm_units,
lengths=length,
is_training=is_training)
offset_probs = slim.fully_connected(offset_outputs,
constants.MIDI_PITCHES,
activation_fn=tf.sigmoid,
scope='offset_probs')
# offset_probs_flat is used during inference.
offset_probs_flat = flatten_maybe_padded_sequences(
offset_probs, length)
if is_training:
offset_labels_flat = flatten_maybe_padded_sequences(
offset_labels, length)
offset_losses = tf_utils.log_loss(offset_labels_flat,
offset_probs_flat)
tf.losses.add_loss(tf.reduce_mean(offset_losses))
losses['offset'] = offset_losses
with tf.variable_scope('velocity'):
velocity_outputs = acoustic_model(
spec,
hparams,
lstm_units=hparams.velocity_lstm_units,
lengths=length,
is_training=is_training)
velocity_values = slim.fully_connected(velocity_outputs,
constants.MIDI_PITCHES,
activation_fn=None,
scope='onset_velocities')
velocity_values_flat = flatten_maybe_padded_sequences(
velocity_values, length)
if is_training:
velocity_labels_flat = flatten_maybe_padded_sequences(
velocity_labels, length)
velocity_loss = tf.reduce_sum(
onset_labels_flat *
tf.square(velocity_labels_flat - velocity_values_flat),
axis=1)
tf.losses.add_loss(tf.reduce_mean(velocity_loss))
losses['velocity'] = velocity_loss
with tf.variable_scope('frame'):
if not hparams.share_conv_features:
# TODO(eriche): this is broken when hparams.frame_lstm_units > 0
activation_outputs = acoustic_model(
spec,
hparams,
lstm_units=hparams.frame_lstm_units,
lengths=length,
is_training=is_training)
activation_probs = slim.fully_connected(
activation_outputs,
constants.MIDI_PITCHES,
activation_fn=tf.sigmoid,
scope='activation_probs')
else:
activation_probs = slim.fully_connected(
onset_outputs,
constants.MIDI_PITCHES,
activation_fn=tf.sigmoid,
scope='activation_probs')
probs = []
if hparams.stop_onset_gradient:
probs.append(tf.stop_gradient(onset_probs))
else:
probs.append(onset_probs)
if hparams.stop_activation_gradient:
probs.append(tf.stop_gradient(activation_probs))
else:
probs.append(activation_probs)
if hparams.stop_offset_gradient:
probs.append(tf.stop_gradient(offset_probs))
else:
probs.append(offset_probs)
combined_probs = tf.concat(probs, 2)
if hparams.combined_lstm_units > 0:
outputs = lstm_layer(
combined_probs,
hparams.batch_size,
hparams.combined_lstm_units,
lengths=length if hparams.use_lengths else None,
stack_size=hparams.combined_rnn_stack_size,
use_cudnn=hparams.use_cudnn,
is_training=is_training,
bidirectional=hparams.bidirectional)
else:
outputs = combined_probs
frame_probs = slim.fully_connected(outputs,
constants.MIDI_PITCHES,
activation_fn=tf.sigmoid,
scope='frame_probs')
frame_probs_flat = flatten_maybe_padded_sequences(frame_probs, length)
if is_training:
frame_labels_flat = flatten_maybe_padded_sequences(
frame_labels, length)
frame_label_weights_flat = flatten_maybe_padded_sequences(
frame_label_weights, length)
if hparams.weight_frame_and_activation_loss:
frame_loss_weights = frame_label_weights_flat
else:
frame_loss_weights = None
frame_losses = tf_utils.log_loss(frame_labels_flat,
frame_probs_flat,
weights=frame_loss_weights)
tf.losses.add_loss(tf.reduce_mean(frame_losses))
losses['frame'] = frame_losses
if hparams.activation_loss:
if hparams.weight_frame_and_activation_loss:
activation_loss_weights = frame_label_weights
else:
activation_loss_weights = None
activation_losses = tf_utils.log_loss(
frame_labels_flat,
flatten_maybe_padded_sequences(activation_probs, length),
weights=activation_loss_weights)
tf.losses.add_loss(tf.reduce_mean(activation_losses))
losses['activation'] = activation_losses
frame_predictions = frame_probs_flat > hparams.predict_frame_threshold
onset_predictions = onset_probs_flat > hparams.predict_onset_threshold
offset_predictions = offset_probs_flat > hparams.predict_offset_threshold
frame_predictions = tf.expand_dims(frame_predictions, axis=0)
onset_predictions = tf.expand_dims(onset_predictions, axis=0)
offset_predictions = tf.expand_dims(offset_predictions, axis=0)
velocity_values = tf.expand_dims(velocity_values_flat, axis=0)
metrics_values = metrics.define_metrics(
frame_probs=frame_probs,
onset_probs=onset_probs,
frame_predictions=frame_predictions,
onset_predictions=onset_predictions,
offset_predictions=offset_predictions,
velocity_values=velocity_values,
length=features.length,
sequence_label=labels.note_sequence,
frame_labels=labels.labels,
sequence_id=features.sequence_id,
hparams=hparams)
for label, loss_collection in losses.items():
loss_label = 'losses/' + label
metrics_values[loss_label] = loss_collection
def predict_sequence():
"""Convert frame predictions into a sequence (TF)."""
def _predict(frame_probs, onset_probs, frame_predictions,
onset_predictions, offset_predictions, velocity_values):
"""Convert frame predictions into a sequence (Python)."""
sequence = infer_util.predict_sequence(
frame_probs=frame_probs,
onset_probs=onset_probs,
frame_predictions=frame_predictions,
onset_predictions=onset_predictions,
offset_predictions=offset_predictions,
velocity_values=velocity_values,
hparams=hparams,
min_pitch=constants.MIN_MIDI_PITCH)
return sequence.SerializeToString()
sequence = tf.py_func(_predict,
inp=[
frame_probs[0],
onset_probs[0],
frame_predictions[0],
onset_predictions[0],
offset_predictions[0],
velocity_values[0],
],
Tout=tf.string,
stateful=False)
sequence.set_shape([])
return tf.expand_dims(sequence, axis=0)
predictions = {
'frame_probs': frame_probs,
'onset_probs': onset_probs,
'frame_predictions': frame_predictions,
'onset_predictions': onset_predictions,
'offset_predictions': offset_predictions,
'velocity_values': velocity_values,
'sequence_predictions': predict_sequence(),
# Include some features and labels in output because Estimator 'predict'
# API does not give access to them.
'sequence_ids': features.sequence_id,
'sequence_labels': labels.note_sequence,
'frame_labels': labels.labels,
'onset_labels': labels.onsets,
}
for k, v in metrics_values.items():
predictions[k] = tf.stack(v)
metric_ops = {k: tf.metrics.mean(v) for k, v in metrics_values.items()}
train_op = None
loss = None
if is_training:
# Creates a pianoroll labels in red and probs in green [minibatch, 88]
images = {}
onset_pianorolls = tf.concat([
onset_labels[:, :, :, tf.newaxis], onset_probs[:, :, :,
tf.newaxis],
tf.zeros(tf.shape(onset_labels))[:, :, :, tf.newaxis]
],
axis=3)
images['OnsetPianorolls'] = onset_pianorolls
offset_pianorolls = tf.concat([
offset_labels[:, :, :, tf.newaxis], offset_probs[:, :, :,
tf.newaxis],
tf.zeros(tf.shape(offset_labels))[:, :, :, tf.newaxis]
],
axis=3)
images['OffsetPianorolls'] = offset_pianorolls
activation_pianorolls = tf.concat([
frame_labels[:, :, :, tf.newaxis], frame_probs[:, :, :,
tf.newaxis],
tf.zeros(tf.shape(frame_labels))[:, :, :, tf.newaxis]
],
axis=3)
images['ActivationPianorolls'] = activation_pianorolls
for name, image in images.items():
tf.summary.image(name, image)
loss = tf.losses.get_total_loss()
tf.summary.scalar('loss', loss)
for label, loss_collection in losses.items():
loss_label = 'losses/' + label
tf.summary.scalar(loss_label, tf.reduce_mean(loss_collection))
train_op = contrib_layers.optimize_loss(
name='training',
loss=loss,
global_step=tf.train.get_or_create_global_step(),
learning_rate=hparams.learning_rate,
learning_rate_decay_fn=functools.partial(
tf.train.exponential_decay,
decay_steps=hparams.decay_steps,
decay_rate=hparams.decay_rate,
staircase=True),
clip_gradients=hparams.clip_norm,
optimizer='Adam')
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op,
eval_metric_ops=metric_ops)
Y.append(input_data[i + window_Size])
i += 1
assert len(X) == len(Y)
return X, Y
X, Y = create_input()
X_train, X_test = X[:int(0.8 * l)], X[int(0.8 * l):]
Y_train, Y_test = Y[:int(0.8 * l)], Y[int(0.8 * l):]
#weights for input gate
weights_input_gate = tf.Variable(
tf.truncated_normal([1, hidden_layer], stddev=0.05))
weights_input_hidden = tf.Variable(
tf.truncated_normal([hidden_layer, hidden_layer], stddev=0.05))
bias_input = tf.Variable(tf.zeros([hidden_layer]))
#weights for forgot gate
weights_forgot_gate = tf.Variable(
tf.truncated_normal([1, hidden_layer], stddev=0.05))
weights_forget_hidden = tf.Variable(
tf.truncated_normal([hidden_layer, hidden_layer], stddev=0.05))
bias_forget = tf.Variable(tf.zeros([hidden_layer]))
#weights for output gate
weights_output_gate = tf.Variable(
tf.truncated_normal([1, hidden_layer], stddev=0.05))
weights_output_hidden = tf.Variable(
tf.truncated_normal([hidden_layer, hidden_layer], stddev=0.05))
bias_output = tf.Variable(tf.zeros([hidden_layer]))
def __init__(self, net_params, batch_size, num_features, num_classes,
prediction_window, seq_len):
for key in net_params.keys():
net_params[key] = _utils.convert_shared_float_array_to_numpy(
net_params[key])
# Suppresses verbosity to only errors
_tf.logging.set_verbosity(_tf.logging.ERROR)
_tf.debugging.set_log_device_placement(False)
_tf.reset_default_graph()
self.num_classes = num_classes
self.batch_size = batch_size
self.seq_len = seq_len
# Vars
self.data = _tf.placeholder(
_tf.float32, [None, prediction_window * seq_len, num_features])
self.weight = _tf.placeholder(_tf.float32, [None, seq_len, 1])
self.target = _tf.placeholder(_tf.int32, [None, seq_len, 1])
self.is_training = _tf.placeholder(_tf.bool)
# Reshaping weights
reshaped_weight = _tf.reshape(self.weight, [self.batch_size, seq_len])
# One hot encoding target
reshaped_target = _tf.reshape(self.target, [self.batch_size, seq_len])
one_hot_target = _tf.one_hot(reshaped_target,
depth=self.num_classes,
axis=-1)
# Weights
self.weights = {
'conv_weight':
_tf.Variable(_tf.zeros([prediction_window, num_features, CONV_H]),
name='conv_weight'),
'dense0_weight':
_tf.Variable(_tf.zeros([LSTM_H, DENSE_H]), name='dense0_weight'),
'dense1_weight':
_tf.Variable(_tf.zeros([DENSE_H, self.num_classes]),
name='dense1_weight')
}
# Biases
self.biases = {
'conv_bias':
_tf.Variable(_tf.zeros([CONV_H]), name='conv_bias'),
'dense0_bias':
_tf.Variable(_tf.zeros([DENSE_H]), name='dense0_bias'),
'dense1_bias':
_tf.Variable(_tf.zeros([num_classes]), name='dense1_bias')
}
# Convolution
conv = _tf.nn.conv1d(self.data,
self.weights['conv_weight'],
stride=prediction_window,
padding='SAME')
conv = _tf.nn.bias_add(conv, self.biases['conv_bias'])
conv = _tf.nn.relu(conv)
dropout = _tf.layers.dropout(conv, rate=0.2, training=self.is_training)
# Long Stem Term Memory
lstm = self.load_lstm_weights_params(net_params)
cells = _tf.nn.rnn_cell.LSTMCell(num_units=LSTM_H,
reuse=_tf.AUTO_REUSE,
forget_bias=0.0,
initializer=_tf.initializers.constant(
lstm, verify_shape=True))
init_state = cells.zero_state(batch_size, _tf.float32)
rnn_outputs, final_state = _tf.nn.dynamic_rnn(cells,
dropout,
initial_state=init_state)
# Dense
dense = _tf.reshape(rnn_outputs, (-1, LSTM_H))
dense = _tf.add(_tf.matmul(dense, self.weights['dense0_weight']),
self.biases['dense0_bias'])
dense = _tf.layers.batch_normalization(
inputs=dense,
beta_initializer=_tf.initializers.constant(net_params['bn_beta'],
verify_shape=True),
gamma_initializer=_tf.initializers.constant(net_params['bn_gamma'],
verify_shape=True),
moving_mean_initializer=_tf.initializers.constant(
net_params['bn_running_mean'], verify_shape=True),
moving_variance_initializer=_tf.initializers.constant(
net_params['bn_running_var'], verify_shape=True),
training=self.is_training)
dense = _tf.nn.relu(dense)
dense = _tf.layers.dropout(dense, rate=0.5, training=self.is_training)
# Output
out = _tf.add(_tf.matmul(dense, self.weights['dense1_weight']),
self.biases['dense1_bias'])
out = _tf.reshape(out, (-1, self.seq_len, self.num_classes))
self.probs = _tf.nn.softmax(out)
# Weights
seq_sum_weights = _tf.reduce_sum(reshaped_weight, axis=1)
binary_seq_sum_weights = _tf.reduce_sum(
_tf.cast(seq_sum_weights > 0, dtype=_tf.float32))
# Loss
loss = _tf.losses.softmax_cross_entropy(
logits=out,
onehot_labels=one_hot_target,
weights=reshaped_weight,
reduction=_tf.losses.Reduction.NONE)
self.loss_per_seq = _tf.reduce_sum(loss,
axis=1) / (seq_sum_weights + 1e-5)
self.loss_op = _tf.reduce_sum(
self.loss_per_seq) / (binary_seq_sum_weights + 1e-5)
# Optimizer
update_ops = _tf.get_collection(_tf.GraphKeys.UPDATE_OPS)
self.set_learning_rate(1e-3)
train_op = self.optimizer.minimize(self.loss_op)
self.train_op = _tf.group([train_op, update_ops])
# Session
self.sess = _tf.Session()
# Initialize all variables
self.sess.run(_tf.global_variables_initializer())
self.sess.run(_tf.local_variables_initializer())
self.load_weights(net_params)
def __init__(self,
*,
lr_kwargs,
layers=None,
preprocessors=None,
learn_std=True,
std_value=0.1,
mu_range=None,
log_std_range=None,
**kwargs):
super(GaussianPolicy, self).__init__(**kwargs)
self.std_value = std_value
self.lr_kwargs = lr_kwargs
self.lr_scheduler = get_scheduler(lr_kwargs)
self.schedulers += (self.lr_scheduler, )
self.learn_std = learn_std
self.mu_range = (-2.0, 2.0) if mu_range is None else mu_range
self.log_std_range = (-10,
0.3) if log_std_range is None else log_std_range
assert not self.discrete_action_space, "Action space for the Gaussian Policy must be continuous!"
# Placeholders
self.lr_ph = tf_v1.placeholder("float32", shape=(), name="lr_ph")
# Create model
if layers is None:
layers = DEFAULT_LAYERS
self.layers = layers
self.preprocessors = preprocessors
self.base_model = NeuralNetwork(self.scope,
input_shapes=[self.obs_shape],
layers=self.layers,
preprocessors=self.preprocessors)
self.mu = tf_v1.keras.layers.Dense(self.action_size, activation=None)(
self.base_model.output)
self.mu = tf_v1.clip_by_value(self.mu, *self.mu_range)
if self.learn_std:
self.log_std = tf_v1.keras.layers.Dense(self.action_size)(
self.base_model.output)
self.log_std = tf_v1.clip_by_value(self.log_std,
*self.log_std_range)
self.std = tf_v1.exp(self.log_std)
self.raw_action_model = tf_v1.keras.Model(
inputs=[self.base_model.input], outputs=[self.mu, self.std])
else:
self.std = tf_v1.constant([std_value] * self.action_size,
dtype="float32")
self.raw_action_model = tf_v1.keras.Model(
inputs=[self.base_model.input], outputs=[self.mu])
batch_size = tf_v1.shape(self.mu)[0]
norm_dist = tfd.Normal(loc=tf_v1.zeros(self.action_size),
scale=tf_v1.ones(self.action_size))
z = norm_dist.sample(batch_size)
raw_actions = self.mu + z * self.std # Reparameterization trick
self.actions = tf_v1.tanh(raw_actions)
self.deterministic_actions = tf_v1.tanh(self.mu)
self.model = tf_v1.keras.Model(inputs=[self.base_model.input],
outputs=[self.actions])
# Loss parameters
self._loss = None
self.train_op = None
# Summary parameters
self.scalar_summaries += ("lr", )
self.scalar_summaries_tf += ("loss", "mean_log_actions", "min_mu",
"mean_mu", "max_mu", "min_std",
"mean_std", "max_std")
self.histogram_summaries_tf += ("actions", "mu", "std", "log_actions")
def reconstruction_loss(self, x_input, x_target, x_length, z=None,
c_input=None):
"""Reconstruction loss calculation.
Args:
x_input: Batch of decoder input sequences of concatenated segmeents for
teacher forcing, sized `[batch_size, max_seq_len, output_depth]`.
x_target: Batch of expected output sequences to compute loss against,
sized `[batch_size, max_seq_len, output_depth]`.
x_length: Length of input/output sequences, sized
`[batch_size, level_lengths[0]]` or `[batch_size]`. If the latter,
each length must either equal `max_seq_len` or 0. In this case, the
segment lengths are assumed to be constant and the total length will be
evenly divided amongst the segments.
z: (Optional) Latent vectors. Required if model is conditional. Sized
`[n, z_size]`.
c_input: (Optional) Batch of control sequences, sized
`[batch_size, max_seq_len, control_depth]`. Required if conditioning on
control sequences.
Returns:
r_loss: The reconstruction loss for each sequence in the batch.
metric_map: Map from metric name to tf.metrics return values for logging.
decode_results: The LstmDecodeResults.
Raises:
ValueError: If `c_input` is provided in re-encoder mode.
"""
if self._hierarchical_encoder and c_input is not None:
raise ValueError(
'Re-encoder mode unsupported when conditioning on controls.')
batch_size = int(x_input.shape[0])
x_length = lstm_utils.maybe_split_sequence_lengths(
x_length, np.prod(self._level_lengths[:-1]), self._total_length)
hier_input = self._reshape_to_hierarchy(x_input)
hier_target = self._reshape_to_hierarchy(x_target)
hier_length = self._reshape_to_hierarchy(x_length)
hier_control = (
self._reshape_to_hierarchy(c_input) if c_input is not None else None)
loss_outputs = []
def base_train_fn(embedding, hier_index):
"""Base function for training hierarchical decoder."""
split_size = self._level_lengths[-1]
split_input = hier_input[hier_index]
split_target = hier_target[hier_index]
split_length = hier_length[hier_index]
split_control = (
hier_control[hier_index] if hier_control is not None else None)
res = self._core_decoder.reconstruction_loss(
split_input, split_target, split_length, embedding, split_control)
loss_outputs.append(res)
decode_results = res[-1]
if self._hierarchical_encoder:
# Get the approximate "sample" from the model.
# Start with the inputs the RNN saw (excluding the start token).
samples = decode_results.rnn_input[:, 1:]
# Pad to be the max length.
samples = tf.pad(
samples,
[(0, 0), (0, split_size - tf.shape(samples)[1]), (0, 0)])
samples.set_shape([batch_size, split_size, self._output_depth])
# Set the final value based on the target, since the scheduled sampling
# helper does not sample the final value.
samples = lstm_utils.set_final(
samples,
split_length,
lstm_utils.get_final(split_target, split_length, time_major=False),
time_major=False)
# Return the re-encoded sample.
return self._hierarchical_encoder.level(0).encode(
sequence=samples,
sequence_length=split_length)
elif self._disable_autoregression:
return None
else:
return tf.concat(tf.nest.flatten(decode_results.final_state), axis=-1)
z = tf.zeros([batch_size, 0]) if z is None else z
self._hierarchical_decode(z, base_train_fn)
# Accumulate the split sequence losses.
r_losses, metric_maps, decode_results = list(zip(*loss_outputs))
# Merge the metric maps by passing through renamed values and taking the
# mean across the splits.
merged_metric_map = {}
for metric_name in metric_maps[0]:
metric_values = []
for i, m in enumerate(metric_maps):
merged_metric_map['segment/%03d/%s' % (i, metric_name)] = m[metric_name]
metric_values.append(m[metric_name][0])
merged_metric_map[metric_name] = (
tf.reduce_mean(metric_values), tf.no_op())
return (tf.reduce_sum(r_losses, axis=0),
merged_metric_map,
self._merge_decode_results(decode_results))
def init_vector(self, shape):
return tf.zeros(shape)
