Hello,
I am trying to integrate a neural network (using Keras) with a probabilistic model (using JointDistributionCoroutineAutoBatched). While the probabilistic model part works fine, I struggle with integrating the neural network. Appreciate any help on this!
Something similar was done here a few years ago using Edward https://willwolf.io/2017/06/15/random-effects-neural-networks/
Here is the working probabilistic (mixed-effects) model, which is pretty much the same as in this tutorial, using the same data (radon) https://www.tensorflow.org/probability/examples/Linear_Mixed_Effects_Model_Variational_Inference
def make_joint_distribution_coroutine(floor, county, n_counties): def model(): # Hyperpriors: # mu_alpha ~ Normal(0,1) mu_alpha = yield tfd.Normal(loc=0., scale=.1, name = 'mu_alpha') # sigma_alpha ~ HalfNormal(1) sigma_alpha = yield tfd.HalfNormal(scale=.1, name = 'sigma_alpha') # Priors: # alpha ~ Normal(mu_alpha, sigma_alpha) alpha = yield tfd.Normal(loc=mu_alpha*tf.ones(n_counties), scale=sigma_alpha, name='alpha') # beta ~ Normal(0,1) beta = yield tfd.Normal(loc=0., scale=.1, name = 'beta') # sigma ~ HalfNormal(1) sigma = yield tfd.HalfNormal(scale=.1, name = 'sigma') # Likelihood fixed_effect = floor*beta random_effect = tf.gather(alpha, county, axis=-1) mu = random_effect + fixed_effect yield tfd.Normal(loc=mu, scale=sigma, name = 'likelihood') return tfd.JointDistributionCoroutineAutoBatched(model) joint = make_joint_distribution_coroutine(floor, county, n_counties) Now the goal is to replace the ‘fixed_effect’ part with a neural network. I tried the following, yielding different error messages. Since everything revolves around distributions, I thought it would make sense to have the Neural Network yield a distribution as well. Here ‘floor*beta’ is replaced by ‘neural_network(floor, n_observations)’
@tf.function def neural_network(floor, n_observations): inputs = tf.keras.layers.Input(shape = (n_observations,), name = "inputs") outputs = tf.keras.layers.Dense(1, activation = 'linear', kernel_regularizer = tf.keras.regularizers.l2(.001))(inputs) return tfp.layers.DistributionLambda(make_distribution_fn= lambda outputs: tfd.Normal(loc=outputs, scale=1)) def make_joint_distribution_coroutine(floor, county, n_counties): def model(): # Hyperpriors: # mu_alpha ~ Normal(0,1) mu_alpha = yield tfd.Normal(loc=0., scale=.1, name = 'mu_alpha') # sigma_alpha ~ HalfNormal(1) sigma_alpha = yield tfd.HalfNormal(scale=.1, name = 'sigma_alpha') # Priors: # alpha ~ Normal(mu_alpha, sigma_alpha) alpha = yield tfd.Normal(loc=mu_alpha*tf.ones(n_counties), scale=sigma_alpha, name='alpha') # sigma ~ HalfNormal(1) sigma = yield tfd.HalfNormal(scale=.1, name = 'sigma') # Likelihood fixed_effect = neural_network(floor, n_observations) random_effect = tf.gather(alpha, county, axis=-1) mu = random_effect + fixed_effect yield tfd.Normal(loc=mu, scale=sigma, name = 'likelihood') return tfd.JointDistributionCoroutineAutoBatched(model) joint = make_joint_distribution_coroutine(floor, county, n_counties) But this yields the following error:
Error in py_call_impl(callable, dots$args, dots$keywords) : TypeError: To be compatible with tf.function, Python functions must return zero or more Tensors or ExtensionTypes or None values; in compilation of <function neural_network at 0x000001E849878310>, found return value of type DistributionLambda, which is not a Tensor or ExtensionType. More in line with the Edward example posted above, I also tried the following:
@tf.function def neural_network(floor, n_observations): inputs = tf.keras.layers.Input(shape = (n_observations,), name = "inputs") outputs = tf.keras.layers.Dense(1, activation = 'linear', kernel_regularizer = tf.keras.regularizers.l2(.001))(inputs) return tf.keras.backend.squeeze(outputs, axis=1) which yields:
Error in py_call_impl(callable, dots$args, dots$keywords) : TypeError: To be compatible with tf.function, Python functions must return zero or more Tensors or ExtensionTypes or None values; in compilation of <function neural_network at 0x000001E849B6B9D0>, found return value of type KerasTensor, which is not a Tensor or ExtensionType.