Speed of tfp.sts.fit_with_hmc

[esperber]

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.

import numpy as np
import  tensorflow_probability as tfp
import tensorflow as tf

time_series_with_nans = [-1., 1., 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, _ = tfp.sts.fit_with_hmc(model, observed_time_series)

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!

[csuter]

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)

[csuter]

There will be some compilation overhead on the first run but should be overall much faster.

[williamjamir]

@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