Open esperber opened 1 year ago
Sorry for the sluggish response. Can you try this?
parameter_samples, _ = tf.function(
func=lambda ots: tfp.sts.fit_with_hmc(model, ots),
jit_compile=True,
autograph=False)(observed_time_series)
There will be some compilation overhead on the first run but should be overall much faster.
@csuter How can I run fit_with_hmc with jit_compile activate?
Following you example
import numpy as np
import tensorflow_probability as tfp
import tensorflow as tf
time_series_with_nans = [-1.0, 1.0, np.nan, 2.4, np.nan, 5]
observed_time_series = tfp.sts.MaskedTimeSeries(
time_series=time_series_with_nans, is_missing=tf.math.is_nan(time_series_with_nans)
)
# Build model using observed time series to set heuristic priors.
linear_trend_model = tfp.sts.LocalLinearTrend(observed_time_series=observed_time_series)
model = tfp.sts.Sum([linear_trend_model], observed_time_series=observed_time_series)
# Fit model to data
parameter_samples, _ = tf.function(
func=lambda ots: tfp.sts.fit_with_hmc(model, ots), jit_compile=True, autograph=False
)(observed_time_series)
Gives me the following error:
parameter_samples, _ = tf.function(
func=lambda ots: tfp.sts.fit_with_hmc(model, ots),
jit_compile=True,
autograph=False)(observed_time_series)
test_jit_hmc.py:262:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
.env/lib/python3.11/site-packages/tensorflow/python/util/traceback_utils.py:153: in error_handler
raise e.with_traceback(filtered_tb) from None
incrementality/prefect/flows/model_training.py:263: in <lambda>
func=lambda ots: tfp.sts.fit_with_hmc(model, ots),
.env/lib/python3.11/site-packages/tensorflow_probability/python/sts/fitting.py:466: in fit_with_hmc
variational_posterior = build_factored_surrogate_posterior(
.env/lib/python3.11/site-packages/tensorflow_probability/python/sts/fitting.py:173: in build_factored_surrogate_posterior
return experimental_vi.build_factored_surrogate_posterior(
.env/lib/python3.11/site-packages/tensorflow_probability/python/internal/trainable_state_util.py:337: in build_stateful_trainable
tf.nest.map_structure(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
t = <tf.Tensor 'fit_with_hmc/build_factored_surrogate_posterior/build_factored_surrogate_posterior/Normal_trainable_variables/normal/stateless_random_normal:0' shape=() dtype=float32>, n = 'loc'
> lambda t, n=name: t if t is None else tf.Variable(t, name=n),
value, expand_composites=True))
E ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information.
.env/lib/python3.11/site-packages/tensorflow_probability/python/internal/trainable_state_util.py:338: ValueError
I'm using
python 3.11
pip list | grep tensorflow
tensorflow 2.17.0
tensorflow-probability 0.24.0
Hi,
I've been building some structural time-series models with Tensorflow Probability over the past week.
I've begun to look into the impute_missing_variable method to smooth in my missing time-series values, but the sts.fit_with_hmc seems incredibly slow. Just running the example on https://www.tensorflow.org/probability/api_docs/python/tfp/sts/impute_missing_values take about 25 seconds.
Is this expected behavior? I'm running this example on CPU and didn't see any performance improvements when I ran the same block of code on GPU.
This is using Tensorflow-Probability==0.19 and Tensorflow==2.11.0
Thanks!