def _estimator_fn(model_dir, exogenous_feature_columns):
return estimators.TimeSeriesRegressor(model=ar_model.ARModel(
periodicities=10,
input_window_size=10,
output_window_size=6,
num_features=1,
exogenous_feature_columns=exogenous_feature_columns,
prediction_model_factory=functools.partial(
ar_model.LSTMPredictionModel, num_units=10)),
config=_SeedRunConfig(),
model_dir=model_dir)
def test_one_shot_prediction_head_export(self):
def _new_temp_dir():
return os.path.join(tf.compat.v1.test.get_temp_dir(), str(ops.uid()))
model_dir = _new_temp_dir()
categorical_column = tf.feature_column.categorical_column_with_hash_bucket(
key="categorical_exogenous_feature", hash_bucket_size=16)
exogenous_feature_columns = [
tf.feature_column.numeric_column(
"2d_exogenous_feature", shape=(2,)),
tf.feature_column.embedding_column(
categorical_column=categorical_column, dimension=10)]
estimator = ts_estimators.TimeSeriesRegressor(
model=ar_model.ARModel(
periodicities=10,
input_window_size=10,
output_window_size=6,
num_features=5,
exogenous_feature_columns=exogenous_feature_columns,
prediction_model_factory=functools.partial(
ar_model.LSTMPredictionModel, num_units=10)),
head_type=ts_head_lib.OneShotPredictionHead,
model_dir=model_dir)
def train_input_fn():
num_range = tf.range(16, dtype=tf.dtypes.int64)
features = {
feature_keys.TrainEvalFeatures.TIMES:
tf.compat.v1.expand_dims(num_range, axis=0),
feature_keys.TrainEvalFeatures.VALUES:
tf.compat.v1.expand_dims(
tf.tile(num_range[:, None], [1, 5]), axis=0),
"2d_exogenous_feature":
tf.ones([1, 16, 2]),
"categorical_exogenous_feature":
tf.compat.v1.expand_dims(
tf.tile(["strkey"], [16])[:, None], axis=0)
}
return features
estimator.train(input_fn=train_input_fn, steps=5)
result = estimator.evaluate(input_fn=train_input_fn, steps=1)
self.assertIn("average_loss", result)
self.assertNotIn(feature_keys.State.STATE_TUPLE, result)
input_receiver_fn = estimator.build_raw_serving_input_receiver_fn()
export_location = estimator.export_saved_model(_new_temp_dir(),
input_receiver_fn)
graph = tf.Graph()
with graph.as_default():
with tf.compat.v1.Session() as session:
signatures = tf.compat.v1.saved_model.load(
session, [tf.saved_model.SERVING], export_location)
self.assertEqual([feature_keys.SavedModelLabels.PREDICT],
list(signatures.signature_def.keys()))
predict_signature = signatures.signature_def[
feature_keys.SavedModelLabels.PREDICT]
six.assertCountEqual(
self,
[feature_keys.FilteringFeatures.TIMES,
feature_keys.FilteringFeatures.VALUES,
"2d_exogenous_feature",
"categorical_exogenous_feature"],
predict_signature.inputs.keys())
features = {
feature_keys.TrainEvalFeatures.TIMES: numpy.tile(
numpy.arange(35, dtype=numpy.int64)[None, :], [2, 1]),
feature_keys.TrainEvalFeatures.VALUES: numpy.tile(numpy.arange(
20, dtype=numpy.float32)[None, :, None], [2, 1, 5]),
"2d_exogenous_feature": numpy.ones([2, 35, 2]),
"categorical_exogenous_feature": numpy.tile(numpy.array(
["strkey"] * 35)[None, :, None], [2, 1, 1])
}
feeds = {
graph.as_graph_element(input_value.name): features[input_key]
for input_key, input_value in predict_signature.inputs.items()}
fetches = {output_key: graph.as_graph_element(output_value.name)
for output_key, output_value
in predict_signature.outputs.items()}
output = session.run(fetches, feed_dict=feeds)
self.assertEqual((2, 15, 5), output["mean"].shape)
# Build a parsing input function, then make a tf.Example for it to parse.
export_location = estimator.export_saved_model(
_new_temp_dir(),
estimator.build_one_shot_parsing_serving_input_receiver_fn(
filtering_length=20, prediction_length=15))
graph = tf.Graph()
with graph.as_default():
with tf.compat.v1.Session() as session:
example = example_pb2.Example()
times = example.features.feature[feature_keys.TrainEvalFeatures.TIMES]
values = example.features.feature[feature_keys.TrainEvalFeatures.VALUES]
times.int64_list.value.extend(range(35))
for i in range(20):
values.float_list.value.extend(
[float(i) * 2. + feature_number
for feature_number in range(5)])
real_feature = example.features.feature["2d_exogenous_feature"]
categortical_feature = example.features.feature[
"categorical_exogenous_feature"]
for i in range(35):
real_feature.float_list.value.extend([1, 1])
categortical_feature.bytes_list.value.append(b"strkey")
# Serialize the tf.Example for feeding to the Session
examples = [example.SerializeToString()] * 2
signatures = tf.compat.v1.saved_model.load(
session, [tf.saved_model.SERVING], export_location)
predict_signature = signatures.signature_def[
feature_keys.SavedModelLabels.PREDICT]
((_, input_value),) = predict_signature.inputs.items()
feeds = {graph.as_graph_element(input_value.name): examples}
fetches = {output_key: graph.as_graph_element(output_value.name)
for output_key, output_value
in predict_signature.outputs.items()}
output = session.run(fetches, feed_dict=feeds)
self.assertEqual((2, 15, 5), output["mean"].shape)
def _train_on_generated_data(
generate_fn, generative_model, train_iterations, seed,
learning_rate=0.1, ignore_params_fn=lambda _: (),
derived_param_test_fn=lambda _: (),
train_input_fn_type=input_pipeline.WholeDatasetInputFn,
train_state_manager=state_management.PassthroughStateManager()):
"""The training portion of parameter recovery tests."""
random_seed.set_random_seed(seed)
generate_graph = ops.Graph()
with generate_graph.as_default():
with session.Session(graph=generate_graph):
generative_model.initialize_graph()
time_series_reader, true_parameters = generate_fn(generative_model)
true_parameters = {
tensor.name: value for tensor, value in true_parameters.items()}
eval_input_fn = input_pipeline.WholeDatasetInputFn(time_series_reader)
eval_state_manager = state_management.PassthroughStateManager()
true_parameter_eval_graph = ops.Graph()
with true_parameter_eval_graph.as_default():
generative_model.initialize_graph()
ignore_params = ignore_params_fn(generative_model)
feature_dict, _ = eval_input_fn()
eval_state_manager.initialize_graph(generative_model)
feature_dict[TrainEvalFeatures.VALUES] = math_ops.cast(
feature_dict[TrainEvalFeatures.VALUES], generative_model.dtype)
model_outputs = eval_state_manager.define_loss(
model=generative_model,
features=feature_dict,
mode=estimator_lib.ModeKeys.EVAL)
with session.Session(graph=true_parameter_eval_graph) as sess:
variables.global_variables_initializer().run()
coordinator = coordinator_lib.Coordinator()
queue_runner_impl.start_queue_runners(sess, coord=coordinator)
true_param_loss = model_outputs.loss.eval(feed_dict=true_parameters)
true_transformed_params = {
param: param.eval(feed_dict=true_parameters)
for param in derived_param_test_fn(generative_model)}
coordinator.request_stop()
coordinator.join()
saving_hook = _SavingTensorHook(
tensors=true_parameters.keys(),
every_n_iter=train_iterations - 1)
class _RunConfig(estimator_lib.RunConfig):
@property
def tf_random_seed(self):
return seed
estimator = estimators.TimeSeriesRegressor(
model=generative_model,
config=_RunConfig(),
state_manager=train_state_manager,
optimizer=adam.AdamOptimizer(learning_rate))
train_input_fn = train_input_fn_type(time_series_reader=time_series_reader)
trained_loss = (estimator.train(
input_fn=train_input_fn,
max_steps=train_iterations,
hooks=[saving_hook]).evaluate(
input_fn=eval_input_fn, steps=1))["loss"]
logging.info("Final trained loss: %f", trained_loss)
logging.info("True parameter loss: %f", true_param_loss)
return (ignore_params, true_parameters, true_transformed_params,
trained_loss, true_param_loss, saving_hook,
true_parameter_eval_graph)