A single step of training.
()
| 120 | captured_scaffold_fn = _CapturedObject() |
| 121 | |
| 122 | def train_step(): |
| 123 | """A single step of training.""" |
| 124 | estimator_spec = self._call_model_fn(features, labels, |
| 125 | model_fn_lib.ModeKeys.TRAIN, params) |
| 126 | |
| 127 | try: |
| 128 | captured_scaffold_fn.capture(estimator_spec.scaffold_fn) |
| 129 | except AttributeError: |
| 130 | captured_scaffold_fn.capture(None) |
| 131 | |
| 132 | # train_step will be run by xla.compile(). xla.compile() only supports |
| 133 | # tensor output while train_op can be either an operation or a tensor. |
| 134 | # Even though xla.compile() automatically adds operation-typed train_op as |
| 135 | # control dependency of other tensor outputs, it doesn't do so for |
| 136 | # tensor-typed train_op. Thus, we need to set it explicitly here. |
| 137 | with ops.control_dependencies([estimator_spec.train_op]): |
| 138 | return array_ops.identity(estimator_spec.loss) |
| 139 | |
| 140 | return train_step, captured_scaffold_fn |
| 141 |
nothing calls this directly
no test coverage detected