def testDistributedModelFit(self, strategy):
if not tf.__internal__.tf2.enabled() and isinstance(
strategy, tf.distribute.experimental.ParameterServerStrategy):
self.skipTest(
"Parameter Server strategy with dataset creator need to be run "
"when eager execution is enabled.")
with strategy.scope():
preprocessing_model = utils.make_preprocessing_model(
self.get_temp_dir())
training_model = utils.make_training_model()
# Merge the two separate models into a single model for training.
inputs = preprocessing_model.inputs
outputs = training_model(preprocessing_model(inputs))
merged_model = tf.keras.Model(inputs, outputs)
merged_model.compile(optimizer="sgd", loss="binary_crossentropy")
def dataset_fn(input_context):
dataset = utils.make_dataset()
dataset = dataset.shard(
input_context.num_input_pipelines,
input_context.input_pipeline_id,
)
batch_size = input_context.get_per_replica_batch_size(
global_batch_size=utils.BATCH_SIZE)
return dataset.batch(batch_size).repeat().prefetch(2)
dataset_creator = tf.keras.utils.experimental.DatasetCreator(
dataset_fn)
merged_model.fit(dataset_creator,
epochs=2,
steps_per_epoch=utils.STEPS)
def testDistributedModelFit(self, strategy):
with strategy.scope():
preprocessing_model = utils.make_preprocessing_model(
self.get_temp_dir())
training_model = utils.make_training_model()
training_model.compile(optimizer="sgd", loss="binary_crossentropy")
dataset = utils.make_dataset()
dataset = dataset.batch(utils.BATCH_SIZE)
dataset = dataset.map(lambda x, y: (preprocessing_model(x), y))
training_model.fit(dataset, epochs=2)
def testDistributedModelFit(self, strategy):
with strategy.scope():
preprocessing_model = utils.make_preprocessing_model(self.get_temp_dir())
training_model = utils.make_training_model()
training_model.compile(optimizer="sgd", loss="binary_crossentropy")
def dataset_fn(input_context):
dataset = utils.make_dataset()
dataset = dataset.shard(input_context.num_input_pipelines,
input_context.input_pipeline_id)
batch_size = input_context.get_per_replica_batch_size(
global_batch_size=utils.BATCH_SIZE)
dataset = dataset.batch(batch_size).repeat().prefetch(2)
return dataset.map(lambda x, y: (preprocessing_model(x), y))
dataset_creator = tf.keras.utils.experimental.DatasetCreator(dataset_fn)
training_model.fit(dataset_creator, epochs=2, steps_per_epoch=utils.STEPS)