Trains model on `dataset` using `optimizer`.
(model, optimizer, dataset, step_counter, log_interval=None,
num_steps=None)
| 138 | # TODO(brianklee): Enable @tf.function on the training loop when zip, enumerate |
| 139 | # are supported by autograph. |
| 140 | def train(model, optimizer, dataset, step_counter, log_interval=None, |
| 141 | num_steps=None): |
| 142 | """Trains model on `dataset` using `optimizer`.""" |
| 143 | start = time.time() |
| 144 | for (batch, (images, labels)) in enumerate(dataset): |
| 145 | if num_steps is not None and batch > num_steps: |
| 146 | break |
| 147 | with tf.contrib.summary.record_summaries_every_n_global_steps( |
| 148 | 10, global_step=step_counter): |
| 149 | # Record the operations used to compute the loss given the input, |
| 150 | # so that the gradient of the loss with respect to the variables |
| 151 | # can be computed. |
| 152 | with tf.GradientTape() as tape: |
| 153 | logits = model(images, training=True) |
| 154 | loss_value = loss(logits, labels) |
| 155 | tf.contrib.summary.scalar('loss', loss_value) |
| 156 | tf.contrib.summary.scalar('accuracy', compute_accuracy(logits, labels)) |
| 157 | grads = tape.gradient(loss_value, model.variables) |
| 158 | optimizer.apply_gradients( |
| 159 | zip(grads, model.variables), global_step=step_counter) |
| 160 | if log_interval and batch % log_interval == 0: |
| 161 | rate = log_interval / (time.time() - start) |
| 162 | print('Step #%d\tLoss: %.6f (%d steps/sec)' % (batch, loss_value, rate)) |
| 163 | start = time.time() |
| 164 | |
| 165 | |
| 166 | def test(model, dataset): |
no test coverage detected