在DenseVariational层后添加自定义层遇数据类型错误，求解决

背景代码 #

贝叶斯神经网络代码 #

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) 

自定义层代码 #

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):
    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)

错误信息 #

TypeError: Value passed to parameter 'x' has DataType int32 not in list of allowed 
values: bfloat16, float16, float32, float64, complex64, complex128

错误原因及解决方法 #

核心错误点 #

1. 输入类型不匹配：DistributionLambda(normal_sp)输出的是tfd.Normal分布对象，不是张量，直接在自定义层中把它当作张量使用（inputs[0]），会导致类型识别混乱。
2. 外部变量未传入层：自定义层中的x是外部全局变量，未作为层的输入传入；如果x是int32类型，和网络输出的float类型参数运算时会触发类型冲突。
3. 公式语法错误：
   • 1/24(...)缺少乘法运算符，应为(1/24)*(...)
   • x^2是Python按位异或操作，不是数学平方，需改为tf.square(x)或x**2

修正后的代码 #

1. 调整模型结构（改为多输入逻辑，传入x和网络参数）

def create_probabilistic_bnn_model(train_size):
    # 输入特征x
    x_input = tf.keras.Input(shape=(1,))
    # BNN部分输出参数分布
    x = tfp.layers.DenseVariational(50, posterior_mean_field, prior, kl_weight=1/train_size, activation='relu', kl_use_exact=True)(x_input)
    x = tfp.layers.DenseVariational(50, posterior_mean_field, prior, kl_weight=1/train_size, activation='relu', kl_use_exact=True)(x)
    x = tfp.layers.DenseVariational(50, posterior_mean_field, prior, kl_weight=1/train_size, activation='relu', kl_use_exact=True)(x)
    params = tfp.layers.DenseVariational(2, posterior_mean_field, prior, kl_weight=1/train_size, kl_use_exact=True)(x)
    distr = tfp.layers.DistributionLambda(normal_sp)(params)
    # 提取分布的均值作为待学习的参数w
    w = distr.mean()
    # 自定义层接收w和原始输入x
    output = MyDenseLayer()([w, x_input])
    # 构建多输入模型
    model = tf.keras.Model(inputs=x_input, outputs=output)
    return model

2. 修正自定义层

class MyDenseLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(MyDenseLayer, self).__init__()

    def call(self, inputs):
        # inputs为列表：[待学习参数w, 原始输入x]
        w = inputs[0]
        x = inputs[1]
        # 确保x转为float32类型，避免int冲突
        x = tf.cast(x, tf.float32)
        # 公式计算（修正语法错误）
        q0 = (3/2) * w
        j0 = 1.0  # 用float类型避免int干扰
        term1 = 0.5 * (1 - q0) * x
        term2 = -(1/24) * (7 - 10*q0 - 9*tf.square(q0)) * tf.square(x)
        sol = term1 + term2 + j0
        return sol

3. 重新编译模型

train_size = x.shape[0]
optimizer = tf.optimizers.Adam(learning_rate=0.0002)

keras_BNN = create_probabilistic_bnn_model(train_size=train_size)
# 此时输出是张量，损失函数需对应调整（不再是NLL，因为不再输出分布）
keras_BNN.compile(optimizer=optimizer, loss='mse')

说明 #

修正后，模型会先通过BNN学习公式中的参数w，再结合原始输入x计算最终的y值，同时解决了数据类型冲突和语法错误问题。

内容的提问来源于stack exchange，提问作者Alucard