自定义TensorFlow层正则化损失未注册及多GPU报错求助

核心问题原因 #

1. 正则化损失未被全局收集 #

你用conv2d.apply(inputs)而非直接调用conv2d(inputs)（层的__call__方法）是关键问题。apply()是TensorFlow早期的底层方法，它不会触发Keras层完整的生命周期钩子——其中就包含正则化损失的全局注册逻辑。而tf.layers.conv2d内部实际调用的是层的__call__，这也是你对比源码没发现差别的原因（你可能没注意到它的调用方式）。

当你手动打印conv2d.losses时能看到损失，是因为层自身已经计算了正则化损失，但这些损失没有被添加到tf.losses的全局集合中，所以get_regularization_losses()返回空。

2. 多GPU环境下打印张量报错 #

在MirroredStrategy的塔上下文（每个GPU的计算分支）中，张量是分布式的MirroredTensor，直接打印会触发张量转换的断言错误——普通打印会尝试将张量转为numpy数组，而分布式张量不允许这种直接的引用转换（对应你遇到的assert not as_ref异常）。

修复步骤 #

1. 替换apply()为层的__call__调用 #

修改你的my_conv2d接口，直接用conv2d(inputs)替代conv2d.apply(inputs)：

def my_conv2d(inputs, kernel_regularizer="glorot_uniform"):
    conv2d = MyConv2D(kernel_regularizer=kernel_regularizer)
    # 用__call__替代apply，触发完整的层生命周期逻辑
    return conv2d(inputs)

__call__方法会自动处理：

• 层的build方法调用（如果未构建）
• 正则化损失的计算与全局注册
• 后续tf.losses.get_regularization_losses()就能正确获取到这些损失

2. 移除塔上下文内的张量打印 #

删掉print(conv2d.losses)这行代码——如果需要调试正则化损失，可以在策略范围之外（比如模型构建完成后）打印，或者用tf.print（它支持分布式张量的打印）：

# 在MirroredStrategy.scope()之外或者训练循环的合适位置
model = tf.keras.Model(inputs=..., outputs=...)
print("正则化损失:", model.losses)
# 或者用tf.print在训练中打印
tf.print("当前正则化损失:", tf.add_n(model.losses))

3. 确保正则化损失在多GPU下被正确聚合 #

当使用MirroredStrategy时，Keras会自动聚合各个GPU上的正则化损失，你只需要在计算总损失时将正则化损失加入即可：

with strategy.scope():
    model = build_model()  # 包含你的my_conv2d层

def compute_loss(labels, predictions):
    per_example_loss = tf.keras.losses.sparse_categorical_crossentropy(labels, predictions)
    loss = tf.nn.compute_average_loss(per_example_loss, global_batch_size=BATCH_SIZE)
    # 添加正则化损失
    loss += tf.add_n(model.losses)
    return loss

验证方法 #

1. 单GPU环境下，调用tf.losses.get_regularization_losses()应该能返回你的自定义层的正则化损失值
2. 多GPU环境下，运行训练代码不再抛出_tensor_conversion_mirrored的断言异常，且正则化损失会被正确计入总损失

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