时间戳

首页 > 解决方案 > 自定义层中的所有中间张量都应该注册为“不可训练的权重”吗？

python - 自定义层中的所有中间张量都应该注册为“不可训练的权重”吗？

问题描述

我正在用Keras. 但是“不可训练的重量”让我感到困惑。假设我有一个简单的层如下（它只做一件事，将平均值除以标准）：

import tensorflow as tf
K = tf.keras.backend
class SharpeLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(SharpeLayer, self).__init__()
        
    def call(self, inputs):
        return K.mean(inputs) / (K.std(inputs) + K.epsilon()) * 15.87450787

有两个中间张量：mean和std。我应该将这两个张量注册为non-trainable-weight吗？

我的想法是：我知道这不会影响最终的训练结果。我的目标是估计我的模型的内存消耗。我不确定中间张量是否也会消耗 GPU 内存。如果是这样，我认为所有中间张量都应该注册为non-trainable-weight，这样model.summary可以反映真实的内存成本。

标签： pythontensorflowkerasdeep-learningout-of-memory

解决方案

您应该只注册 Layer 类的实例变量。中间张量是由 tensorflow 创建的，你不应该关心它们。

将张量注册为可训练/不可训练不会改变它们的内存消耗。它决定是否在梯度的反向传播过程中更新这些张量。

推荐阅读