wehub-resource-sync
73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
from typing import Dict
|
|
|
|
import ray.tune
|
|
from ray.train.tensorflow import TensorflowCheckpoint
|
|
from ray.train.tensorflow.keras import RayReportCallback
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
_DEPRECATION_MESSAGE = (
|
|
"The `ray.tune.integration.keras` module is deprecated in favor of "
|
|
"`ray.train.tensorflow.keras.ReportCheckpointCallback`."
|
|
)
|
|
|
|
|
|
class TuneReportCallback:
|
|
"""Deprecated.
|
|
Use :class:`ray.train.tensorflow.keras.ReportCheckpointCallback` instead."""
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
raise DeprecationWarning(_DEPRECATION_MESSAGE)
|
|
|
|
|
|
class _TuneCheckpointCallback:
|
|
"""Deprecated.
|
|
Use :class:`ray.train.tensorflow.keras.ReportCheckpointCallback` instead."""
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
raise DeprecationWarning(_DEPRECATION_MESSAGE)
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class TuneReportCheckpointCallback(RayReportCallback):
|
|
"""Keras callback for Ray Tune reporting and checkpointing.
|
|
|
|
.. note::
|
|
Metrics are always reported with checkpoints, even if the event isn't specified
|
|
in ``report_metrics_on``.
|
|
|
|
Example:
|
|
.. code-block:: python
|
|
|
|
############# Using it in Ray Tune ###############
|
|
from ray.tune.integrations.keras import TuneReportCheckpointCallback
|
|
|
|
def train_fn():
|
|
model = build_model()
|
|
model.fit(dataset_shard, callbacks=[TuneReportCheckpointCallback()])
|
|
|
|
tuner = tune.Tuner(train_fn)
|
|
results = tuner.fit()
|
|
|
|
Args:
|
|
metrics: Metrics to report. If this is a list, each item describes
|
|
the metric key reported to Keras, and it's reported under the
|
|
same name. If this is a dict, each key is the name reported
|
|
and the respective value is the metric key reported to Keras.
|
|
If this is None, all Keras logs are reported.
|
|
report_metrics_on: When to report metrics. Must be one of
|
|
the Keras event hooks (less the ``on_``), e.g.
|
|
"train_start" or "predict_end". Defaults to "epoch_end".
|
|
checkpoint_on: When to save checkpoints. Must be one of the Keras event hooks
|
|
(less the ``on_``), e.g. "train_start" or "predict_end". Defaults to
|
|
"epoch_end".
|
|
|
|
"""
|
|
|
|
def _save_and_report_checkpoint(
|
|
self, metrics: Dict, checkpoint: TensorflowCheckpoint
|
|
):
|
|
ray.tune.report(metrics, checkpoint=checkpoint)
|
|
|
|
def _report_metrics(self, metrics: Dict):
|
|
ray.tune.report(metrics)
|