如何在TensorFlow CNN模型中自定义初始化权重？

你遇到的这种权重初始化敏感导致NaN的情况在CNN里挺常见的——尤其是当你用ReLU这类激活函数时，如果初始权重太大，很容易造成神经元饱和，反向传播时梯度爆炸直接变成NaN；太小的话又会梯度消失，模型根本学不动。下面我给你详细讲几种自定义权重初始化的方法，顺便提几个能帮你避开NaN的实用技巧：

一、用TensorFlow内置初始化器自定义参数 #

TensorFlow本身提供了很多经过验证的初始化器，你可以直接调整它们的参数来适配你的模型，比从头写自定义函数简单得多：

1. He初始化（ReLU激活的黄金搭档） #

这是专门针对ReLU家族激活函数设计的初始化器，能有效缓解梯度消失/爆炸问题，直接调用tf.keras.initializers.HeNormal()或者HeUniform()就行，比如在卷积层里这么用：

model.add(tf.keras.layers.Conv2D(32, (3, 3), activation='relu',
                                 kernel_initializer=tf.keras.initializers.HeNormal(),
                                 bias_initializer=tf.keras.initializers.Zeros()))

2. Xavier/Glorot初始化（适配sigmoid/tanh） #

如果你的模型用的是sigmoid、tanh这类饱和激活函数，用这个初始化器会更合适：

model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='tanh',
                                 kernel_initializer=tf.keras.initializers.GlorotNormal(),
                                 bias_initializer=tf.keras.initializers.Constant(0.1)))

二、完全自定义初始化函数 #

如果内置初始化器满足不了你的需求，你可以自己写初始化逻辑，比如自定义服从特定分布的权重，或者加入截断逻辑避免极端值（极端值是NaN的常见诱因）：

1. 简单自定义张量初始化函数 #

def custom_initializer(shape, dtype=None):
    # 生成截断正态分布，均值0，标准差0.01，超出2倍标准差的值会被重新采样
    return tf.random.truncated_normal(shape, mean=0.0, stddev=0.01, dtype=dtype)

# 在卷积层中使用
model.add(tf.keras.layers.Conv2D(32, (3, 3), activation='relu',
                                 kernel_initializer=custom_initializer))

2. 继承Initializer类（规范复用方式） #

如果需要复用初始化逻辑，或者要加入更复杂的自定义规则，建议继承tf.keras.initializers.Initializer类，还能支持模型序列化：

class MyCustomInitializer(tf.keras.initializers.Initializer):
    def __init__(self, mean=0.0, stddev=0.02):
        self.mean = mean
        self.stddev = stddev

    def __call__(self, shape, dtype=None):
        # 这里可以加自定义逻辑，比如根据输入形状动态调整标准差
        adjusted_stddev = self.stddev * tf.sqrt(2.0 / tf.reduce_prod(shape[:-1]))
        return tf.random.normal(shape, mean=self.mean, stddev=adjusted_stddev, dtype=dtype)

    def get_config(self):
        # 保存配置，方便模型保存和加载
        return {'mean': self.mean, 'stddev': self.stddev}

# 使用示例
model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu',
                                 kernel_initializer=MyCustomInitializer(mean=0.0, stddev=0.03)))

三、额外避免NaN的小技巧 #

除了初始化，这些操作也能帮你从根源上减少NaN问题：

• 梯度裁剪：在编译模型时加入梯度裁剪，限制梯度的最大范数，防止梯度爆炸：

  optimizer = tf.keras.optimizers.Adam(clipnorm=1.0)
  model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy')
  

• 检查损失函数：如果用了自定义损失，要避免除以0或者取对数0的情况，比如加一个极小的epsilon：tf.math.log(x + 1e-8)
• 监控权重和梯度：训练时可以打印权重的均值、方差，或者用TensorBoard监控，看看是不是某些层的权重出现了极端值

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