Python tf.keras.utils.experimental.DatasetCreator用法及代碼示例

調用時返回 tf.data.Dataset 的對象。

用法

tf.keras.utils.experimental.DatasetCreator(
    dataset_fn, input_options=None
)

tf.keras.utils.experimental.DatasetCreator 被指定為 x 或 tf.keras.Model.fit 中的輸入的支持類型。當使用返回 tf.data.Dataset 的可調用(帶有 input_context 參數)時，將此類的實例傳遞給 fit 。

model = tf.keras.Sequential([tf.keras.layers.Dense(10)])
model.compile(tf.keras.optimizers.SGD(), loss="mse")

def dataset_fn(input_context):
  global_batch_size = 64
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat()
  dataset = dataset.shard(
      input_context.num_input_pipelines, input_context.input_pipeline_id)
  dataset = dataset.batch(batch_size)
  dataset = dataset.prefetch(2)
  return dataset

input_options = tf.distribute.InputOptions(
    experimental_fetch_to_device=True,
    experimental_per_replica_buffer_size=2)
model.fit(tf.keras.utils.experimental.DatasetCreator(
    dataset_fn, input_options=input_options), epochs=10, steps_per_epoch=10)

Model.fit 與DatasetCreator 一起使用旨在適用於所有tf.distribute.Strategy，隻要在模型創建時使用Strategy.scope：

strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver)
with strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(10)])
model.compile(tf.keras.optimizers.SGD(), loss="mse")

def dataset_fn(input_context):
  ...

input_options = ...
model.fit(tf.keras.utils.experimental.DatasetCreator(
    dataset_fn, input_options=input_options), epochs=10, steps_per_epoch=10)

注意：在Model.fit 中使用DatasetCreator , steps_per_epoch 參數時，必須提供此類輸入的基數，因為無法推斷此類輸入的基數。