A Trainer model with summaries for testing purposes.
| 212 | |
| 213 | |
| 214 | class TestTrainerWithSummaries(standard_runner.StandardTrainer): |
| 215 | """A Trainer model with summaries for testing purposes.""" |
| 216 | |
| 217 | def __init__(self): |
| 218 | self.strategy = tf.distribute.get_strategy() |
| 219 | self.model = create_model() |
| 220 | self.optimizer = tf_keras.optimizers.RMSprop(learning_rate=0.1) |
| 221 | self.global_step = self.optimizer.iterations |
| 222 | self.train_loss = tf_keras.metrics.Mean("train_loss", dtype=tf.float32) |
| 223 | train_dataset = self.strategy.distribute_datasets_from_function(dataset_fn) |
| 224 | standard_runner.StandardTrainer.__init__( |
| 225 | self, |
| 226 | train_dataset, |
| 227 | options=standard_runner.StandardTrainerOptions( |
| 228 | use_tpu_summary_optimization=True)) |
| 229 | |
| 230 | def build_train_dataset(self): |
| 231 | return self.strategy.distribute_datasets_from_function(dataset_fn) |
| 232 | |
| 233 | def train_step(self, iterator): |
| 234 | |
| 235 | def _replicated_step(inputs): |
| 236 | """Replicated training step.""" |
| 237 | inputs, targets = inputs |
| 238 | with tf.GradientTape() as tape: |
| 239 | outputs = self.model(inputs) |
| 240 | loss = tf.reduce_mean(tf_keras.losses.MSE(targets, outputs)) |
| 241 | tf.summary.scalar("loss", loss) |
| 242 | grads = tape.gradient(loss, self.model.variables) |
| 243 | self.optimizer.apply_gradients(zip(grads, self.model.variables)) |
| 244 | self.train_loss.update_state(loss) |
| 245 | |
| 246 | self.strategy.run(_replicated_step, args=(next(iterator),)) |
| 247 | |
| 248 | |
| 249 | class ControllerTest(tf.test.TestCase, parameterized.TestCase): |
no outgoing calls