Spectrum augmentation: Avoid creating TF global step variable statically.

PiperOrigin-RevId: 737138292

Author: agutkin

protoscribe/corpus/reader/tasks.py

@@ -28,6 +28,7 @@
 from protoscribe.glyphs import glyph_vocab as glyph_lib
 from protoscribe.sketches.utils import stroke_tokenizer as tokenizer_lib
 from protoscribe.speech import audio_tokenizer
+from protoscribe.speech import augmentation
 import seqio
 import tensorflow as tf
 
@@ -454,6 +455,10 @@ def register(
   sketch_stroke_stats = ds_lib.get_sketch_stroke_stats(config)
   stroke_tokenizer = ds_lib.get_stroke_tokenizer(config)
 
+  if is_training and speech_spectrum_augmentation:
+    # Initialize spectrum augmentation.
+    augmentation.tf_spec_augment_init()
+
   speech_tokenizer = None
   if speech_tokenizer_name_or_path:
     speech_tokenizer = audio_tokenizer.get_tokenizer(

protoscribe/speech/augmentation.py

@@ -32,6 +32,14 @@ def _default_spec_augment_config() -> ml_collections.FrozenConfigDict:
   })
 
 
+def tf_spec_augment_init() -> None:
+  """Global initialization for spectrum augmenter.
+
+  Should be called once ideally to create a global step variable.
+  """
+  impl_lib.tf_spec_augment_init()
+
+
 def tf_spec_augment(
     spectrum: tf.Tensor,
     config: ml_collections.FrozenConfigDict | None = None,

protoscribe/speech/augmentation_lingvo.py

@@ -28,9 +28,6 @@
 from protoscribe.speech import tf_utils
 import tensorflow as tf
 
-# This creates global step variable.
-tf_utils.get_or_create_global_step_var()
-
 
 @dataclasses.dataclass
 class AugmenterConfig:
@@ -1028,6 +1025,14 @@ def _augmentation_network(
   return inputs
 
 
+def tf_spec_augment_init() -> None:
+  """Global initialization for spectrum augmenter.
+
+  Should be called once ideally to create a global step variable.
+  """
+  tf_utils.get_or_create_global_step_var()
+
+
 def tf_spec_augment_lingvo(
     config: AugmenterConfig,
     inputs: tf.Tensor,

protoscribe/speech/augmentation_lingvo_test.py

@@ -21,6 +21,10 @@
 
 class AugmentationTest(tf.test.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    lib.tf_spec_augment_init()
+
   def test_with_time_mask(self):
     tf.random.set_seed(127)
     batch_size = 5

protoscribe/speech/augmentation_test.py

@@ -25,6 +25,10 @@
 
 class AugmentationTest(tf.test.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    lib.tf_spec_augment_init()
+
   def test_with_default_config(self):
     input_shape = [_NUM_TIME_STEPS, _NUM_FREQ_BINS]
     inputs = tf.ones(input_shape, dtype=tf.float32)