替换Conv1D/Dense为Flipout层的1D贝叶斯CNN无法训练求解

我有一个用于波形分类的CNN模型，输入形状为(601,3)（601为时间步，3为通道数），可将波形分为噪声或信号。该模型训练正常，第12轮收敛，最终准确率超99%。但将其中的Conv1D和Dense层分别替换为Convolution1DFlipout和DenseFlipout层构建1D贝叶斯CNN后，模型无法收敛，训练过程中准确率始终徘徊在50%左右，loss下降缓慢且波动。

以下是原模型、贝叶斯模型代码及训练日志：

原模型代码 #

# imports
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

random_seed = 42

tf.random.set_seed(random_seed)

model = keras.Sequential([
    layers.Input(shape=(601, 3)),  # Input shape for 1D data
    layers.Conv1D(32, kernel_size=16, activation='relu'),
    layers.Conv1D(64, kernel_size=16, activation='relu'),
    layers.Conv1D(128, kernel_size=16, activation='relu'),
    layers.Flatten(),
    layers.Dense(80, activation='relu'),
    layers.Dense(80, activation='relu'),
    layers.Dense(2, activation='softmax')
])

optimizer = keras.optimizers.Adam(learning_rate=0.001)
model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])

num_epochs = 40
batch_size = 48

history = model.fit(X_train, y_train_encoded, epochs=num_epochs, batch_size=batch_size,
                    validation_data=(X_test, y_test_encoded), verbose=2)

# X_train shape: (num_train_samples,601,3)
# X_test shape: (num_test_samples,601,3)
# y_train_encoded shape: (num_train_samples,2)
# y_test_encoded shape: (num_test_samples,2)

贝叶斯CNN模型代码 #

# imports
import tensorflow_probability as tfp

tfd = tfp.distributions
tfpl = tfp.layers

random_seed = 42

tf.random.set_seed(random_seed)

num_training_samples = X_train.shape[0]
kl_divergence_fn = lambda q, p, _: tfd.kl_divergence(q, p) / num_training_samples

model = keras.Sequential([
    layers.Input(shape=(601, 3)),
    tfpl.Convolution1DFlipout(
        32, kernel_size=16, activation=tf.nn.relu, kernel_divergence_fn=kl_divergence_fn, bias_divergence_fn=kl_divergence_fn),
    tfpl.Convolution1DFlipout(
        64, kernel_size=16, activation=tf.nn.relu, kernel_divergence_fn=kl_divergence_fn, bias_divergence_fn=kl_divergence_fn),
    tfpl.Convolution1DFlipout(
        128, kernel_size=16, activation=tf.nn.relu, kernel_divergence_fn=kl_divergence_fn, bias_divergence_fn=kl_divergence_fn),
    layers.MaxPooling1D(pool_size=2),
    layers.Flatten(),
    tfpl.DenseFlipout(80, activation=tf.nn.relu, kernel_divergence_fn=kl_divergence_fn, bias_divergence_fn=kl_divergence_fn),
    tfpl.DenseFlipout(80, activation=tf.nn.relu, kernel_divergence_fn=kl_divergence_fn, bias_divergence_fn=kl_divergence_fn),
    tfpl.DenseFlipout(2, activation=tf.nn.softmax, kernel_divergence_fn=kl_divergence_fn, bias_divergence_fn=kl_divergence_fn)
])

optimizer = keras.optimizers.Adam(learning_rate=0.001)
model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])

# Train the model (same as before)
num_epochs = 40
batch_size = 48

history = model.fit(X_train, y_train_encoded, epochs=num_epochs, batch_size=batch_size,
                    validation_data=(X_test, y_test_encoded), verbose=2)

训练日志 #

Layer (type)                Output Shape              Param #   
=================================================================
 conv1d_flipout (Conv1DFlipo  (None, 586, 32)          3104      
 ut)                                                             
                                                                 
 conv1d_flipout_1 (Conv1DFli  (None, 571, 64)          65600     
 pout)                                                           
                                                                 
 conv1d_flipout_2 (Conv1DFli  (None, 556, 128)         262272    
 pout)                                                           
                                                                 
 flatten (Flatten)           (None, 71168)             0         
                                                                 
 dense_flipout (DenseFlipout  (None, 80)               11386960  
 )                                                               
                                                                 
 dense_flipout_1 (DenseFlipo  (None, 80)               12880     
 ut)                                                             
                                                                 
 dense_flipout_2 (DenseFlipo  (None, 2)                322       
 ut)                                                             
                                                                 
=================================================================
Total params: 11,731,138
Trainable params: 11,731,138
Non-trainable params: 0
_________________________________________________________________
Epoch 1/40
3572/3572 - 1937s - loss: 44.5086 - accuracy: 0.4996 - val_loss: 10.1754 - val_accuracy: 0.5016 - 1937s/epoch - 542ms/step
Epoch 2/40
3572/3572 - 1933s - loss: 3.7701 - accuracy: 0.4993 - val_loss: 1.8013 - val_accuracy: 0.4946 - 1933s/epoch - 541ms/step
...
Epoch 40/40
3572/3572 - 1945s - loss: 0.7480 - accuracy: 0.5011 - val_loss: 0.7336 - val_accuracy: 0.4969 - 1945s/epoch - 544ms/step

问题排查与解决 #

1. 结构不一致导致特征断裂 #

贝叶斯模型额外添加了MaxPooling1D层，这和原模型的特征提取路径完全不同。原模型经过3层Conv1D后直接Flatten，而贝叶斯模型缩小了特征维度，导致模型无法学习到原有的有效特征。解决方案：移除MaxPooling1D层，保持和原模型结构一致。

2. Loss计算未正确包含KL散度 #

贝叶斯模型的总损失应该是分类交叉熵 + KL散度正则项，但你直接使用categorical_crossentropy作为Loss，没有将KL散度纳入计算。虽然TFP的Flipout层会自动将KL损失添加到model.losses中，但需要手动合并到总Loss里。解决方案：自定义Loss函数，将交叉熵和KL损失相加。

3. 学习率过高 #

贝叶斯模型的参数数量是原模型的2倍（每个参数对应均值和标准差两个可训练变量），原学习率0.001对于贝叶斯模型来说过大，容易导致梯度爆炸或训练不稳定。解决方案：将学习率降低到0.0001或更小。

4. 输出层激活函数冲突 #

在Flipout层后使用softmax激活，会叠加随机性和非线性，导致训练难度增加。建议输出层用linear激活，在Loss计算中手动应用softmax，或者使用TFP专门的分类层。解决方案：修改输出层激活为linear，在自定义Loss中计算softmax。

修正后的贝叶斯模型代码 #

# imports
import tensorflow_probability as tfp

tfd = tfp.distributions
tfpl = tfp.layers

random_seed = 42

tf.random.set_seed(random_seed)

num_training_samples = X_train.shape[0]
kl_divergence_fn = lambda q, p, _: tfd.kl_divergence(q, p) / num_training_samples

model = keras.Sequential([
    layers.Input(shape=(601, 3)),
    tfpl.Convolution1DFlipout(
        32, kernel_size=16, activation='relu', kernel_divergence_fn=kl_divergence_fn),
    tfpl.Convolution1DFlipout(
        64, kernel_size=16, activation='relu', kernel_divergence_fn=kl_divergence_fn),
    tfpl.Convolution1DFlipout(
        128, kernel_size=16, activation='relu', kernel_divergence_fn=kl_divergence_fn),
    # 移除额外的MaxPooling1D
    layers.Flatten(),
    tfpl.DenseFlipout(80, activation='relu', kernel_divergence_fn=kl_divergence_fn),
    tfpl.DenseFlipout(80, activation='relu', kernel_divergence_fn=kl_divergence_fn),
    # 输出层使用linear激活
    tfpl.DenseFlipout(2, activation='linear', kernel_divergence_fn=kl_divergence_fn)
])

# 自定义总Loss：交叉熵 + KL散度
def total_loss(y_true, y_pred):
    cross_entropy = keras.losses.categorical_crossentropy(y_true, tf.nn.softmax(y_pred))
    kl_loss = sum(model.losses)
    return cross_entropy + kl_loss

# 降低学习率
optimizer = keras.optimizers.Adam(learning_rate=0.0001)
model.compile(loss=total_loss, optimizer=optimizer, metrics=['accuracy'])

# 可适当增加epochs，给贝叶斯模型更多训练时间
num_epochs = 60
batch_size = 48

history = model.fit(X_train, y_train_encoded, epochs=num_epochs, batch_size=batch_size,
                    validation_data=(X_test, y_test_encoded), verbose=2)

额外优化建议 #

• KL散度退火：训练初期将KL散度的权重设为0，逐渐增加到1，帮助模型先拟合数据特征，再引入正则约束。
• 参数初始化：给贝叶斯层的参数设置较小的初始标准差（比如0.01），避免初始分布过宽导致训练波动。
• 批量大小：如果显存允许，适当增大batch_size，稳定梯度计算，减少训练波动。

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