Keras自定义层中的中间张量是否需注册为non-trainable-weight？

首先直接给你结论：不需要把call方法里计算的mean和std注册为non-trainable-weight，原因和你的内存顾虑分析如下：

• 先明确核心概念：non-trainable-weight是层的持久化参数，比如BatchNormalization层里的moving_mean和moving_variance，这类参数会被保存在模型中、伴随整个训练/推理过程，只是不参与梯度更新。而你在call里计算的mean和std是临时中间张量——它们只是前向传播时临时生成的计算结果，完成当前步骤的计算后就会被自动释放（除非被后续计算节点引用），本质上不属于层的“权重”范畴。

• 关于GPU内存占用：这些中间张量确实会在计算时占用GPU内存，但model.summary()的内存统计只针对层的可训练/不可训练权重，不会统计临时计算张量的内存。如果想要估算真实的内存开销，model.summary()的结果是不够的，你可以用TensorFlow的Profiler工具（比如tf.profiler.experimental.start()）来追踪计算过程中的内存使用，或者直接通过nvidia-smi（如果使用NVIDIA GPU）查看实时显存占用。

• 额外提醒：如果强行把这些临时张量注册为non-trainable-weight，反而会导致它们变成持久化的张量，一直占用内存——这完全没必要，因为每次前向传播的输入都可能不同，对应的mean和std也是动态计算的，根本不需要保存下来。

总结一下，你当前的代码写法是正确的，不需要做额外修改。如果要精准分析内存消耗，建议使用TensorFlow的专业内存分析工具，而不是依赖model.summary()的统计。

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