使用tf.map_fn时TensorFlow梯度为None的问题排查

问题背景 #

我针对虚构的随机图像分类数据集模拟梯度反向传播，场景如下：

• 数据集提供(images,labels)数据对；
• 模型输出形状为(3,batch_size,C)的logits张量（C为类别数量）；
• 需要为每个(i,batch_size,C)（i=0,1,2）分别计算损失。

我实现了三种逻辑一致但实现方式不同的损失计算函数（分别用tf.map_fn、Python列表循环、tf.TensorArray+tf.while_loop），但均无法得到有效梯度，梯度始终返回None。想确认是TensorFlow的Bug还是我的梯度计算逻辑存在错误？要求代码必须在图模式下高效运行。

原始代码 #

import tensorflow as tf


# 模型定义
class MyModel(tf.keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self._l1 = tf.keras.Sequential(
            [
                tf.keras.layers.Conv2D(filters=64, kernel_size=3),
                tf.keras.layers.GlobalAveragePooling2D(keepdims=False),
                tf.keras.layers.Dense(units=10)
            ]
        )
        self._l2 = tf.keras.Sequential(
            [
                tf.keras.layers.Conv2D(filters=64, kernel_size=3),
                tf.keras.layers.GlobalAveragePooling2D(keepdims=False),
                tf.keras.layers.Dense(units=10)
            ]
        )
        self._l3 = tf.keras.Sequential(
            [
                tf.keras.layers.Conv2D(filters=64, kernel_size=3),
                tf.keras.layers.GlobalAveragePooling2D(keepdims=False),
                tf.keras.layers.Dense(units=10)
            ]
        )

    def call(self, inputs, training=None, mask=None):
        y1 = self._l1(inputs)
        y2 = self._l2(inputs)
        y3 = self._l3(inputs)
        out = tf.stack([y1, y2, y3])
        # 输出形状：(3, batch_size, 10)
        return out


model = MyModel()


# 生成随机数据集
def get_dataset():
    db = tf.data.Dataset.range(100)
    db = db.map(
        lambda x: (tf.random.uniform(shape=(224, 224, 3), dtype=tf.float32),
                   tf.one_hot(tf.random.uniform(shape=(), minval=0, maxval=9, dtype=tf.int64), depth=10)
                   )
    )
    db = db.batch(5)
    return db


dataset = get_dataset()


# 损失计算版本1：tf.map_fn实现
def get_loss_v1(logits, labels):
    loss = tf.map_fn(
        fn=lambda x: tf.reduce_mean(
            tf.keras.losses.categorical_crossentropy(
                from_logits=True,
                y_pred=tf.gather(logits, x),
                y_true=labels
            )
        ),
        elems=tf.range(3),
        fn_output_signature=tf.float32
    )
    return loss


# 损失计算版本2：Python列表循环实现
def get_loss_v2(logits, labels):
    losses = list()
    for i in tf.range(3):
        y_pred = tf.gather(logits, i)
        losses.append(
            tf.reduce_mean(
                tf.keras.losses.categorical_crossentropy(
                    from_logits=True,
                    y_pred=y_pred,
                    y_true=labels
                )
            )
        )
    return losses


# 损失计算版本3：tf.TensorArray+tf.while_loop实现
def get_loss_v3(logits, labels):
    losses = tf.TensorArray(dtype=tf.float32,
                            size=0, dynamic_size=True, clear_after_read=False, element_shape=()
                            )
    index = tf.constant(0)

    def body(counter_var, log, lbl, l):
        l = l.write(l.size(), tf.reduce_mean(
            tf.keras.losses.categorical_crossentropy(
                from_logits=True,
                y_pred=tf.gather(log, counter_var),
                y_true=lbl
            )
        ))
        return counter_var + 1, log, lbl, l

    output = tf.while_loop(
        cond=lambda i, *_: tf.less(i, 3),
        loop_vars=(index, logits, labels, losses),
        body=body,
        parallel_iterations=1
    )
    loss_val = output[-1].stack()
    return loss_val


# 训练步骤
def train_step(images, labels):
    with tf.GradientTape(persistent=True) as tape:
        logits = model(images, training=True)
        loss_val = get_loss_v1(logits, labels)
        # loss_val = get_loss_v2(logits, labels)
        # loss_val = get_loss_v3(logits, labels)
        print(loss_val)  # 正向传播损失能正常打印
    grads = list()
    for i in tf.range(3):
        tgt = tf.gather(loss_val, i)
        grads.append(
            tape.gradient(
                target=tgt,
                sources=model.trainable_variables
            )
        )
    print(grads)  # 始终打印None


for data in dataset:
    images, labels = data
    train_step(images, labels)

问题原因分析 #

核心问题出在梯度追踪链路断裂：

1. tape上下文外的张量操作：train_step中用tf.gather(loss_val, i)获取单个损失作为梯度计算目标，这个操作是在tf.GradientTape的上下文之外执行的，tape没有追踪该操作与原始损失张量的关联，无法回溯到模型变量的梯度。
2. 图模式下的循环逻辑错误：用tf.range生成的张量作为Pythonfor循环的迭代对象，在图模式下，这个循环是在图构建阶段静态执行的，无法动态追踪每个迭代中张量的依赖关系，导致梯度链路中断。
3. 高阶API的梯度追踪限制：tf.map_fn、tf.while_loop这类动态控制流操作，在梯度追踪时需要额外的上下文维护，若实现不当容易导致梯度无法正确传递。

修复方案 #

方案1：直接在tape上下文内拆分计算损失 #

跳过独立的损失函数，直接在train_step中拆分模型输出计算每个分支的损失，确保每个损失张量都被tape完整追踪：

def train_step(images, labels):
    with tf.GradientTape(persistent=True) as tape:
        logits = model(images, training=True)
        # 直接拆分三个分支计算损失
        loss_0 = tf.reduce_mean(tf.keras.losses.categorical_crossentropy(
            y_true=labels, y_pred=logits[0], from_logits=True
        ))
        loss_1 = tf.reduce_mean(tf.keras.losses.categorical_crossentropy(
            y_true=labels, y_pred=logits[1], from_logits=True
        ))
        loss_2 = tf.reduce_mean(tf.keras.losses.categorical_crossentropy(
            y_true=labels, y_pred=logits[2], from_logits=True
        ))
        loss_val = tf.stack([loss_0, loss_1, loss_2])
        print(loss_val)
    
    # 直接用tape上下文中生成的损失张量计算梯度
    grads_0 = tape.gradient(loss_0, model.trainable_variables)
    grads_1 = tape.gradient(loss_1, model.trainable_variables)
    grads_2 = tape.gradient(loss_2, model.trainable_variables)
    grads = [grads_0, grads_1, grads_2]
    print(grads)
    # 释放persistent tape
    del tape

方案2：向量化损失计算+正确循环方式 #

用向量化操作替代动态控制流，同时用Python整数循环（而非张量循环）遍历损失分支：

# 优化后的向量化损失函数
def get_loss_v4(logits, labels):
    # 将labels扩展为与logits匹配的形状(3, batch_size, 10)
    labels_expanded = tf.tile(tf.expand_dims(labels, 0), [3, 1, 1])
    # 批量计算所有分支的样本级损失
    per_sample_loss = tf.keras.losses.categorical_crossentropy(
        y_true=labels_expanded, y_pred=logits, from_logits=True
    )
    # 对每个分支的样本损失取均值
    loss_val = tf.reduce_mean(per_sample_loss, axis=1)
    return loss_val

def train_step(images, labels):
    with tf.GradientTape(persistent=True) as tape:
        logits = model(images, training=True)
        loss_val = get_loss_v4(logits, labels)
        print(loss_val)
    
    grads = []
    # 用Python整数循环（3是固定值，图模式下会被静态展开）
    for i in range(3):
        grads.append(tape.gradient(loss_val[i], model.trainable_variables))
    print(grads)
    del tape

关键注意事项 #

• 图模式下，优先使用向量化操作替代动态控制流（如tf.map_fn、while_loop），不仅效率更高，梯度追踪也更稳定。
• tf.GradientTape的梯度计算目标必须是tape上下文内直接生成的张量，避免在tape外对张量做索引、拼接等操作后再计算梯度。
• 当循环次数固定时，用Python整数循环替代张量循环，确保图模式下能正确展开并追踪梯度。

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