My Bayesian CNN model will not learn if I have more than 1 dense variational layer and I can't figure out why. I am working on a bayesian cnn from a tensorflow 2 tutorial: https://goodboychan.github.io/pytho...01-Bayesian-Convolutional-Neural-Network.html My current model is this: def build_bayesian_model_mnist(x_train, number_classes): divergence_fn = lambda q, p, _: tfd.kl_divergence(q, p) / x_train.shape[0] bayesian_model = Sequential([ InputLayer(input_shape=(x_train.shape[1], x_train.shape[2], x_train.shape[3])), get_convolutional_reparameterization_layer(16, (3, 3), divergence_fn), MaxPooling2D(pool_size=(3, 3)), get_convolutional_reparameterization_layer(32, (3, 3), divergence_fn), GlobalMaxPooling2D(), # Flatten(), # For whatever reason adding this casues the model to stop learning # get_dense_variational_layer(get_prior, get_posterior, kl_weight=1/x_train.shape[0], num_classes=32), get_dense_variational_layer(get_prior, get_posterior, kl_weight=1 / x_train.shape[0], num_classes=number_classes), tfpl.OneHotCategorical(number_classes, convert_to_tensor_fn=tfd.Distribution.mode) ]) bayesian_model.compile(loss=nll, optimizer=RMSprop(), metrics=['accuracy'], experimental_run_tf_function=False) return bayesian_model If I uncomment the dense variation layer with 32 units the model will not get past 10% accuracy. I've also made sure that the training data is the correct format. It will train on mnist with 1 dense layer and achieve 90+ accuracy, and on cifar with 1 dense layer the model will learn. However, I need to add extra dense layers to further improve performance on other datasets. If there is anything I should check out/modify, let me know! Thank you! As well, if it helps, here is the extra code I have from building the model: def nll(y_true, y_pred): """ This function should return the negative log-likelihood of each sample in y_true given the predicted distribution y_pred. If y_true is of shape [B, E] and y_pred has batch shape and event_shape [E], the output should be a Tensor of shape . """ return -y_pred.log_prob(y_true) def get_convolutional_reparameterization_layer(filters, kernel_size, divergence_fn): layer = tfpl.Convolution2DReparameterization( filters=filters, kernel_size=kernel_size, activation= tf.nn.softplus, # 'relu', padding='SAME', # Preserves spatial resolution kernel_prior_fn=tfpl.default_multivariate_normal_fn, kernel_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False), kernel_divergence_fn=divergence_fn, bias_prior_fn=tfpl.default_multivariate_normal_fn, bias_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False), bias_divergence_fn=divergence_fn ) return layer def spike_and_slab(event_shape, dtype): distribution = tfd.Mixture( cat=tfd.Categorical(probs=[0.5, 0.5]), components=[ tfd.Independent(tfd.Normal( loc=tf.zeros(event_shape, dtype=dtype), scale=1.0 * tf.ones(event_shape, dtype=dtype)), reinterpreted_batch_ndims=1), tfd.Independent(tfd.Normal( loc=tf.zeros(event_shape, dtype=dtype), scale=10.0 * tf.ones(event_shape, dtype=dtype)), reinterpreted_batch_ndims=1)], name='spike_and_slab') return distribution def get_prior(kernel_size, bias_size, dtype=None): """ This function should create the prior distribution, consisting of the "spike and slab" distribution that is described above. The distribution should be created using the kernel_size, bias_size and dtype function arguments above. The function should then return a callable, that returns the prior distribution. """ n = kernel_size + bias_size prior_model = Sequential([tfpl.DistributionLambda(lambda t: spike_and_slab(n, dtype))]) return prior_model def get_posterior(kernel_size, bias_size, dtype=None): """ This function should create the posterior distribution as specified above. The distribution should be created using the kernel_size, bias_size and dtype function arguments above. The function should then return a callable, that returns the posterior distribution. """ n = kernel_size + bias_size return Sequential([ tfpl.VariableLayer(tfpl.IndependentNormal.params_size(n), dtype=dtype), tfpl.IndependentNormal(n) ]) def get_dense_variational_layer(prior_fn, posterior_fn, kl_weight, num_classes): """ This function should create an instance of a DenseVariational layer according to the above specification. The function takes the prior_fn, posterior_fn and kl_weight as arguments, which should be used to define the layer. Your function should then return the layer instance. """ return tfpl.DenseVariational( units=num_classes, make_posterior_fn=posterior_fn, make_prior_fn=prior_fn, kl_weight=kl_weight ) Continue reading...
