fix: ga

examples/configs/librispeech/data.yml.j2

@@ -10,7 +10,7 @@ data_config:
     tfrecords_shards: 32
     shuffle: True
     cache: False
-    buffer_size: 100
+    buffer_size: 1000
     drop_remainder: True
     stage: train
     metadata: {{metadata}}
@@ -24,9 +24,9 @@ data_config:
       - {{datadir}}/dev-other/transcripts.tsv
     tfrecords_dir: {{datadir}}/tfrecords
     tfrecords_shards: 2
-    shuffle: True
+    shuffle: False
     cache: False
-    buffer_size: 100
+    buffer_size: 1000
     drop_remainder: True
     stage: eval
     metadata: {{metadata}}

examples/train.py

@@ -36,11 +36,12 @@ def main(
     jit_compile: bool = False,
     ga_steps: int = None,
     verbose: int = 1,
+    tpu_address: str = None,
     repodir: str = os.path.realpath(os.path.join(os.path.dirname(__file__), "..")),
 ):
     keras.backend.clear_session()
     env_util.setup_seed()
-    strategy = env_util.setup_strategy(devices)
+    strategy = env_util.setup_strategy(devices, tpu_address=tpu_address)
     env_util.setup_mxp(mxp=mxp)
 
     config = Config(config_path, training=True, repodir=repodir, datadir=datadir, modeldir=modeldir)
@@ -68,10 +69,7 @@ def main(
     ga_steps = ga_steps or config.learning_config.ga_steps or 1
 
     train_data_loader = train_dataset.create(train_batch_size, ga_steps=ga_steps, padded_shapes=padded_shapes)
-    if train_dataset.use_ga:
-        logger.info(f"train_data_loader.element_spec = {json.dumps(train_data_loader.element_spec.element_spec, indent=2, default=str)}")
-    else:
-        logger.info(f"train_data_loader.element_spec = {json.dumps(train_data_loader.element_spec, indent=2, default=str)}")
+    logger.info(f"train_data_loader.element_spec = {json.dumps(train_data_loader.element_spec, indent=2, default=str)}")
 
     eval_data_loader = eval_dataset.create(eval_batch_size, padded_shapes=padded_shapes)
     if eval_data_loader:

tensorflow_asr/datasets.py

@@ -64,7 +64,6 @@
 import logging
 import os
 from dataclasses import asdict, dataclass
-from typing import List
 
 import numpy as np
 import tqdm
@@ -386,15 +385,6 @@ def process(
 
         # only apply for training dataset, eval and test dataset should not use GA
         if ga_steps > 1 and self.stage == "train":
-
-            def _key_fn(i, _):
-                return i // ga_steps
-
-            def _reduce_fn(_, ds):
-                elem = ds.map(lambda _, x: x)
-                return tf.data.Dataset.from_tensors(elem)
-
-            dataset = dataset.enumerate().group_by_window(key_func=_key_fn, reduce_func=_reduce_fn, window_size=ga_steps)
             self.use_ga = True
 
         # PREFETCH to improve speed of input length

tensorflow_asr/models/base_model.py

@@ -15,7 +15,6 @@
 
 import importlib
 import logging
-import typing
 
 import numpy as np
 
@@ -152,6 +151,7 @@ def compile(
         if isinstance(ga_steps, int) and ga_steps > 1:
             self.use_ga = True
             self.ga = GradientAccumulator(ga_steps=ga_steps)
+            kwargs["steps_per_execution"] = 1
             logger.info(f"Using gradient accumulation with accumulate steps = {ga_steps}")
         else:
             self.use_ga = False
@@ -203,30 +203,24 @@ def _train_step(self, data: schemas.TrainData):
 
         return gradients
 
-    def train_step(self, data_list: typing.Union[schemas.TrainData, typing.Iterable[schemas.TrainData]]):
-        if not self.use_ga:
-            data = data_list
-            gradients = self._train_step(data)
-        else:
-            iterator = iter(data_list)
-            data = next(iterator)
-            gradients = self._train_step(data)
-
-            for _ in range(1, self.ga.total_steps):
-                try:
-                    data = next(iterator)
-                except StopIteration:
-                    break
-                per_ga_gradients = self._train_step(data)
-                gradients = self.ga.accumulate(gradients, per_ga_gradients)
-
+    def _apply_gradients(self, gradients):
         if self.gradn is not None:
             gradients = self.gradn(step=self.optimizer.iterations, gradients=gradients)
         self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
 
+    def train_step(self, data):
+        gradients = self._train_step(data)
+        self._apply_gradients(gradients)
         metrics = self.get_metrics_result()
         return metrics
 
+    def train_step_ga(self, data, prev_gradients):
+        gradients = self._train_step(data)
+        if prev_gradients is not None:
+            gradients = self.ga.accumulate(prev_gradients, gradients)
+        metrics = self.get_metrics_result()
+        return metrics, gradients
+
     def _test_step(self, data: schemas.TrainData):
         x = data[0]
         y, _ = data_util.set_length(data[1].labels, data[1].labels_length)
@@ -278,6 +272,18 @@ def one_step_on_data(data):
         if not self.run_eagerly:
             one_step_on_data = tf.function(one_step_on_data, reduce_retracing=True, jit_compile=self.jit_compile)
 
+        @tf.autograph.experimental.do_not_convert
+        def one_ga_step_on_data(data, prev_gradients):
+            """Runs a single training step on a batch of data."""
+            outputs, gradients = self.train_step_ga(data, prev_gradients)
+            # Ensure counter is updated only if `train_step` succeeds.
+            with tf.control_dependencies(_minimum_control_deps(outputs)):
+                self._train_counter.assign_add(1)
+            return outputs, gradients
+
+        if not self.run_eagerly:
+            one_ga_step_on_data = tf.function(one_ga_step_on_data, reduce_retracing=True, jit_compile=self.jit_compile)
+
         @tf.autograph.experimental.do_not_convert
         def one_step_on_iterator(iterator):
             """Runs a single training step given a Dataset iterator."""
@@ -292,11 +298,31 @@ def one_step_on_iterator(iterator):
 
         @tf.autograph.experimental.do_not_convert
         def multi_step_on_iterator(iterator):
-            for _ in range(self.steps_per_execution):
-                outputs = one_step_on_iterator(iterator)
-            return outputs
+            for _ in range(self.steps_per_execution.numpy().item()):
+                outputs, data = one_step_on_iterator(iterator)
+            return outputs, data
 
-        if self.steps_per_execution > 1:
+        @tf.autograph.experimental.do_not_convert
+        def ga_step_in_iterator(iterator):
+            data = next(iterator)
+            outputs, gradients = self.distribute_strategy.run(one_ga_step_on_data, args=(data, None))
+            for _ in range(1, self.ga.total_steps):
+                try:
+                    data = next(iterator)
+                    outputs, gradients = self.distribute_strategy.run(one_ga_step_on_data, args=(data, gradients))
+                except StopIteration:
+                    break
+            self.distribute_strategy.run(self._apply_gradients, args=(gradients,))
+            outputs = keras_util.reduce_per_replica(
+                outputs,
+                self.distribute_strategy,
+                reduction=self.distribute_reduction_method,
+            )
+            return outputs, data
+
+        if self.use_ga:
+            train_function = ga_step_in_iterator
+        elif self.steps_per_execution > 1:
             train_function = multi_step_on_iterator
         else:
             train_function = one_step_on_iterator
@@ -347,7 +373,7 @@ def one_step_on_iterator(iterator):
 
         @tf.autograph.experimental.do_not_convert
         def multi_step_on_iterator(iterator):
-            for _ in range(self.steps_per_execution):
+            for _ in range(self.steps_per_execution.numpy().item()):
                 outputs = one_step_on_iterator(iterator)
             return outputs
 

tensorflow_asr/optimizers/accumulation.py

@@ -20,6 +20,9 @@ def __init__(self, ga_steps, name="ga"):
     def total_steps(self):
         return self._ga_steps
 
+    def is_apply_step(self, step):
+        return tf.math.equal(step % self._ga_steps, 0)
+
     def accumulate(self, gradients, per_ga_gradients):
         """Accumulates :obj:`gradients` on the current replica."""
         with tf.name_scope(self.name):

tensorflow_asr/utils/env_util.py

@@ -83,10 +83,7 @@ def setup_devices(
 def setup_tpu(
     tpu_address=None,
 ):
-    if tpu_address is None:
-        resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
-    else:
-        resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="grpc://" + tpu_address)
+    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu_address)
     tf.tpu.experimental.initialize_tpu_system(resolver)
     return tf.distribute.TPUStrategy(resolver)