in tensorflow_estimator/python/estimator/canned/timeseries/ar_model.py [0:0]
def predict(self, features):
"""Computes predictions multiple steps into the future.
Args:
features: A dictionary with the following key/value pairs:
PredictionFeatures.TIMES: A [batch size, predict window size] integer
Tensor of times, after the window of data indicated by `STATE_TUPLE`,
to make predictions for.
PredictionFeatures.STATE_TUPLE: A tuple of (times, values), times with
shape [batch size, self.input_window_size], values with shape [batch
size, self.input_window_size, self.num_features] representing a
segment of the time series before `TIMES`. This data is used to start
of the autoregressive computation. This should have data for at least
self.input_window_size timesteps. And any exogenous features, with
shapes prefixed by shape of `TIMES`.
Returns:
A dictionary with keys, "mean", "covariance". The
values are Tensors of shape [batch_size, predict window size,
num_features] and correspond to the values passed in `TIMES`.
"""
if not self._graph_initialized:
self.initialize_graph()
predict_times = tf.cast(
ops.convert_to_tensor(features[PredictionFeatures.TIMES]),
tf.dtypes.int32)
exogenous_regressors = self._process_exogenous_features(
times=predict_times,
features={
key: value for key, value in features.items() if key not in [
TrainEvalFeatures.TIMES, TrainEvalFeatures.VALUES,
PredictionFeatures.STATE_TUPLE
]
})
with tf.control_dependencies([
tf.compat.v1.debugging.assert_equal(
tf.compat.v1.shape(predict_times)[1],
tf.compat.v1.shape(exogenous_regressors)[1])
]):
exogenous_regressors = tf.identity(exogenous_regressors)
batch_size = tf.compat.v1.shape(predict_times)[0]
num_predict_values = tf.compat.v1.shape(predict_times)[1]
prediction_iterations = (
(num_predict_values + self.output_window_size - 1) //
self.output_window_size)
# Pad predict_times and exogenous regressors so as to have exact multiple of
# self.output_window_size values per example.
padding_size = (
prediction_iterations * self.output_window_size - num_predict_values)
predict_times = tf.compat.v1.pad(predict_times, [[0, 0], [0, padding_size]])
exogenous_regressors = tf.compat.v1.pad(exogenous_regressors,
[[0, 0], [0, padding_size], [0, 0]])
state = features[PredictionFeatures.STATE_TUPLE]
(state_times, state_values, state_exogenous_regressors) = state
state_times = tf.cast(ops.convert_to_tensor(state_times), tf.dtypes.int32)
state_values = ops.convert_to_tensor(state_values, dtype=self.dtype)
state_exogenous_regressors = ops.convert_to_tensor(
state_exogenous_regressors, dtype=self.dtype)
initial_input_times = predict_times[:, :self.output_window_size]
initial_input_exogenous_regressors = (
exogenous_regressors[:, :self.output_window_size, :])
if self.input_window_size > 0:
initial_input_times = tf.concat(
[state_times[:, -self.input_window_size:], initial_input_times], 1)
values_size = tf.compat.v1.shape(state_values)[1]
times_size = tf.compat.v1.shape(state_times)[1]
with tf.control_dependencies([
tf.compat.v1.debugging.assert_greater_equal(values_size,
self.input_window_size),
tf.compat.v1.debugging.assert_equal(values_size, times_size)
]):
initial_input_values = state_values[:, -self.input_window_size:, :]
initial_input_exogenous_regressors = tf.concat([
state_exogenous_regressors[:, -self.input_window_size:, :],
initial_input_exogenous_regressors[:, :self.output_window_size, :]
],
axis=1)
else:
initial_input_values = 0
# Iterate over the predict_times, predicting self.output_window_size values
# in each iteration.
def _while_condition(iteration_number, *unused_args):
return tf.math.less(iteration_number, prediction_iterations)
def _while_body(iteration_number, input_times, input_values,
input_exogenous_regressors, mean_ta, covariance_ta):
"""Predict self.output_window_size values."""
prediction_ops = self.prediction_ops(input_times, input_values,
input_exogenous_regressors)
predicted_mean = prediction_ops["mean"]
predicted_covariance = prediction_ops["covariance"]
offset = self.output_window_size * tf.math.minimum(
iteration_number + 1, prediction_iterations - 1)
if self.input_window_size > 0:
if self.output_window_size < self.input_window_size:
new_input_values = tf.concat(
[input_values[:, self.output_window_size:, :], predicted_mean], 1)
new_input_exogenous_regressors = tf.concat([
input_exogenous_regressors[:, -self.input_window_size:, :],
exogenous_regressors[
:, offset:offset + self.output_window_size, :]
], axis=1)
new_input_times = tf.concat([
input_times[:, -self.input_window_size:],
predict_times[:, offset:offset + self.output_window_size]
], 1)
else:
new_input_values = predicted_mean[:, -self.input_window_size:, :]
new_input_exogenous_regressors = exogenous_regressors[
:,
offset - self.input_window_size:offset + self.output_window_size,
:]
new_input_times = predict_times[
:,
offset - self.input_window_size:offset + self.output_window_size]
else:
new_input_values = input_values
new_input_exogenous_regressors = exogenous_regressors[
:, offset:offset + self.output_window_size, :]
new_input_times = predict_times[:,
offset:offset + self.output_window_size]
new_input_times.set_shape(initial_input_times.get_shape())
new_input_exogenous_regressors.set_shape(
initial_input_exogenous_regressors.get_shape())
new_mean_ta = mean_ta.write(iteration_number, predicted_mean)
if isinstance(covariance_ta, tf.TensorArray):
new_covariance_ta = covariance_ta.write(iteration_number,
predicted_covariance)
else:
new_covariance_ta = covariance_ta
return (iteration_number + 1, new_input_times, new_input_values,
new_input_exogenous_regressors, new_mean_ta, new_covariance_ta)
# Note that control_flow_ops.while_loop doesn't seem happy with None. Hence
# using 0 for cases where we don't want to predict covariance.
covariance_ta_init = (
tf.TensorArray(dtype=self.dtype, size=prediction_iterations)
if self.loss != ARModel.SQUARED_LOSS else 0.)
mean_ta_init = tf.TensorArray(dtype=self.dtype, size=prediction_iterations)
_, _, _, _, mean_ta, covariance_ta = tf.compat.v1.while_loop(
_while_condition, _while_body, [
0, initial_input_times, initial_input_values,
initial_input_exogenous_regressors, mean_ta_init, covariance_ta_init
])
def _parse_ta(values_ta):
"""Helper function to parse the returned TensorArrays."""
if not isinstance(values_ta, tf.TensorArray):
return None
predictions_length = prediction_iterations * self.output_window_size
# Shape [prediction_iterations, batch_size, self.output_window_size,
# self.num_features]
values_packed = values_ta.stack()
# Transpose to move batch dimension outside.
output_values = tf.reshape(
tf.compat.v1.transpose(values_packed, [1, 0, 2, 3]),
tf.stack([batch_size, predictions_length, -1]))
# Clip to desired size
return output_values[:, :num_predict_values, :]
predicted_mean = _parse_ta(mean_ta)
predicted_covariance = _parse_ta(covariance_ta)
if predicted_covariance is None:
predicted_covariance = tf.compat.v1.ones_like(predicted_mean)
# Transform and scale the mean and covariance appropriately.
predicted_mean = self._scale_back_data(predicted_mean)
predicted_covariance = self._scale_back_variance(predicted_covariance)
return {"mean": predicted_mean, "covariance": predicted_covariance}