8648d437573e2e41b92b8fea13b4afe5e01851e8,baseline/tf/classify/training/eager.py,ClassifyTrainerEagerTf,_train,#ClassifyTrainerEagerTf#Any#Any#,76

Before Change

        for features, y in pg(loader):
            lossv = self.optimizer.update(self.model, features, y).numpy()
            batchsz = int(y.shape[0])
            report_lossv = lossv * batchsz
            epoch_loss += report_lossv
            epoch_div += batchsz

After Change

        epoch_loss = tf.keras.metrics.Sum()
        epoch_div = tf.keras.metrics.Sum()
        nstep_loss = tf.keras.metrics.Sum()
        nstep_div = tf.keras.metrics.Sum()
        self.nstep_start = time.time()

        @tf.function
        def _train_step(inputs):
            Replicated training step.

            features, y = inputs
            loss = self.optimizer.update(self.model, features, y)
            batchsz = tf.cast(get_shape_as_list(y)[0], tf.float32)
            report_loss = loss * batchsz
            epoch_loss.update_state(report_loss)
            nstep_loss.update_state(report_loss)
            epoch_div.update_state(batchsz)
            nstep_div.update_state(batchsz)

        for inputs in pg(loader):
            _train_step(inputs)
            step = self.optimizer.global_step.numpy() + 1

            if step % self.nsteps == 0:
                metrics = self.calc_metrics(nstep_loss.result().numpy(), nstep_div.result().numpy())
                self.report(
                    step, metrics, self.nstep_start,
                    "Train", "STEP", reporting_fns, self.nsteps
                )

Instances