无法使用 GradientTape 重现 model.fit

Posted 2023-02-16

【中文标题】无法使用 GradientTape 重现 model.fit

【英文标题】：can't reproduce model.fit with GradientTape

【发布时间】：2020-03-07 03:07:39

【问题描述】：

我一直在尝试调查原因（例如，通过在训练期间检查权重、梯度和激活）为什么具有 0.001 学习率的 SGD 在训练中有效，而 Adam 却没有这样做。 （请参阅我之前的帖子 [这里](Why is my loss (binary cross entropy) converging on ~0.6? (Task: Natural Language Inference)"为什么我的损失（二进制交叉熵）收敛到 ~0.6？（任务：自然语言推理）")）

注意：我在这里也使用了与我之前的帖子相同的模型。

---

使用 tf.keras，我使用 model.fit() 训练了神经网络：

model.compile(optimizer=SGD(learning_rate=0.001),
            loss='binary_crossentropy',
            metrics=['accuracy'])

model.fit(x=ds,
        epoch=80,
        validation_data=ds_val)

这导致了如下图所示的 epoch loss，在第一个 epoch 内，它达到了 0.46 的 train_loss，然后最终导致 train_loss 为 0.1241 和 val_loss 为 0.2849。

我会使用tf.keras.callbacks.Tensorboard(histogram_freq=1) 来训练网络，同时使用 SGD(0.001) 和 Adam 进行调查，但它在变量：0 上抛出 InvalidArgumentError，这是我无法破译的。所以我尝试使用 GradientTape 编写一个自定义训练循环并绘制值。

---

使用 tf.GradientTape()，我尝试使用完全相同的模型和数据集来重现结果，但是 epoch 损失的训练速度非常慢，在 15 个 epoch 后达到了 0.676 的训练损失（参见下图），我的实现有问题吗？ （代码如下）

@tf.function
def compute_grads(train_batch: Dict[str,tf.Tensor], target_batch: tf.Tensor, 
                 loss_fn: Loss, model: tf.keras.Model):
    with tf.GradientTape(persistent=False) as tape:
        # forward pass
        outputs = model(train_batch)
        # calculate loss
        loss = loss_fn(y_true=target_batch, y_pred=outputs)

    # calculate gradients for each param
    grads = tape.gradient(loss, model.trainable_variables)
    return grads, loss

BATCH_SIZE = 8
EPOCHS = 15

bce = BinaryCrossentropy()
optimizer = SGD(learning_rate=0.001)

for epoch in tqdm(range(EPOCHS), desc='epoch'):
    # - accumulators
    epoch_loss = 0.0

    for (i, (train_batch, target_dict)) in tqdm(enumerate(ds_train.shuffle(1024).batch(BATCH_SIZE)), desc='step'):

        (grads, loss) = compute_grads(train_batch, target_dict['target'], bce, model)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        epoch_loss += loss

    avg_epoch_loss = epoch_loss/(i+1)
    tensorboard_scalar(writer, name='epoch_loss', data=avg_epoch_loss, step=epoch)  # custom helper function
    print("Epoch : epoch_loss = ".format(epoch, avg_epoch_loss))

提前致谢！

【参考方案1】：

检查您是否对数据集进行了混洗，那么问题可能来自使用 tf.Dataset 方法进行的混洗。它当时只在数据集一个桶中洗牌。使用 Keras.Model.fit 会产生更好的结果，因为它可能会增加另一个洗牌。 通过添加numpy.random.shuffle 的改组，它可以提高训练性能。 From this reference.

将其应用于数据集生成的示例是：

numpy_data = np.hstack([index_rows.reshape(-1, 1), index_cols.reshape(-1, 1), index_data.reshape(-1, 1)])

np.random.shuffle(numpy_data)

indexes = np.array(numpy_data[:, :2], dtype=np.uint32)
labels = np.array(numpy_data[:, 2].reshape(-1, 1), dtype=np.float32)

train_ds = data.Dataset.from_tensor_slices(
    (indexes, labels)
).shuffle(100000).batch(batch_size, drop_remainder=True)

如果这不起作用，您可能需要使用 Dataset .repeat(epochs_number) and .shuffle(..., reshuffle_each_iteration=True):

train_ds = data.Dataset.from_tensor_slices(
    (np.hstack([index_rows.reshape(-1, 1), index_cols.reshape(-1, 1)]), index_data)
    ).shuffle(100000, reshuffle_each_iteration=True
    ).batch(batch_size, drop_remainder=True
    ).repeat(epochs_number)

for ix, (examples, labels) in train_ds.enumerate():
    train_step(examples, labels)
    current_epoch = ix // (len(index_data) // batch_size)

这种解决方法既不美观也不自然，目前您可以使用它来随机播放每个时期。这是一个已知问题，将得到修复，将来您可以使用for epoch in range(epochs_number) 而不是.repeat()

here 提供的解决方案也可能有很大帮助。你可能想检查一下。

如果不是这样，您可能需要加速 TF2.0 GradientTape。这可以是解决方案： TensorFlow 2.0 引入了functions 的概念，将 Eager 代码转换为图形代码。

用法非常简单。唯一需要更改的是所有相关函数（like compute_loss 和 apply_gradients）都必须用 @tf.function. 注释