I have made a small bayesian neural network with few dense variational layers.
import numpy as np
from sklearn.model_selection import train_test_split
from tqdm.notebook import tqdm
import tensorflow_probability as tfp
import tensorflow as tf
from tensorflow.keras.layers import Input
from tensorflow.keras.optimizers import Adam
tfd= tfp.distributions
def prior(kernel_size, bias_size, dtype=None):
n = kernel_size bias_size
# Independent Normal Distribution
return lambda t: tfd.Independent(tfd.Normal(loc=tf.zeros(n, dtype=dtype),scale=1), reinterpreted_batch_ndims=1)
train_size = X_train.shape[0] batch_size = 256
def NLL(y, distr):
return -distr.log_prob(y)
def normal_sp(params):
return tfd.Normal(loc=params[:,:1],scale=1e-5 0.001*tf.keras.backend.exp(params[:,1:]))# both parameters are learnable
def random_gaussian_initializer(shape, dtype):
n = int(shape / 2)
loc_norm = tf.random_normal_initializer(mean=0., stddev=0.1)
loc = tf.Variable(initial_value=loc_norm(shape=(n,), dtype=dtype) )
scale_norm = tf.random_normal_initializer(mean=-3., stddev=0.1)
scale = tf.Variable(initial_value=scale_norm(shape=(n,), dtype=dtype))
return tf.concat([loc, scale], 0)
def posterior_mean_field(kernel_size, bias_size=0, dtype=None):
n = kernel_size bias_size
return tf.keras.Sequential([
tfp.layers.VariableLayer(2 * n, dtype=dtype,initializer=lambda shape, dtype: random_gaussian_initializer(shape, dtype), trainable=True),
tfp.layers.DistributionLambda(lambda t: tfd.Independent(tfd.Normal(loc=t[..., :n],scale=1e-5 0.02*tf.nn.softplus(0.04 t[..., n:])),reinterpreted_batch_ndims=1)),
])
def create_probabilistic_bnn_model(train_size):
model = tf.keras.Sequential([
tf.keras.Input(shape=(1,)),
tfp.layers.DenseVariational(50, posterior_mean_field, prior, kl_weight=1/train_size, activation='relu', kl_use_exact=True),
tfp.layers.DenseVariational(50, posterior_mean_field, prior, kl_weight=1/train_size, activation='relu', kl_use_exact=True),
tfp.layers.DenseVariational(50, posterior_mean_field, prior, kl_weight=1/train_size, activation='relu', kl_use_exact=True),
tfp.layers.DenseVariational(2, posterior_mean_field, prior, kl_weight=1/train_size, kl_use_exact=True),
tfp.layers.DistributionLambda(normal_sp)
])
return model
train_size = x.shape[0]
batch_size = 256
optimizer = tf.optimizers.Adam(learning_rate=0.0002)
keras_BNN = create_probabilistic_bnn_model(train_size=train_size)
keras_BNN.compile(optimizer=optimizer,loss=NLL)
I taught to add a custom layer at the end so that the bnn would learn not directly the target y but a parameter w of a simple equation $y= f(x,w)$.
def __init__(self, num_outputs):
super(MyDenseLayer, self).__init__()
self.num_outputs = num_outputs
def build(self, input_shape):
self.kernel = self.add_weight("kernel",
shape=[int(input_shape[-1]),
self.num_outputs])
def call(self, inputs):
q0= 3/2*inputs[0]
j0=1
sol= 0.5*(1-q0)*x -1/24(7-10*q0 -9*q0**2 )*x^2 j0
return sol
layer = MyDenseLayer(1)
But when I tried to add MyDenseLayer(1) after tfp.layers.DistributionLambda(normal_sp) I got the error message:
TypeError: Value passed to parameter 'x' has DataType int32 not in list of allowed
values: bfloat16, float16, float32, float64, complex64, complex128
What do I do wrong?
CodePudding user response:
You should stick to one data type when doing your calculations and it should be fine. Here is an example, but I cannot verify that the logic is correct (that is up to you):
import tensorflow as tf
import tensorflow_probability as tfp
tfk = tf.keras
tfkl = tf.keras.layers
tfd = tfp.distributions
tfpl = tfp.layers
class MyDenseLayer(tf.keras.layers.Layer):
def __init__(self, num_outputs):
super(MyDenseLayer, self).__init__()
self.num_outputs = num_outputs
def build(self, input_shape):
self.kernel = self.add_weight("kernel",
shape=[int(input_shape[-1]),
self.num_outputs])
def call(self, inputs):
x = 4
q0= 3/2*inputs[0]
j0 = 1
xor = tf.cast(x, dtype=tf.int32)^2
sol= 0.5*(1-q0)*x -1/24 * (7-10*q0 -9*q0**2 )* tf.cast(xor, dtype=tf.float32) j0
return sol
def normal_sp(params):
return tfd.Normal(loc=params[:,:1],scale=1e-5 0.001*tf.keras.backend.exp(params[:,1:]))# both parameters are learnable
layer = MyDenseLayer(1)
model = tfk.Sequential([
tfkl.Dense(2, input_dim=2),
tfpl.DistributionLambda(normal_sp),
layer
])
model(tf.random.normal((2, 2)))
Note that ^
that does not work with floats that is why I am explicitly casting to int
, assuming that x
is later some float tensor. If not, you do not have to cast.