更改优化器时@tf.function 错误

如何解决更改优化器时@tf.function 错误

我正在尝试使用 tf.function 处理简单模型的训练步骤。我的 train_step 函数如下所示：

def loss_computation(model,x,y):
    y_logits = model(x)
    cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=y_logits,labels=y)
    return tf.reduce_mean(cross_ent)

@tf.function
def train_step(optimizer,model,y):
    with tf.GradientTape() as tape:
        loss = loss_computation(model,y)
    gradients = tape.gradient(loss,model.trainable_variables)
    optimizer.apply_gradients(zip(gradients,model.trainable_variables))

然后我循环调用 train_step ：

optimizer = tf.keras.optimizers.Adam(1e-4)

batch_size=10
dataset = tf.data.Dataset.from_tensor_slices((scaled_images_np_for_train,train_targets_np)).batch(batch_size)

for epoch in range(3):
    for batch in pb.progressbar(dataset,max_value=scaled_images_np_for_train.shape[0]//batch_size):
    train_step(optimizer,batch[0],batch[1])

但是，假设我想更改优化器或优化器步骤。我发现如果我创建一个新的优化器，然后像这样运行 train_step：

optimizer_2 = tf.keras.optimizers.Adam(1e-5)

for epoch in range(3):
    for batch in pb.progressbar(dataset,max_value=scaled_images_np_for_train.shape[0]//batch_size):
        train_step(optimizer_2,batch[1])

我收到以下错误：

ValueError: in user code:

    <ipython-input-83-d8473382c9f9>:6 train_step  *
        optimizer.apply_gradients(zip(gradients,model.trainable_variables))
    /data0/matthew/Software/ML_Learning/venv/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:478 apply_gradients  **
        self._create_all_weights(var_list)
    /data0/matthew/Software/ML_Learning/venv/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:661 _create_all_weights
        _ = self.iterations
    /data0/matthew/Software/ML_Learning/venv/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:668 __getattribute__
    --------
    Rest of the callstack
    --------
    /data0/matthew/Software/ML_Learning/venv/lib/python3.8/site-packages/tensorflow/python/ops/variables.py:259 __call__
        return cls._variable_v1_call(*args,**kwargs)
    /data0/matthew/Software/ML_Learning/venv/lib/python3.8/site-packages/tensorflow/python/ops/variables.py:205 _variable_v1_call
        return prevIoUs_getter(
    /data0/matthew/Software/ML_Learning/venv/lib/python3.8/site-packages/tensorflow/python/ops/variables.py:66 getter
        return captured_getter(captured_prevIoUs,**kwargs)
    /data0/matthew/Software/ML_Learning/venv/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py:510 invalid_creator_scope
        raise ValueError(

    ValueError: tf.function-decorated function tried to create variables on non-first call.

调用堆栈的顶部表明行 optimizer.apply_gradients(zip(gradients,model.trainable_variables)) 似乎是问题所在，具体而言，优化器正在创建 tensorflow 变量以使用具有此调用的给定模型在“第一次”调用时保持梯度：{{ 1}}

但是，我很困惑.. 难道self._create_all_weights(var_list) 不应该认识到我已经将不同的优化器传递给tf.function 并因此创建并跟踪一个新图吗？

如果我创建一个可能具有不同潜在变量维度的新模型，并尝试使用与之前模型相同的优化器，则会发生同样的事情。同样，train_step 是否应该检测到 tf.function 正在使用不同的模型调用并创建一个新图的跟踪？

可能是我没有正确使用优化器，并且每个优化器只能与特定模型一起使用一次，我不知道，但我对 tf.function 的直觉告诉我，任何错误都不应该t 来自 tf.function 进程..

train_step 是不是不够智能，无法自动检测比简单张量 dtype 变化或张量形状变化更复杂的变化？我可以为 tf.function 指定一个特殊的签名来“帮助”它，让它做我想做的事吗？

对这个问题的任何澄清都会很棒。