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 参数时，必须提供此类输入的基数，因为无法推断此类输入的基数。