r'''Decorator to set params in a model function.
(model_fn, model_dir, config)
| 755 | gfile.Copy(src, target) |
| 756 | |
| 757 | def wraps_model_fn(model_fn, model_dir, config): |
| 758 | r'''Decorator to set params in a model function. |
| 759 | ''' |
| 760 | def wrapped_model_fn(features, labels, mode, params): |
| 761 | r'''Wrapped model function. |
| 762 | ''' |
| 763 | with scope(): |
| 764 | estimator_spec = model_fn(features, labels, mode, params) |
| 765 | if estimator_spec.scaffold.saver is None: |
| 766 | estimator_spec.scaffold._saver = Saver( # pylint: disable=protected-access |
| 767 | max_to_keep=config.keep_checkpoint_max, |
| 768 | keep_checkpoint_every_n_hours=config.keep_checkpoint_every_n_hours, |
| 769 | defer_build=True, |
| 770 | save_relative_paths=True) |
| 771 | training_hooks = list(estimator_spec.training_hooks) or [] |
| 772 | training_chief_hooks = list( |
| 773 | estimator_spec.training_chief_hooks) or [] |
| 774 | estimator_spec = estimator_spec._replace( # pylint: disable=protected-access |
| 775 | training_hooks=training_hooks, |
| 776 | training_chief_hooks=training_chief_hooks) |
| 777 | return estimator_spec |
| 778 | return wrapped_model_fn |
| 779 | |
| 780 | |
| 781 | def start_std_server(config): |