能否通过条件控制Python with语句的启用/禁用？TensorFlow Keras场景

问题描述 #

我正在为TensorFlow Keras模型编写自定义训练代码，需要区分更新权重的train_step和仅推理的test_step，希望实现一个能同时适配两者的统一函数。我的初步思路是通过条件控制tf.GradientTape的with语句启用/禁用，同时保留块内代码执行，具体设想代码如下：

def _train_or_evaluate(self, inputs, gt1, gt2, is_training=False):
    with tf.GradientTape() as tape1 'if is_training': # 期望通过条件启用/禁用with语句
        model1_output = self.model1(inputs)
        model1_loss = self.loss_obj(gt1, model1_output)

    inputs2 = self.process_output(model1_output)

    with tf.GradientTape() as tape2 'if is_training': # 期望通过条件启用/禁用with语句
        model2_output = self.model2(inputs2)
        model2_loss = self.loss_obj2(gt2, model2_output)

    if is_training:
        model1_gradients = tape1.gradient(model1_loss, self.model1.trainable_variables)
        self.optimizer1.apply_gradients(model1_gradients, self.model1.trainable_variables)

        model2_gradients = tape2.gradient(model2_loss, self.model2.trainable_variables)
        self.optimizer2.apply_gradients(model2_gradients, self.model2.trainable_variables)

    return model1_loss, model2_loss

def train_step(self, inputs):
    inputs = inputs, (gt1, gt2)
    return self._train_or_evaluate(inputs, gt1, gt2, True)

def test_step(self, inputs):
    inputs = inputs, (gt1, gt2)
    return self._train_or_evaluate(inputs, gt1, gt2, False)

核心需求是：

• 当is_training=True时，等效于执行完整训练逻辑（带GradientTape记录梯度并更新权重）
• 当is_training=False时，等效于仅执行推理计算损失（无GradientTape包裹）

可行实现方案 #

Python不支持直接在with语句后添加条件判断来启用/禁用上下文，但可以通过以下两种方式实现你的需求：

方案1：利用tf.GradientTape的条件监视 #

通过控制是否让GradientTape监视可训练变量，同时始终保留with块，这样在评估模式下Tape不会记录梯度，效果等同于无with包裹：

def _train_or_evaluate(self, inputs, gt1, gt2, is_training=False):
    # 处理model1部分
    with tf.GradientTape(persistent=False) as tape1:
        # 仅训练模式下监视model1的可训练变量
        if is_training:
            tape1.watch(self.model1.trainable_variables)
        model1_output = self.model1(inputs)
        model1_loss = self.loss_obj(gt1, model1_output)

    inputs2 = self.process_output(model1_output)

    # 处理model2部分
    with tf.GradientTape(persistent=False) as tape2:
        # 仅训练模式下监视model2的可训练变量
        if is_training:
            tape2.watch(self.model2.trainable_variables)
        model2_output = self.model2(inputs2)
        model2_loss = self.loss_obj2(gt2, model2_output)

    if is_training:
        # 计算并应用梯度
        model1_gradients = tape1.gradient(model1_loss, self.model1.trainable_variables)
        self.optimizer1.apply_gradients(zip(model1_gradients, self.model1.trainable_variables))

        model2_gradients = tape2.gradient(model2_loss, self.model2.trainable_variables)
        self.optimizer2.apply_gradients(zip(model2_gradients, self.model2.trainable_variables))

    return model1_loss, model2_loss

def train_step(self, inputs):
    inputs, (gt1, gt2) = inputs  # 修正原代码的解构错误
    return self._train_or_evaluate(inputs, gt1, gt2, is_training=True)

def test_step(self, inputs):
    inputs, (gt1, gt2) = inputs  # 修正原代码的解构错误
    return self._train_or_evaluate(inputs, gt1, gt2, is_training=False)

方案2：条件分支包裹with块，复用核心计算代码 #

把模型前向传播和损失计算的代码抽出来，通过条件判断决定是否用GradientTape包裹，避免重复代码：

def _compute_model1(self, inputs):
    model1_output = self.model1(inputs)
    return model1_output

def _compute_model2(self, inputs2):
    model2_output = self.model2(inputs2)
    return model2_output

def _train_or_evaluate(self, inputs, gt1, gt2, is_training=False):
    if is_training:
        with tf.GradientTape() as tape1:
            model1_output = self._compute_model1(inputs)
            model1_loss = self.loss_obj(gt1, model1_output)
    else:
        model1_output = self._compute_model1(inputs)
        model1_loss = self.loss_obj(gt1, model1_output)

    inputs2 = self.process_output(model1_output)

    if is_training:
        with tf.GradientTape() as tape2:
            model2_output = self._compute_model2(inputs2)
            model2_loss = self.loss_obj2(gt2, model2_output)
    else:
        model2_output = self._compute_model2(inputs2)
        model2_loss = self.loss_obj2(gt2, model2_output)

    if is_training:
        model1_gradients = tape1.gradient(model1_loss, self.model1.trainable_variables)
        self.optimizer1.apply_gradients(zip(model1_gradients, self.model1.trainable_variables))

        model2_gradients = tape2.gradient(model2_loss, self.model2.trainable_variables)
        self.optimizer2.apply_gradients(zip(model2_gradients, self.model2.trainable_variables))

    return model1_loss, model2_loss

def train_step(self, inputs):
    inputs, (gt1, gt2) = inputs
    return self._train_or_evaluate(inputs, gt1, gt2, is_training=True)

def test_step(self, inputs):
    inputs, (gt1, gt2) = inputs
    return self._train_or_evaluate(inputs, gt1, gt2, is_training=False)

关键说明 #

• 方案1更简洁，利用GradientTape的watch方法精准控制监视对象，评估模式下Tape不会记录任何梯度计算，性能开销极低。
• 方案2通过抽离公共代码避免重复，逻辑更直观，适合对TensorFlow API不太熟悉的场景。
• 注意修正原代码中train_step和test_step里的变量解构错误（原代码inputs = inputs, (gt1, gt2)是赋值而非解构，会导致参数传递错误）。

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