✍️ tensorflow-io does not support tpu yet

Author: nglehuy

examples/conformer/train_tpu_keras_subword_conformer.py

@@ -125,5 +125,4 @@
     conformer.fit(
         train_data_loader, epochs=config.learning_config.running_config.num_epochs,
         validation_data=eval_data_loader, callbacks=callbacks,
-        steps_per_epoch=train_dataset.total_steps
     )

tensorflow_asr/featurizers/speech_featurizers.py

@@ -21,9 +21,11 @@
 import tensorflow as tf
 import tensorflow_io as tfio
 
-from ..utils.utils import log10
+from ..utils.utils import log10, has_tpu
 from .gammatone import fft_weights
 
+tpu = has_tpu()
+
 
 def load_and_convert_to_wav(path: str) -> tf.Tensor:
     wave, rate = librosa.load(os.path.expanduser(path), sr=None, mono=True)
@@ -48,8 +50,10 @@ def read_raw_audio(audio, sample_rate=16000):
 
 def tf_read_raw_audio(audio: tf.Tensor, sample_rate=16000):
     wave, rate = tf.audio.decode_wav(audio, desired_channels=1, desired_samples=-1)
-    resampled = tfio.audio.resample(wave, rate_in=tf.cast(rate, dtype=tf.int64), rate_out=sample_rate)
-    return tf.reshape(resampled, shape=[-1])  # reshape for using tf.signal
+    if not tpu:
+        resampled = tfio.audio.resample(wave, rate_in=tf.cast(rate, dtype=tf.int64), rate_out=sample_rate)
+        return tf.reshape(resampled, shape=[-1])  # reshape for using tf.signal
+    return tf.reshape(wave, shape=[-1])  # reshape for using tf.signal
 
 
 def slice_signal(signal, window_size, stride=0.5) -> np.ndarray:

tensorflow_asr/utils/utils.py

@@ -169,6 +169,12 @@ def has_gpu_or_tpu():
     return True
 
 
+def has_tpu():
+    tpus = tf.config.list_logical_devices("TPU")
+    if len(tpus) == 0: return False
+    return True
+
+
 def find_max_length_prediction_tfarray(tfarray: tf.TensorArray) -> tf.Tensor:
     with tf.name_scope("find_max_length_prediction_tfarray"):
         index = tf.constant(0, dtype=tf.int32)