(self)
| 95 | """Mock AsyncTrainer to test the _AsyncTrainer class.""" |
| 96 | |
| 97 | def __init__(self): |
| 98 | self._strategy = tf.distribute.get_strategy() |
| 99 | self.init_async() |
| 100 | |
| 101 | self.global_step = tf.Variable( |
| 102 | 0, |
| 103 | dtype=tf.int64, |
| 104 | name='global_step', |
| 105 | trainable=False, |
| 106 | aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA) |
| 107 | self.eval_global_step = tf.Variable( |
| 108 | 0, |
| 109 | dtype=tf.int64, |
| 110 | name='eval_global_step', |
| 111 | trainable=False, |
| 112 | aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA) |
| 113 | |
| 114 | train_dataset = self.distribute_dataset(dataset_fn) |
| 115 | orbit.StandardTrainer.__init__( |
| 116 | self, train_dataset, options=orbit.StandardTrainerOptions()) |
| 117 | |
| 118 | validation_dataset = self.distribute_dataset(dataset_fn) |
| 119 | orbit.StandardEvaluator.__init__( |
| 120 | self, |
| 121 | validation_dataset, |
| 122 | options=orbit.StandardEvaluatorOptions(use_tf_while_loop=True)) |
| 123 | |
| 124 | def train_loop_begin(self): |
| 125 | self.global_step.assign(0) |
nothing calls this directly
no test coverage detected