⚡ Supported Gradients Accumulation

Author: nglehuy

README.md

@@ -19,6 +19,7 @@ TensorFlowASR implements some automatic speech recognition architectures such as
 
 ## What's New?
 
+- (11/14/2020) Supported Gradient Accumulation for Training in Larger Batch Size
 - (11/3/2020) Reduce differences between `librosa.stft` and `tf.signal.stft`
 - (10/31/2020) Update DeepSpeech2 and Supported Jasper [https://arxiv.org/abs/1904.03288](https://arxiv.org/abs/1904.03288)
 - (10/18/2020) Supported Streaming Transducer [https://arxiv.org/abs/1811.06621](https://arxiv.org/abs/1811.06621)

examples/conformer/train_ga_conformer.py

@@ -41,6 +41,9 @@
 parser.add_argument("--ebs", type=int, default=None,
                     help="Evaluation batch size per replica")
 
+parser.add_argument("--acs", type=int, default=None,
+                    help="Train accumulation steps")
+
 parser.add_argument("--devices", type=int, nargs="*", default=[0],
                     help="Devices' ids to apply distributed training")
 
@@ -125,4 +128,5 @@
 conformer_trainer.compile(model=conformer, optimizer=optimizer,
                           max_to_keep=args.max_ckpts)
 
-conformer_trainer.fit(train_dataset, eval_dataset, train_bs=args.tbs, eval_bs=args.ebs)
+conformer_trainer.fit(train_dataset, eval_dataset,
+                      train_bs=args.tbs, eval_bs=args.ebs, train_acs=args.acs)

examples/conformer/train_ga_subword_conformer.py

@@ -41,6 +41,9 @@
 parser.add_argument("--ebs", type=int, default=None,
                     help="Evaluation batch size per replica")
 
+parser.add_argument("--acs", type=int, default=None,
+                    help="Train accumulation steps")
+
 parser.add_argument("--devices", type=int, nargs="*", default=[0],
                     help="Devices' ids to apply distributed training")
 
@@ -141,4 +144,5 @@
 conformer_trainer.compile(model=conformer, optimizer=optimizer,
                           max_to_keep=args.max_ckpts)
 
-conformer_trainer.fit(train_dataset, eval_dataset, train_bs=args.tbs, eval_bs=args.ebs)
+conformer_trainer.fit(train_dataset, eval_dataset,
+                      train_bs=args.tbs, eval_bs=args.ebs, train_acs=args.acs)

examples/streaming_transducer/train_ga_streaming_transducer.py

@@ -40,6 +40,9 @@
 parser.add_argument("--ebs", type=int, default=None,
                     help="Evaluation batch size per replica")
 
+parser.add_argument("--acs", type=int, default=None,
+                    help="Train accumulation steps")
+
 parser.add_argument("--devices", type=int, nargs="*", default=[0],
                     help="Devices' ids to apply distributed training")
 
@@ -116,4 +119,5 @@
 streaming_transducer_trainer.compile(model=streaming_transducer, optimizer=optimizer,
                                      max_to_keep=args.max_ckpts)
 
-streaming_transducer_trainer.fit(train_dataset, eval_dataset, train_bs=args.tbs, eval_bs=args.ebs)
+streaming_transducer_trainer.fit(train_dataset, eval_dataset,
+                                 train_bs=args.tbs, eval_bs=args.ebs, train_acs=args.acs)

examples/streaming_transducer/train_ga_subword_streaming_transducer.py

@@ -40,6 +40,9 @@
 parser.add_argument("--ebs", type=int, default=None,
                     help="Evaluation batch size per replica")
 
+parser.add_argument("--acs", type=int, default=None,
+                    help="Train accumulation steps")
+
 parser.add_argument("--devices", type=int, nargs="*", default=[0],
                     help="Devices' ids to apply distributed training")
 
@@ -132,4 +135,5 @@
 streaming_transducer_trainer.compile(model=streaming_transducer, optimizer=optimizer,
                                      max_to_keep=args.max_ckpts)
 
-streaming_transducer_trainer.fit(train_dataset, eval_dataset, train_bs=args.tbs, eval_bs=args.ebs)
+streaming_transducer_trainer.fit(train_dataset, eval_dataset,
+                                 train_bs=args.tbs, eval_bs=args.ebs, train_acs=args.acs)

setup.py

@@ -18,7 +18,7 @@
     long_description = fh.read()
 
 requirements = [
-    # "tensorflow>=2.3.0",
+    "tensorflow>=2.3.0",
     "tensorflow-datasets>=3.2.1,<4.0.0",
     "tensorflow-addons>=0.10.0",
     "setuptools>=47.1.1",
@@ -38,7 +38,7 @@
 
 setuptools.setup(
     name="TensorFlowASR",
-    version="0.2.10",
+    version="0.3.0",
     author="Huy Le Nguyen",
     author_email="[email protected]",
     description="Almost State-of-the-art Automatic Speech Recognition using Tensorflow 2",

tensorflow_asr/optimizers/accumulation.py

@@ -20,12 +20,14 @@ def __init__(self, trainable_variables):
         self.gradients = [
             tf.Variable(
                 tf.zeros_like(g),
+                trainable=False,
                 synchronization=tf.VariableSynchronization.ON_READ
             ) for g in trainable_variables
         ]
 
     def reset(self):
-        for g in self.gradients: g.assign(tf.zeros_like(g))
+        for i, g in enumerate(self.gradients):
+            self.gradients[i].assign(tf.zeros_like(g))
 
     def accumulate(self, step_gradients):
         for i, g in enumerate(step_gradients):

tensorflow_asr/runners/__init__.py

@@ -28,8 +28,8 @@ def save_from_checkpoint(func,
         max_to_keep: number of checkpoints to keep
         **kwargs: contains built models, optimizers
     """
-    steps = tf.Variable(0, dtype=tf.int64)  # Step must be int64
-    epochs = tf.Variable(1)
+    steps = tf.Variable(0, trainable=False, dtype=tf.int64)  # Step must be int64
+    epochs = tf.Variable(1, trainable=False)
     checkpoint_dir = os.path.join(outdir, "checkpoints")
     if not os.path.exists(checkpoint_dir):
         raise ValueError(f"checkpoint directory not found: {checkpoint_dir}")

tensorflow_asr/runners/base_runners.py

@@ -72,7 +72,8 @@ def __init__(self,
         super(BaseTrainer, self).__init__(config)
         self.set_strategy(strategy)
         # Steps and Epochs start from 0
-        self.steps = tf.Variable(0, dtype=tf.int64)  # Step must be int64 to use tf.summary
+        # Step must be int64 to use tf.summary
+        self.steps = tf.Variable(0, trainable=False, dtype=tf.int64)
         self.train_steps_per_epoch = None
         self.eval_steps_per_epoch = None
         # Dataset
@@ -120,13 +121,14 @@ def set_train_data_loader(self, train_dataset, train_bs=None, train_acs=None):
         self.config.batch_size = train_bs  # Update batch size fed from arguments
 
         if not train_acs: train_acs = self.config.accumulation_steps
-        assert train_bs % train_acs == 0, "Batch size must be a multiple of Accumulation Steps"
-        self.accumulation_bs = train_bs // train_acs
         self.config.accumulation_steps = train_acs  # update accum steps fed from arguments
 
         self.train_data = train_dataset.create(self.global_batch_size)
         self.train_data_loader = self.strategy.experimental_distribute_dataset(self.train_data)
-        self.train_steps_per_epoch = train_dataset.total_steps
+        if hasattr(self, "accumulation"):
+            self.train_steps_per_epoch = train_dataset.total_steps // self.config.accumulation_steps
+        else:
+            self.train_steps_per_epoch = train_dataset.total_steps
 
     def set_eval_data_loader(self, eval_dataset, eval_bs=None):
         """ Set eval data loader (MUST).

tensorflow_asr/runners/transducer_runners.py

@@ -90,48 +90,39 @@ def compile(self,
 class TransducerTrainerGA(TransducerTrainer):
     """ Transducer Trainer that uses Gradients Accumulation """
 
-    @tf.function(experimental_relax_shapes=True)
-    def _train_step(self, batch):
-        _, bfeatures, binput_length, blabels, blabel_length, bpred_inp = batch
-
+    @tf.function
+    def _train_function(self, iterator):
+        for _ in range(self.config.accumulation_steps):
+            batch = next(iterator)
+            self.strategy.run(self._train_step, args=(batch,))
+        self.strategy.run(self._apply_gradients, args=())
+
+    @tf.function
+    def _apply_gradients(self):
+        self.optimizer.apply_gradients(
+            zip(self.accumulation.gradients, self.model.trainable_variables))
         self.accumulation.reset()
 
-        for accum_step in range(self.config.accumulation_steps):
+    @tf.function(experimental_relax_shapes=True)
+    def _train_step(self, batch):
+        _, features, input_length, labels, label_length, pred_inp = batch
 
-            indices = tf.expand_dims(
-                tf.range(
-                    accum_step * self.accumulation_bs,
-                    (accum_step + 1) * self.accumulation_bs,
-                    dtype=tf.int32
-                ),
-                axis=-1
+        with tf.GradientTape() as tape:
+            logits = self.model([features, pred_inp], training=True)
+            tape.watch(logits)
+            per_train_loss = rnnt_loss(
+                logits=logits, labels=labels, label_length=label_length,
+                logit_length=(input_length // self.model.time_reduction_factor),
+                blank=self.text_featurizer.blank
+            )
+            train_loss = tf.nn.compute_average_loss(
+                per_train_loss,
+                global_batch_size=self.global_batch_size
             )
 
-            features = tf.gather_nd(bfeatures, indices)
-            input_length = tf.gather_nd(binput_length, indices)
-            labels = tf.gather_nd(blabels, indices)
-            label_length = tf.gather_nd(blabel_length, indices)
-            pred_inp = tf.gather_nd(bpred_inp, indices)
-
-            with tf.GradientTape() as tape:
-                logits = self.model([features, pred_inp], training=True)
-                tape.watch(logits)
-                per_train_loss = rnnt_loss(
-                    logits=logits, labels=labels, label_length=label_length,
-                    logit_length=(input_length // self.model.time_reduction_factor),
-                    blank=self.text_featurizer.blank
-                )
-                train_loss = tf.nn.compute_average_loss(
-                    per_train_loss,
-                    global_batch_size=self.global_batch_size
-                )
-
-            step_gradients = tape.gradient(train_loss, self.model.trainable_variables)
-            self.accumulation.accumulate(step_gradients)
-            self.train_metrics["transducer_loss"].update_state(per_train_loss)
-
-        self.optimizer.apply_gradients(
-            zip(self.accumulation.gradients, self.model.trainable_variables))
+        gradients = tape.gradient(train_loss, self.model.trainable_variables)
+        self.accumulation.accumulate(gradients)
+        self.train_metrics["transducer_loss"].update_state(per_train_loss)
 
     def compile(self,
                 model: Transducer,