如何测试带可训练权重的Keras自定义层代码正确性？

我之前在项目里经常需要写Keras自定义层，踩过不少测试的坑，总结了一套实用的测试方法，分享给你：

1. 基础功能验证：前向传播输出是否符合预期 #

最核心的第一步是验证层的前向计算逻辑是否正确。我们可以通过固定权重/偏置的初始值，手动计算预期输出，再和层的实际输出对比，避免随机初始化带来的不确定性。

用Python标准的unittest框架写个例子：

import tensorflow as tf
from tensorflow.keras.layers import Layer
import unittest

# 自定义一个简单的线性层
class LinearLayer(Layer):
    def __init__(self, units=32):
        super().__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer='random_normal',
            trainable=True
        )
        self.b = self.add_weight(
            shape=(self.units,),
            initializer='zeros',
            trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

# 单元测试类
class TestLinearLayer(unittest.TestCase):
    def test_forward_pass(self):
        # 初始化层并指定输出单元数
        layer = LinearLayer(units=1)
        # 手动构建层（指定输入形状）
        layer.build(input_shape=(None, 2))
        # 固定权重和偏置的值，方便计算预期结果
        layer.w.assign(tf.constant([[3.0], [4.0]]))
        layer.b.assign(tf.constant([5.0]))
        
        # 输入测试数据
        inputs = tf.constant([[1.0, 2.0]])
        output = layer(inputs)
        
        # 手动计算预期输出：1*3 + 2*4 +5 = 16
        expected_output = tf.constant([[16.0]])
        # 验证输出误差在可接受范围内
        self.assertTrue(tf.reduce_all(tf.abs(output - expected_output) < 1e-6))

if __name__ == '__main__':
    unittest.main()

2. 可训练权重的正确性验证 #

自定义层的核心价值往往在于可训练权重，所以必须验证权重是否能正确更新、梯度计算是否准确。

方法一：手动推导梯度，对比自动求导结果 #

用tf.GradientTape记录梯度，再和手动计算的梯度对比，确保反向传播逻辑正确：

def test_weight_gradients(self):
    layer = LinearLayer(units=1)
    layer.build(input_shape=(None, 2))
    # 固定初始权重
    layer.w.assign(tf.constant([[3.0], [4.0]]))
    layer.b.assign(tf.constant([5.0]))
    
    inputs = tf.constant([[1.0, 2.0]])
    target = tf.constant([[20.0]])
    
    # 用GradientTape记录计算过程
    with tf.GradientTape() as tape:
        output = layer(inputs)
        loss = tf.reduce_mean(tf.square(output - target))
    
    # 获取权重和偏置的梯度
    grads = tape.gradient(loss, [layer.w, layer.b])
    grad_w, grad_b = grads
    
    # 手动计算梯度：损失=(16-20)^2=16，对w的梯度是2*(16-20)*inputs = [-8, -16]，对应形状(2,1)
    expected_grad_w = tf.constant([[-8.0], [-16.0]])
    # 对b的梯度是2*(16-20)*1 = -8
    expected_grad_b = tf.constant([-8.0])
    
    self.assertTrue(tf.reduce_all(tf.abs(grad_w - expected_grad_w) < 1e-6))
    self.assertTrue(tf.reduce_all(tf.abs(grad_b - expected_grad_b) < 1e-6))

方法二：验证训练过程中权重是否更新 #

用优化器执行一步训练，检查权重是否和初始值不同，确保权重确实是可训练的：

def test_weight_update(self):
    layer = LinearLayer(units=1)
    layer.build(input_shape=(None, 2))
    # 保存初始权重的副本
    initial_w = tf.identity(layer.w)
    initial_b = tf.identity(layer.b)
    
    optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
    inputs = tf.constant([[1.0, 2.0]])
    target = tf.constant([[20.0]])
    
    # 执行一步训练
    with tf.GradientTape() as tape:
        output = layer(inputs)
        loss = tf.reduce_mean(tf.square(output - target))
    grads = tape.gradient(loss, layer.trainable_weights)
    optimizer.apply_gradients(zip(grads, layer.trainable_weights))
    
    # 验证权重已更新
    self.assertFalse(tf.reduce_all(tf.equal(layer.w, initial_w)))
    self.assertFalse(tf.reduce_all(tf.equal(layer.b, initial_b)))

3. 边缘情况与兼容性测试 #

还要覆盖一些容易忽略的场景，确保层的鲁棒性：

• 不同输入形状/批量大小：比如测试输入形状为(None, 3)、批量大小为0或1的情况
• 序列化/反序列化：验证层能否正确保存和加载，避免部署时出问题

序列化测试的例子：

def test_serialization(self):
    layer = LinearLayer(units=5)
    # 构建层
    layer.build(input_shape=(None, 3))
    # 生成测试输入
    inputs = tf.random.normal((2, 3))
    original_output = layer(inputs)
    
    # 序列化和反序列化层
    serialized_layer = tf.keras.layers.serialize(layer)
    deserialized_layer = tf.keras.layers.deserialize(serialized_layer)
    # 反序列化后需要重新构建层
    deserialized_layer.build(input_shape=(None, 3))
    # 复制原层的权重
    deserialized_layer.set_weights(layer.get_weights())
    
    # 验证输出一致
    deserialized_output = deserialized_layer(inputs)
    self.assertTrue(tf.reduce_all(tf.abs(original_output - deserialized_output) < 1e-6))

4. 利用Keras内置工具简化测试 #

Keras提供了tf.keras.testing_utils.layer_test工具，能自动测试层的基本功能（前向传播、序列化、不同输入形状兼容性等），省了不少重复代码：

def test_with_keras_test_utils(self):
    from tensorflow.keras import testing_utils
    # 自动测试LinearLayer的多种场景
    testing_utils.layer_test(
        LinearLayer,
        kwargs={'units': 5},
        input_shape=(3, 4)  # 输入形状(批量大小3，特征数4)
    )

以上方法覆盖了自定义层的核心测试场景，基本上能保证你的层逻辑正确、训练正常、兼容性良好。

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