fix: ga step

Author: nglehuy

tensorflow_asr/models/base_model.py

@@ -184,14 +184,19 @@ def train_step(self, data):
         metrics = self.get_metrics_result()
         return metrics
 
-    def train_step_ga(self, data_buffer):
-        first_data, *rest_data = data_buffer
-        gradients = self._train_step(first_data)
-        for data in rest_data:
-            next_gradients = self._train_step(data)
-            gradients = self.ga.accumulate(gradients, next_gradients)
-        self._apply_gradients(gradients)
+    def train_step_ga(self, data):
+        gradients = self._train_step(data)
         metrics = self.get_metrics_result()
+        return metrics, gradients
+
+    def train_step_ga_next(self, data, prev_gradients):
+        metrics, gradients = self.train_step_ga(data)
+        gradients = self.ga.accumulate(prev_gradients, gradients)
+        return metrics, gradients
+
+    def train_step_ga_last(self, data, prev_gradients):
+        metrics, gradients = self.train_step_ga_next(data, prev_gradients)
+        self._apply_gradients(gradients)
         return metrics
 
     def _test_step(self, data: schemas.TrainData):
@@ -244,29 +249,17 @@ def make_train_function(self, force=False):
             return self.train_function
 
         @tf.autograph.experimental.do_not_convert
-        def one_ga_step_on_data(data_buffer):
-            outputs = self.distribute_strategy.run(self.train_step_ga, args=(data_buffer,))
+        def one_ga_step_on_data(iterator):
+            outputs, gradients = self.distribute_strategy.run(self.train_step_ga, args=(next(iterator),))
+            for i, data in zip(range(1, self.ga.total_steps - 1), iterator):
+                outputs, gradients = self.distribute_strategy.run(self.train_step_ga_next, args=(data, gradients))
+            outputs = self.distribute_strategy.run(self.train_step_ga_last, args=(next(iterator), gradients))
             outputs = reduce_per_replica(
                 outputs,
                 self.distribute_strategy,
                 reduction="auto",
             )
             return outputs
-            # data = next(iterator)
-            # outputs, gradients = self.distribute_strategy.run(self.train_step_ga, args=(data, None))
-            # for _ in range(1, self.ga.total_steps):
-            #     try:
-            #         data = next(iterator)
-            #         outputs, gradients = self.distribute_strategy.run(self.train_step_ga, args=(data, gradients))
-            #     except StopIteration:
-            #         break
-            # self.distribute_strategy.run(self._apply_gradients, args=(gradients,))
-            # outputs = reduce_per_replica(
-            #     outputs,
-            #     self.distribute_strategy,
-            #     reduction="auto",
-            # )
-            # return outputs
 
         if not self.run_eagerly:
             one_ga_step_on_data = tf.function(
@@ -276,10 +269,8 @@ def one_ga_step_on_data(data_buffer):
             )
 
         def function(iterator):
-            data_buffer = []
-            for _, data in zip(range(self.ga.total_steps), iterator):
-                data_buffer.append(data)
-            return one_ga_step_on_data(data_buffer)
+            outputs = one_ga_step_on_data(iterator)
+            return outputs
 
         self.train_function = function
         return self.train_function