def train_func(config):
batch_size = config.get("batch_size", 64)
epochs = config.get("epochs", 3)
strategy = tf.distribute.MultiWorkerMirroredStrategy()
with strategy.scope():
# Model building/compiling need to be within `strategy.scope()`.
multi_worker_model = build_and_compile_model(config)
dataset_pipeline = train.get_dataset_shard()
dataset_iterator = dataset_pipeline.iter_epochs()
results = []
for _ in range(epochs):
dataset = next(dataset_iterator)
tf_dataset = prepare_dataset_shard(
dataset.to_tf(
label_column="y",
output_signature=(
tf.TensorSpec(shape=(None, 1), dtype=tf.float32),
tf.TensorSpec(shape=(None), dtype=tf.float32),
),
batch_size=batch_size,
))
history = multi_worker_model.fit(tf_dataset,
callbacks=[TrainReportCallback()])
results.append(history.history)
return results
def train_func(config: dict):
strategy = tf.distribute.MultiWorkerMirroredStrategy()
with strategy.scope():
# Model building/compiling need to be within `strategy.scope()`.
multi_worker_model = build_model()
multi_worker_model.compile(
optimizer=tf.keras.optimizers.SGD(
learning_rate=config.get("lr", 1e-3)),
loss=tf.keras.losses.mean_squared_error,
metrics=[tf.keras.metrics.mean_squared_error],
)
dataset = session.get_dataset_shard("train")
for _ in range(config.get("epoch", 3)):
tf_dataset = prepare_dataset_shard(
dataset.to_tf(
label_column="y",
output_signature=(
tf.TensorSpec(shape=(None, 1), dtype=tf.float32),
tf.TensorSpec(shape=(None), dtype=tf.float32),
),
batch_size=32,
))
multi_worker_model.fit(tf_dataset, callbacks=[Callback()])
def train_func(config: dict):
batch_size = config.get("batch_size", 64)
epochs = config.get("epochs", 3)
strategy = tf.distribute.MultiWorkerMirroredStrategy()
with strategy.scope():
# Model building/compiling need to be within `strategy.scope()`.
multi_worker_model = build_model()
multi_worker_model.compile(
optimizer=tf.keras.optimizers.SGD(learning_rate=config.get("lr", 1e-3)),
loss=tf.keras.losses.mean_squared_error,
metrics=[tf.keras.metrics.mean_squared_error],
)
dataset = train.get_dataset_shard("train")
results = []
for _ in range(epochs):
tf_dataset = prepare_dataset_shard(
dataset.to_tf(
label_column="y",
output_signature=(
tf.TensorSpec(shape=(None, 1), dtype=tf.float32),
tf.TensorSpec(shape=(None), dtype=tf.float32),
),
batch_size=batch_size,
)
)
history = multi_worker_model.fit(
tf_dataset, callbacks=[TrainCheckpointReportCallback()]
)
results.append(history.history)
return results
def to_tf_dataset(dataset, batch_size):
def to_tensor_iterator():
data_iterator = dataset.iter_batches(
batch_format="numpy", batch_size=batch_size
)
for d in data_iterator:
yield (
# "concat_out" is the output column of the Concatenator.
tf.convert_to_tensor(d["concat_out"], dtype=tf.float32),
tf.convert_to_tensor(d["target"], dtype=tf.float32),
)
output_signature = (
tf.TensorSpec(shape=(None, num_features), dtype=tf.float32),
tf.TensorSpec(shape=(None), dtype=tf.float32),
)
tf_dataset = tf.data.Dataset.from_generator(
to_tensor_iterator, output_signature=output_signature
)
return prepare_dataset_shard(tf_dataset)
def train_func(config: dict):
per_worker_batch_size = config.get("batch_size", 64)
epochs = config.get("epochs", 3)
dataset_shard = session.get_dataset_shard("train")
strategy = tf.distribute.MultiWorkerMirroredStrategy()
with strategy.scope():
# Model building/compiling need to be within `strategy.scope()`.
multi_worker_model = build_autoencoder_model()
learning_rate = config.get("lr", 0.001)
multi_worker_model.compile(
loss=tf.keras.losses.BinaryCrossentropy(),
optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
metrics=[
"binary_crossentropy",
],
)
results = []
for epoch in range(epochs):
tf_dataset = prepare_dataset_shard(
dataset_shard.to_tf(
feature_columns=["image"],
label_column="label",
output_signature=(
tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
),
batch_size=per_worker_batch_size,
))
history = multi_worker_model.fit(
tf_dataset, callbacks=[TrainCheckpointReportCallback()])
results.append(history.history)
return results