Keras模型全层权重与方差可视化问题求助

我尝试可视化以下Keras孪生模型各层的权重均值与标准差，但参考代码仅能展示第一层的权重，无法查看所有层。希望实现所有层的权重可视化，观察训练过程中每个epoch内模型各层的更新情况。

模型定义 #

import tensorflow as tf
import tensorflow.keras.layers
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Flatten, Dense, Dropout, Lambda
import keras


def initialize_base_network():
    # 输入层传递
    input = Input(shape=(100,), name="base_input")
    x = Flatten(name="flatten_input")(input)
    x = Dropout(0.2, name="first_dropout")(x)
    x = Dense(128, activation='relu', name="first_base_dense")(x)
    x = Dense(128, activation='relu', name="second_base_dense")(x)
    x = Dropout(0.1, name="second_dropout")(x)
    x = Dense(128, activation='relu', name="third_base_dense")(x)

    # 返回基础网络模型
    return Model(inputs=input, outputs=x)


def euclidean_distance(vects):
    x, y = vects
    sum_square = K.sum(K.square(x - y), axis=1, keepdims=True)
    return K.sqrt(K.maximum(sum_square, K.epsilon()))


def eucl_dist_output_shape(shapes):
    shape1, shape2 = shapes
    print(shape1[0], shape1, shape2)
    return (shape1[0], 1)


base_network = initialize_base_network()
base_network.summary()

# 创建左输入并连接到基础网络
input_a = Input(shape=(100,), name="left_input")
vect_output_a = base_network(input_a)

# 创建右输入并连接到基础网络
input_b = Input(shape=(100,), name="right_input")
vect_output_b = base_network(input_b)

# 计算两个输出向量的欧式距离
output = Lambda(euclidean_distance, name="output_layer", output_shape=eucl_dist_output_shape)(
    [vect_output_a, vect_output_b])

# 定义完整模型
model = Model([input_a, input_b], output)

model.summary()

模型summary输出：

模型编译与训练代码 #

import keras
from keras import backend as K

from tensorflow.keras.callbacks import Callback
 
class WeightCapture(Callback):
    "捕获模型各层权重"
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.weights = []
        self.epochs = []
 
    def on_epoch_end(self, epoch, logs=None):
        self.epochs.append(epoch) # 记录epoch
        weight = {}
        for layer in model.layers:
            if not layer.weights:
                continue
            name = layer.weights[0].name.split("/")[0]
            weight[name] = layer.weights[0].numpy()
        self.weights.append(weight)
    
    def on_epoch_end(self, epoch, logs=None):
        self.epochs.append(epoch) # 记录epoch
        weight = {}
        for layer in model.layers:
            if not layer.weights:
                continue
            name = layer.weights[0].name.split("/")[0]
            weight[name] = layer.weights[0].numpy()
        self.weights.append(weight)
 
# 自定义准确率计算
def accuracy(y_true, y_pred):
    '''
    基于距离阈值计算分类准确率
    '''
    pred = y_pred.ravel() < 0.5 # 距离小于0.5视为相似样本对
    return np.mean(pred == y_true)


from tensorflow.python.ops.numpy_ops import np_config
np_config.enable_numpy_behavior()

# 自定义对比损失类
class ContrastivLoss(Loss):
    def __init__(self, margin =1):
        super().__init__()
        self.margin = margin
    def call(self, y_true, y_pred):
        square_pred = K.square(y_pred) # y_true为0/1，标记样本对是否相似
        margin_square = K.square(K.maximum(self.margin - y_pred, 0))
        return K.mean(y_true * square_pred + (1 - y_true) * margin_square)


decay_steps = 1000
es = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=30)
capture_cb = WeightCapture(model)
capture_cb.on_epoch_end(-1)
callbacks_list = [capture_cb]

rms = tf.keras.optimizers.Adam(learning_rate=0.0001)
model.compile(loss=ContrastivLoss(margin=1), optimizer=rms, metrics=["accuracy"])
history = model.fit([train_data[:,0], train_data[:,1]], labels_train,
                    validation_data=([test_data[:,0], test_data[:,1]], labels_test), epochs=10, 
                    batch_size=64, callbacks=callbacks_list, verbose=2)

当前问题 #

运行以下绘图代码后，仅能可视化第一层的权重均值与标准差：

def plotweight(capture_cb):
    "绘制各epoch权重的均值和标准差"
    fig, ax = plt.subplots(2, 1, sharex=True, constrained_layout=True, figsize=(8, 10))
    ax[0].set_title("Mean weight")
    for key in capture_cb.weights[0]:
        ax[0].plot(capture_cb.epochs, [w[key].mean() for w in capture_cb.weights], label=key)
    ax[0].legend()
    ax[1].set_title("S.D.")
    for key in capture_cb.weights[0]:
        ax[1].plot(capture_cb.epochs, [w[key].std() for w in capture_cb.weights], label=key)
    ax[1].legend()
    plt.show()
    
plotweight(capture_cb)

绘图输出：

问题根源 #

1. 回调方法重复定义：WeightCapture类中重复写了两次on_epoch_end，只有最后一次定义会生效。
2. 未遍历子模型内部层：外层模型model的层仅包含输入层、base_network子模型和输出层，真正的可训练Dense层都在base_network内部，原回调只遍历外层模型的层，所以仅捕获到了子模型的整体权重，而非内部各层。

修改后的代码 #

1. 修复权重捕获回调

class WeightCapture(Callback):
    "捕获所有层（包括子模型内部层）的权重"
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.weights = []
        self.epochs = []
 
    def on_epoch_end(self, epoch, logs=None):
        self.epochs.append(epoch)
        weight = {}
        # 遍历所有层，递归处理子模型
        for layer in self.model.layers:
            # 如果是子模型，遍历其内部层
            if isinstance(layer, Model):
                for sub_layer in layer.layers:
                    if sub_layer.weights:
                        # 仅保留Dense层的权重（忽略Input、Dropout等无训练权重的层）
                        if 'dense' in sub_layer.name.lower():
                            weight[sub_layer.name] = sub_layer.weights[0].numpy()
            else:
                if layer.weights and 'dense' in layer.name.lower():
                    weight[layer.name] = layer.weights[0].numpy()
        self.weights.append(weight)

2. 优化绘图函数

import matplotlib.pyplot as plt

def plotweight(capture_cb):
    "绘制所有Dense层在各epoch的权重均值与标准差"
    layer_names = list(capture_cb.weights[0].keys())
    # 为每个Dense层单独绘制子图
    fig, axes = plt.subplots(len(layer_names), 2, figsize=(12, 4*len(layer_names)), constrained_layout=True)
    
    for idx, layer_name in enumerate(layer_names):
        # 绘制权重均值变化
        axes[idx,0].plot(capture_cb.epochs, [w[layer_name].mean() for w in capture_cb.weights], color='darkblue')
        axes[idx,0].set_title(f"{layer_name} - 权重均值")
        axes[idx,0].set_xlabel("Epoch")
        axes[idx,0].set_ylabel("均值")
        
        # 绘制权重标准差变化
        axes[idx,1].plot(capture_cb.epochs, [w[layer_name].std() for w in capture_cb.weights], color='darkorange')
        axes[idx,1].set_title(f"{layer_name} - 权重标准差")
        axes[idx,1].set_xlabel("Epoch")
        axes[idx,1].set_ylabel("标准差")
    
    plt.show()

3. 训练时的回调使用

# 初始化回调
capture_cb = WeightCapture(model)
callbacks_list = [capture_cb]

# 编译与训练代码不变
rms = tf.keras.optimizers.Adam(learning_rate=0.0001)
model.compile(loss=ContrastivLoss(margin=1), optimizer=rms, metrics=["accuracy"])
history = model.fit([train_data[:,0], train_data[:,1]], labels_train,
                    validation_data=([test_data[:,0], test_data[:,1]], labels_test), epochs=10, 
                    batch_size=64, callbacks=callbacks_list, verbose=2)

# 生成可视化图
plotweight(capture_cb)

说明 #

• 修改后的回调会递归遍历base_network子模型的内部层，仅捕获Dense层的权重（忽略无训练参数的层）。
• 绘图函数为每个Dense层单独生成均值和标准差的变化曲线，更清晰地观察每层的权重更新趋势。
• 移除了回调中重复定义的on_epoch_end方法，避免逻辑冲突。

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