[Tune] Add export_formats option to export policy graphs (#3868)

In earlier PRs, PR#3585 and PR#3637, export_policy_model and export_policy_checkpoint were introduced for users to export TensorFlow model and checkpoint.

For Ray Tune users, these APIs are not accessible through YAML configurations.

In this pull request, export_formats option is provided to enable users to choose the desired export format.
This commit is contained in:
Tianming Xu
2019-02-01 09:07:27 +08:00
committed by Richard Liaw
parent b9eed2e86c
commit 1302fafc0b
11 changed files with 174 additions and 7 deletions
+11
View File
@@ -346,3 +346,14 @@ class RayTrialExecutor(TrialExecutor):
logger.exception("Error restoring runner for Trial %s.", trial)
self.set_status(trial, Trial.ERROR)
return False
def export_trial_if_needed(self, trial):
"""Exports model of this trial based on trial.export_formats.
Return:
A dict that maps ExportFormats to successfully exported models.
"""
if trial.export_formats and len(trial.export_formats) > 0:
return ray.get(
trial.runner.export_model.remote(trial.export_formats))
return {}