Keras中train_step内梯度计算前修改权重的实现问询

现有train_step代码 #

def train_step(self, data):
    # Unpack the data. Its structure depends on your model and
    # on what you pass to `fit()`.
    x, y = data

    with tf.GradientTape() as tape:
        y_pred = self(x, training=True)  # Forward pass
        # Compute the loss value
        # (the loss function is configured in `compile()`)
        loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)

    # Compute gradients
    trainable_vars = self.trainable_variables
    gradients = tape.gradient(loss, trainable_vars)
    # Update weights
    self.optimizer.apply_gradients(zip(gradients, trainable_vars))
    # Update metrics (includes the metric that tracks the loss)
    self.compiled_metrics.update_state(y, y_pred)
    # Return a dict mapping metric names to current value
    return {m.name: m.result() for m in self.metrics}

需求说明 #

需要在train_step函数内修改模型权重：用原始权重计算预测值与损失后，在计算梯度及执行权重更新前修改权重，基于原始权重对应的损失对新权重进行优化。例如：假设有两组相关权重A和B，希望用权重A计算得到的损失来优化权重B。

尝试方案及问题 #

• 方案1：在计算损失后、梯度计算前调用self.set_weights(x)，触发错误：RuntimeError: Cannot get value inside Tensorflow graph function.
• 方案2：使用run_eagerly=True可解决报错，但出现TensorFlow重追踪警告，且结果与预期差异较大，无法确认方案有效性。

可行实现方式 #

可以实现，核心是避开依赖Eager模式的self.set_weights()操作，改用图模式支持的变量赋值API来修改权重。

具体修改示例 #

以下代码以“用权重A的损失优化权重B”为场景，展示修改逻辑：

def train_step(self, data):
    x, y = data

    # 保存所有原始权重的副本（图模式下合法操作）
    original_weights = [var.read_value() for var in self.trainable_variables]

    with tf.GradientTape() as tape:
        # 用原始权重计算预测值和损失
        y_pred = self(x, training=True)
        loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)

    # 按需求修改目标权重（示例：筛选名称包含"weight_b"的变量进行修改）
    for idx, var in enumerate(self.trainable_variables):
        if "weight_b" in var.name:
            # 替换为你的实际权重修改逻辑，比如基于权重A生成新权重B
            new_weight = original_weights[idx] * 1.1  # 示例操作，按需替换
            var.assign(new_weight)

    # 基于原始损失，对修改后的权重计算梯度
    trainable_vars = self.trainable_variables
    gradients = tape.gradient(loss, trainable_vars)
    # 用梯度更新修改后的权重
    self.optimizer.apply_gradients(zip(gradients, trainable_vars))

    self.compiled_metrics.update_state(y, y_pred)
    return {m.name: m.result() for m in self.metrics}

关键注意点 #

• var.read_value()：图模式下安全读取变量当前值的方法，不会触发报错
• var.assign(new_weight)：图模式下支持的变量赋值操作，直接修改权重张量
• 梯度计算逻辑：tf.GradientTape记录的是原始权重的前向计算过程，最终梯度是基于原始损失对修改后权重计算的，完全符合需求

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