自定义损失函数中Tensor超出作用域问题求解

问题 #

处理批量图像序列的5D张量（形状通常为[batch_size, sequence_length, height, width, channels]）时，想要实现总变差损失，于是编写了自定义损失函数avg_total_variation，遍历序列中的每个4D张量，计算tf.image.total_variation并求和后取平均。代码如下：

@tf.function
def avg_total_variation(x, batch_of_sequences_y):
    """
    tf.image.total_variation accepts either a 3D tensor for an image
    or a 4D tensor for a batch of images.
    But we have a 5D tensor for a batch of sequence of images.
    This custom total variation averages the total_variation for each sequence
    """
    variations = []
    for s in batch_of_sequences_y:
        variations.append(tf.reduce_sum(tf.image.total_variation(s)))
    print(variations)
    return tf.reduce_mean(variations)

在模型训练阶段执行：

model.fit(x, y, epochs=EPOCHS, batch_size=BATCH_SIZE, validation_split=0.10)

出现以下错误：

InaccessibleTensorError: <tf.Tensor 'while/Sum:0' shape=() dtype=float32> is out of scope 
and cannot be used here.
Use return values, explicit Python locals or TensorFlow collections to access it.

解决方法 #

错误根源是@tf.function编译时，Python原生的for循环和列表操作会被转换为TensorFlow图结构，循环内创建的张量被限制在循环作用域内，无法被外部的tf.reduce_mean访问。需要用TensorFlow原生张量操作替代Python循环，以下是两种可行方案：

方案1：用tf.map_fn处理序列维度 #

@tf.function
def avg_total_variation(x, batch_of_sequences_y):
    # 定义单序列的总变差计算逻辑
    def compute_seq_tv(seq):
        return tf.reduce_sum(tf.image.total_variation(seq))
    
    # 批量遍历每个序列计算总变差和
    variations = tf.map_fn(compute_seq_tv, batch_of_sequences_y)
    # 对所有序列的总变差取平均
    return tf.reduce_mean(variations)

方案2：直接利用维度Reduce操作（更高效） #

tf.image.total_variation支持直接处理5D张量，返回形状为[batch_size, sequence_length]的结果（对应每个序列中每张图的总变差），可以直接通过维度求和、取平均完成计算：

@tf.function
def avg_total_variation(x, batch_of_sequences_y):
    # 计算每个图像的总变差，得到[batch_size, sequence_length]张量
    tv_per_image = tf.image.total_variation(batch_of_sequences_y)
    # 对每个序列的所有图像总变差求和，得到[batch_size]张量
    tv_per_seq = tf.reduce_sum(tv_per_image, axis=1)
    # 对批量内所有序列的总变差取平均
    return tf.reduce_mean(tv_per_seq)

补充说明 #

• 方案2无需显式循环，完全基于TensorFlow的广播和Reduce操作实现，计算效率更高，也更贴合TensorFlow的图计算逻辑。
• 自定义损失函数的参数顺序需注意：第一个参数是模型输入x，第二个是真实标签y，即使损失计算不需要x，也必须保留参数位置以适配Keras的损失函数接口。

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