Repository: cosmoquester/speech-recognition

Issue:
LAS training so slow
- LAS training too slow on GPU
- LAS training TPU not support

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun) Report
> Merging [#36](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun) (159a577) into [master](https://codecov.io/gh/cosmoquester/speech-recognition/commit/64420143977280c151afe08ee5e3687d21ff17c9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun) (6442014) will **increase** coverage by `4.32%`.
> The diff coverage is `75.00%`.

[![Impacted file tree graph](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36/graphs/tree.svg?width=650&height=150&src=pr&token=veHoLRzJum&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun)](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun)

```diff
@@            Coverage Diff             @@
##           master      #36      +/-   ##
==========================================
+ Coverage   52.48%   56.81%   +4.32%     
==========================================
  Files          16       16              
  Lines        1128     1116      -12     
  Branches      119      118       -1     
==========================================
+ Hits          592      634      +42     
+ Misses        518      461      -57     
- Partials       18       21       +3     
```


| [Impacted Files](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun) | Coverage Δ | |
|---|---|---|
| [speech\_recognition/measure.py](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun#diff-c3BlZWNoX3JlY29nbml0aW9uL21lYXN1cmUucHk=) | `43.90% <0.00%> (ø)` | |
| [speech\_recognition/models/las.py](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun#diff-c3BlZWNoX3JlY29nbml0aW9uL21vZGVscy9sYXMucHk=) | `91.72% <80.00%> (+33.76%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun). Last update [6442014...159a577](https://codecov.io/gh/cosmoquester/speech-recognition/pull/36?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ParkSangJun).


Patch:
diff --git a/speech_recognition/measure.py b/speech_recognition/measure.py
index 62515b8..e1fd6ce 100644
--- a/speech_recognition/measure.py
+++ b/speech_recognition/measure.py
@@ -17,7 +17,7 @@ class SparseCategoricalCrossentrophy(tf.keras.losses.Loss):
 
     def call(self, y_true, y_pred):
         loss = tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred, from_logits=self.from_logits)
-        loss = tf.boolean_mask(loss[:, tf.newaxis], y_true != self.ignore_index)
+        loss = tf.boolean_mask(loss, y_true != self.ignore_index)
         return loss
 
 
diff --git a/speech_recognition/models/las.py b/speech_recognition/models/las.py
index b53afd3..3b198a6 100644
--- a/speech_recognition/models/las.py
+++ b/speech_recognition/models/las.py
@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import tensorflow as tf
 from tensorflow.keras.layers import GRU, LSTM, BatchNormalization, Conv2D, Dense, Dropout, Embedding, SimpleRNN
@@ -340,59 +340,39 @@ class LAS(ModelProto):
             rnn_type, vocab_size, decoder_hidden_dim, num_decoder_layers, dropout, pad_id, name="attend_and_speller"
         )
 
-    def call(self, inputs):
-        """
-        `call` method is not used for training. This is only for model build and print summary.
-        Training will be executed with `train_step` function.
-        """
+    def call(self, inputs: Tuple[tf.Tensor, tf.Tensor], training: Optional[bool] = None) -> tf.Tensor:
+        # audio: [BatchSize, TimeStep, DimAudio], decoder_input: [BatchSize, NumTokens]
         audio_input, decoder_input = inputs
-        audio_output, attention_mask, *states = self.listener(audio_input)
-        output, *states = self.attend_and_speller(audio_output, decoder_input[:, 0], attention_mask, states)
 
-    @tf.function
-    def _step(self, inputs, training=None):
-        # audio: [BatchSize, TimeStep, DimAudio]
-        # decoder_input, target: [BatchSize, NumTokens]
-        (audio_input, decoder_input), target = inputs
+        # Use range on TPU because of issue https://github.com/tensorflow/tensorflow/issues/49469
+        if decoder_input.shape[1]:
+            token_length = decoder_input.shape[1]
+            index_iter = range(token_length)
+        else:
+            token_length = tf.shape(decoder_input)[1]
+            index_iter = tf.range(token_length)
 
-        token_length = target.shape[1] or tf.shape(target)[1]
-        audio_output, attention_mask, *states = self.listener(audio_input, training=training)
+        audio_output, attention_mask, *states = self.listener(audio_input)
+        outputs = tf.TensorArray(
+            audio_output.dtype, size=token_length, infer_shape=False, element_shape=[None, self.vocab_size]
+        )
 
-        loss = tf.constant(0.0, dtype=audio_output.dtype)
         use_teacher_forcing = tf.random.uniform((), 0, 1) < self.teacher_forcing_rate
         output = tf.zeros([tf.shape(audio_output)[0], self.vocab_size], dtype=audio_output.dtype)
-        for i in tf.range(token_length):
+        for i in index_iter:
             if use_teacher_forcing or i == 0:
                 decoder_input_t = tf.gather(decoder_input, i, axis=1)
             else:
                 decoder_input_t = tf.argmax(output, axis=-1, output_type=tf.int32)
 
-            output, *states = self.attend_and_speller(
-                audio_output, decoder_input_t, attention_mask, states, training=training
-            )
-
-            loss += tf.reduce_mean(self.compiled_loss(target[:, i], output))
-            self.compiled_metrics.update_state(decoder_input[:, i], output)
-        loss /= tf.cast(token_length, loss.dtype)
-        return loss
-
-    def train_step(self, inputs: Tuple[Tuple[tf.Tensor, tf.Tensor], tf.Tensor]) -> Dict[str, float]:
-        # Use inner tf.function because of https://github.com/tensorflow/tensorflow/issues/42119 tf.range cannot be compiled without tf.function
-        # And https://github.com/tensorflow/models/issues/9972 error on multi device with use tf.function of whole train_step
-        with tf.GradientTape() as tape:
-            loss = self._step(inputs, training=True)
-        variables = self.trainable_variables
-        gradients = tape.gradient(loss, variables)
-        self.optimizer.apply_gradients(zip(gradients, variables))
-
-        return {"loss": loss, **{m.name: m.result() for m in self.metrics}}
+            output, *states = self.attend_and_speller(audio_output, decoder_input_t, attention_mask, states)
+            outputs = outputs.write(i, output)
 
-    def test_step(self, inputs: Tuple[Tuple[tf.Tensor, tf.Tensor], tf.Tensor]) -> Dict[str, float]:
-        loss = self._step(inputs, training=False)
-        return {"loss": loss, **{m.name: m.result() for m in self.metrics}}
+        result = tf.transpose(outputs.stack(), [1, 0, 2])
+        return result
 
     def get_loss_fn(self):
-        return SparseCategoricalCrossentrophy(self.pad_id, reduction=tf.keras.losses.Reduction.NONE)
+        return SparseCategoricalCrossentrophy(self.pad_id)
 
     def get_metrics(self):
         return [SparseCategoricalAccuracy(self.pad_id)]


Test patch:
diff --git a/tests/models/test_las.py b/tests/models/test_las.py
index a2eadca..ea0f1eb 100644
--- a/tests/models/test_las.py
+++ b/tests/models/test_las.py
@@ -78,4 +78,5 @@ def test_las(
     las.compile(loss=las.get_loss_fn(), metrics=las.get_metrics())
     audio = tf.random.normal([batch_size, audio_sequence_length, audio_dim, 3])
     tokens = tf.random.uniform([batch_size, num_tokens], 0, vocab_size, tf.int32)
-    loss = las._step(((audio, tokens[:, :-1]), tokens[:, 1:]))
+    output = las((audio, tokens))
+    tf.debugging.assert_equal(tf.shape(output), [batch_size, num_tokens, vocab_size])