model based on TensorFlow Probability with keras.fit

I am use tfp *-Flipout layers to construct a Bayesian neural network (BNN) and combine it with keras.fit to train. I am using a very similar way to define a BNN structure as a CNN but the keras.fit() function returns an issue about None gradient as

ValueError: Variable <tf.Variable 'conv2d_flipout/kernel_posterior_loc:0' shape=(3, 3, 1, 32) dtype=float32> has `None` for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: K.argmax, K.round, K.eval.

I am using the following versions of tfp and tf:

tfp.__version__ == '0.7.0'

tf.__version__ == '1.14.0'

Below is a minimal working example on the MNIST dataset. Feel free to comment the working CNN part to see the BNN error above (either bcnn_model_1 or bcnn_model_2 throws the above None gradient error when calling their fit functions):

import os
os.environ['KERAS_BACKEND'] = 'tensorflow'  # set up tensorflow backend for keras

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.layers import DenseVariational, DenseReparameterization, DenseFlipout, Convolution2DFlipout, Convolution2DReparameterization
from tensorflow_probability.python.layers import DistributionLambda
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.layers import Input, Dense, Conv2D, Flatten, BatchNormalization, Activation, LeakyReLU
from tensorflow.keras.utils import plot_model
from tensorflow.keras.optimizers import *

tf.enable_eager_execution()
tfd = tfp.distributions

import numpy as np
import matplotlib.pyplot as plt

def neg_log_likelihood(y_true, y_pred):
    return -y_pred.log_prob(y_true)

def get_neg_log_likelihood_fn(bayesian=False):
    """
    Get the negative log-likelihood function
    # Arguments
        bayesian(bool): Bayesian neural network (True) or point-estimate neural network (False)

    # Returns
        a negative log-likelihood function
    """
    if bayesian:
        def neg_log_likelihood_bayesian(y_true, y_pred):
            labels_distribution = tfp.distributions.Categorical(logits=y_pred)
            log_likelihood = labels_distribution.log_prob(tf.argmax(input=y_true, axis=1))
            loss = -tf.reduce_mean(input_tensor=log_likelihood)
            return loss
        return neg_log_likelihood_bayesian
    else:
        def neg_log_likelihood(y_true, y_pred):
            y_pred_softmax = keras.layers.Activation('softmax')(y_pred)  # logits to softmax
            loss = keras.losses.categorical_crossentropy(y_true, y_pred_softmax)
            return loss
        return neg_log_likelihood


n_class = 10

(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()
X_train = np.expand_dims(X_train, -1)
n_train = X_train.shape[0]
X_test = np.expand_dims(X_test, -1)
n_test = X_test.shape[0]

# Normalize data
X_train = X_train.astype('float32') / 255
X_test = X_test.astype('float32') / 255

print("X_train.shape =", X_train.shape)
print("y_train.shape =", y_train.shape)
print("X_test.shape =", X_test.shape)
print("y_test.shape =", y_test.shape)

plt.imshow(X_train[0, :, :, 0], cmap='gist_gray')

lr = 1e-3

def build_cnn_model(input_shape):
    model_in = Input(shape=input_shape)
    x = Conv2D(32, kernel_size=3, padding="same", strides=2)(model_in)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Conv2D(64, kernel_size=3, padding="same", strides=2)(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Flatten()(x)
    x = Dense(512, activation='relu')(x)
    model_out = Dense(10, activation='softmax')(x)  # softmax
    model = Model(model_in, model_out)
    return model

def build_bayesian_cnn_model_1(input_shape):
    model_in = Input(shape=input_shape)
    x = Convolution2DFlipout(32, kernel_size=3, padding="same", strides=2)(model_in)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Convolution2DFlipout(64, kernel_size=3, padding="same", strides=2)(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Flatten()(x)
    x = DenseFlipout(512, activation='relu')(x)
    model_out = DenseFlipout(10, activation=None)(x)  # logits
    model = Model(model_in, model_out)
    return model

def build_bayesian_cnn_model_2(input_shape):
    model_in = Input(shape=input_shape)
    x = Convolution2DFlipout(32, kernel_size=3, padding="same", strides=2)(model_in)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Convolution2DFlipout(64, kernel_size=3, padding="same", strides=2)(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Flatten()(x)
    x = DenseFlipout(512, activation='relu')(x)
    x = DenseFlipout(10, activation=None)(x)  # logits
    model_out = DistributionLambda(lambda t: tfd.Categorical(logits=t))(x)  # distribution
    model = Model(model_in, model_out)
    return model

cnn_model = build_cnn_model(X_train.shape[1:])
cnn_model.compile(loss='sparse_categorical_crossentropy', optimizer=Adam(lr), metrics=['accuracy'])
print('CNN Model:')
cnn_model.summary()

bcnn_model_1 = build_bayesian_cnn_model_1(X_train.shape[1:])
bcnn_model_1.compile(loss=get_neg_log_likelihood_fn(bayesian=True), optimizer=Adam(lr), metrics=['accuracy'])
print("BCNN Model 1:")
bcnn_model_1.summary()

bcnn_model_2 = build_bayesian_cnn_model_2(X_train.shape[1:])
bcnn_model_2.compile(loss=neg_log_likelihood, optimizer=Adam(lr), metrics=['accuracy'])
print("BCNN Model 2:")
bcnn_model_2.summary()

batch_size = 128
n_epochs = 30
hist_cnn = cnn_model.fit(X_train, y_train, batch_size=batch_size, epochs=n_epochs, verbose=1)
hist_bcnn_1 = bcnn_model_1.fit(X_train, y_train, batch_size=batch_size, epochs=n_epochs, verbose=1)
hist_bcnn_2 = bcnn_model_2.fit(X_train, y_train, batch_size=batch_size, epochs=n_epochs, verbose=1)

Any idea why is keras.fit() not able to work for such BNN models?

I also implemented a ResNet with tfp layers, as shown here: https://github.com/zhulingchen/tfp-resnet/blob/master/tfp_resnet.py. That really did work. So it starts to confuse me.

Start by running the minimal MNIST example with TensorFlow 1.14.0 and TensorFlow Probability 0.7.0, focusing on build_bayesian_cnn_model_1, build_bayesian_cnn_model_2, and their fit calls. Compare the Bayesian models with build_cnn_model and trace the None gradient; done means identifying the compatibility or loss-path cause and confirming a working training path.