tensorflow - tf model.fit() 中的 batch_size 与 tf.data.Dataset 中的 batch_size

Question

我有一个可以放入主机内存的大型数据集。但是，当我使用 tf.keras 训练模型时，它会产生 GPU 内存不足问题。然后我查看 tf.data.Dataset 并希望使用它的 batch() 方法来批处理训练数据集，以便它可以在 GPU 中执行 model.fit()。根据其文档，示例如下:

train_dataset = tf.data.Dataset.from_tensor_slices((train_examples, train_labels))
test_dataset = tf.data.Dataset.from_tensor_slices((test_examples, test_labels))

BATCH_SIZE = 64
SHUFFLE_BUFFER_SIZE = 100

train_dataset = train_dataset.shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE)
test_dataset = test_dataset.batch(BATCH_SIZE)

dataset.from_tensor_slices().batch() 中的 BATCH_SIZE 与 tf.keras modelt.fit() 中的 batch_size 是否相同？
我应该如何选择 BATCH_SIZE 以便 GPU 有足够的数据来高效运行，但其内存不会溢出？

Answer 1

您不需要通过 batch_size model.fit() 中的参数在这种情况下。它将自动使用您在 tf.data.Dataset().batch() 中使用的 BATCH_SIZE .
至于你的另一个问题:批量大小超参数确实需要仔细调整。另一方面，如果你看到 OOM 错误，你应该减少它，直到你没有得到 OOM(通常以这种方式 32 --> 16 --> 8 ...)。
在您的情况下，我会从 2 的 batch_size 开始，然后将其增加 2 的幂，然后检查是否仍然出现 OOM。
您不需要提供 batch_size参数，如果您使用 tf.data.Dataset().batch()方法。
其实连官方documentation声明如下:

batch_size : Integer or None. Number of samples per gradient update.If unspecified, batch_size will default to 32. Do not specify thebatch_size if your data is in the form of datasets, generators, orkeras.utils.Sequence instances (since they generate batches).