TensorFlow 2胶囊网络在FashionMNIST上准确率停滞10%求助

核心问题总结 #

• 模型训练后准确率始终≤10%（随机猜测水平）
• 自定义损失函数训练1个epoch后损失变为负数
• 更换margin loss后准确率仍无提升，且出现无梯度错误
• 调整模型输出为胶囊长度时出现维度不匹配错误

一、损失函数的关键问题修复 #

1. 自定义损失为负的根源 #

胶囊网络的损失（如margin loss）基于胶囊长度设计，若计算逻辑出现符号或阈值错误，会导致损失为负，模型反向传播完全失效：
正确的margin loss实现示例：

def margin_loss(y_true, y_pred):
    # y_true: one-hot编码标签，shape (batch, 10)
    # y_pred: 输出胶囊的L2长度，shape (batch, 10)
    pos_loss = y_true * tf.square(tf.maximum(0.0, 0.9 - y_pred))
    neg_loss = (1 - y_true) * tf.square(tf.maximum(0.0, y_pred - 0.1))
    return tf.reduce_mean(pos_loss + 0.5 * neg_loss)

若你的自定义损失将阈值（0.9/0.1）搞反、符号错误，或把损失逻辑写成了奖励（比如用减法代替平方损失），会直接导致损失为负，模型无法学习。

2. 无梯度错误的排查 #

无梯度错误几乎都是因为链路中存在不可微分操作：

• 禁止对胶囊长度使用tf.round()或numpy的round操作，这类操作会切断梯度传播；
• 动态路由部分必须用TF2原生可微分控制流（如tf.while_loop），避免使用Python原生for循环或不可微分的变量更新逻辑；
• 检查胶囊层的权重是否被正确标记为可训练变量，而非普通张量。

二、维度匹配问题修复 #

FashionMNIST为10分类任务，需确保模型输出与标签维度严格对齐：

1. 模型输出处理：输出胶囊层的原始输出是(batch_size, 10, capsule_dim)（如16维胶囊），必须先计算每个胶囊的L2长度：

   capsule_lengths = tf.norm(output_capsules, axis=-1)  # shape: (batch_size, 10)
   

   再将该张量传入损失函数，直接传入原始胶囊向量会导致维度不匹配。
2. 标签格式：确保标签是one-hot编码（shape: (batch_size,10)），若使用整数标签（shape: (batch_size,)），需修改损失函数适配稀疏标签：

   def sparse_margin_loss(y_true, y_pred):
       y_true_one_hot = tf.one_hot(y_true, depth=10)
       pos_loss = y_true_one_hot * tf.square(tf.maximum(0.0, 0.9 - y_pred))
       neg_loss = (1 - y_true_one_hot) * tf.square(tf.maximum(0.0, y_pred - 0.1))
       return tf.reduce_mean(pos_loss + 0.5 * neg_loss)
   

三、训练参数与TF1→TF2适配调整 #

1. 学习率调整：你当前使用的1e-05过小，TF2优化器的实现与TF1存在差异，胶囊网络的合理学习率范围为1e-3 ~ 1e-4，过小的学习率会导致模型参数几乎不更新，始终停留在随机初始化状态。
2. 动态路由实现对齐：TF1版本的动态路由依赖变量作用域和tf.while_loop，TF2中需确保路由权重是可训练的tf.Variable，且迭代过程完全在tf.function内执行，避免跨图操作。
3. 权重初始化对齐：检查胶囊层的权重初始化是否与TF1版本一致（如使用Xavier初始化），TF2默认的GlorotUniform与TF1的xavier_initializer行为一致，但需确保卷积层、胶囊层的初始化参数完全匹配。

四、简化模型验证步骤 #

先从极简模型开始验证，逐步排查问题：

1. 剥离动态路由：先实现静态路由的胶囊网络（即输出胶囊直接由初级胶囊加权求和得到，不做路由迭代），验证模型是否能突破随机准确率（>10%）；
2. 损失函数独立测试：用普通Dense层替代输出胶囊层，测试损失函数是否能正常计算正损失值，且梯度可传播；
3. 数据验证：再次确认输入数据的归一化是否正确（像素值范围[0,1]），标签与输入的对应关系无误。

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