Added a TF-Hub-compatible SavedModel exporter for VGGish. (#9081)

* Added a TF-Hub-compatible SavedModel exporter for VGGish.

* Responded to DAn's comments, switched linspace to arange everywhere.

research/audioset/vggish/vggish_export_tfhub.py

+"""Exports VGGish as a SavedModel for publication to TF Hub.
+
+The exported SavedModel accepts a 1-d float32 Tensor of arbitrary shape
+containing an audio waveform (assumed to be mono 16 kHz samples in the [-1, +1]
+range) and returns a 2-d float32 batch of 128-d VGGish embeddings, one per
+0.96s example generated from the waveform.
+
+Usage:
+  export_tfhub.py <path/to/VGGish/checkpoint> <path/to/tfhub/export>
+"""
+
+import sys
+sys.path.append('..')  # Lets us import yamnet modules from sibling directory.
+
+import numpy as np
+import resampy
+import tensorflow as tf
+assert tf.version.VERSION >= '2.0.0', (
+    'Need at least TF 2.0, you have TF v{}'.format(tf.version.VERSION))
+import tensorflow_hub as tfhub
+
+import vggish_input
+import vggish_params
+import vggish_slim
+from yamnet import features as yamnet_features
+from yamnet import params as yamnet_params
+
+
+def vggish_definer(variables, checkpoint_path):
+  """Defines VGGish with variables tracked and initialized from a checkpoint."""
+  reader = tf.compat.v1.train.NewCheckpointReader(checkpoint_path)
+
+  def var_tracker(next_creator, **kwargs):
+    """Variable creation hook that assigns initial values from a checkpoint."""
+    var_name = kwargs['name']
+    var_value = reader.get_tensor(var_name)
+    kwargs.update({'initial_value': var_value})
+    var = next_creator(**kwargs)
+    variables.append(var)
+    return var
+
+  def waveform_to_features(waveform):
+    """Creates VGGish features using the YAMNet feature extractor."""
+    yamnet_params.SAMPLE_RATE = vggish_params.SAMPLE_RATE
+    yamnet_params.STFT_WINDOW_SECONDS = vggish_params.STFT_WINDOW_LENGTH_SECONDS
+    yamnet_params.STFT_HOP_SECONDS = vggish_params.STFT_HOP_LENGTH_SECONDS
+    yamnet_params.MEL_BANDS = vggish_params.NUM_MEL_BINS
+    yamnet_params.MEL_MIN_HZ = vggish_params.MEL_MIN_HZ
+    yamnet_params.MEL_MAX_HZ = vggish_params.MEL_MAX_HZ
+    yamnet_params.LOG_OFFSET = vggish_params.LOG_OFFSET
+    yamnet_params.PATCH_WINDOW_SECONDS = vggish_params.EXAMPLE_WINDOW_SECONDS
+    yamnet_params.PATCH_HOP_SECONDS = vggish_params.EXAMPLE_HOP_SECONDS
+    log_mel_spectrogram = yamnet_features.waveform_to_log_mel_spectrogram(
+        waveform, yamnet_params)
+    return yamnet_features.spectrogram_to_patches(
+        log_mel_spectrogram, yamnet_params)
+
+  def define_vggish(waveform):
+    with tf.variable_creator_scope(var_tracker):
+      features = waveform_to_features(waveform)
+      return vggish_slim.define_vggish_slim(features, training=False)
+
+  return define_vggish
+
+
+class VGGish(tf.Module):
+  """A TF2 Module wrapper around VGGish."""
+  def __init__(self, checkpoint_path):
+    super().__init__()
+    self._variables = []
+    self._vggish_fn = tf.compat.v1.wrap_function(
+        vggish_definer(self._variables, checkpoint_path),
+        signature=(tf.TensorSpec(shape=[None], dtype=tf.float32),))
+
+  @tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.float32),))
+  def __call__(self, waveform):
+    return self._vggish_fn(waveform)
+
+
+def check_model(model_fn):
+  """Applies vggish_smoke_test's sanity check to an instance of VGGish."""
+  num_secs = 3
+  freq = 1000
+  sr = 44100
+  t = np.arange(0, num_secs, 1 / sr)
+  waveform = np.sin(2 * np.pi * freq * t)
+
+  waveform = resampy.resample(waveform, sr, vggish_params.SAMPLE_RATE)
+  embeddings = model_fn(waveform)
+
+  expected_embedding_mean = -0.0333
+  expected_embedding_std = 0.380
+  rel_error = 0.1
+  np.testing.assert_allclose(
+      [np.mean(embeddings), np.std(embeddings)],
+      [expected_embedding_mean, expected_embedding_std],
+      rtol=rel_error)
+
+
+def main(args):
+  # Create a TF2 wrapper around VGGish.
+  vggish_checkpoint_path = args[0]
+  vggish = VGGish(vggish_checkpoint_path)
+  check_model(vggish)
+
+  # Make TF-Hub export.
+  vggish_tfhub_export_path = args[1]
+  tf.saved_model.save(vggish, vggish_tfhub_export_path)
+
+  # Check export in TF2.
+  model = tfhub.load(vggish_tfhub_export_path)
+  check_model(model)
+
+  # Check export in TF1.
+  with tf.compat.v1.Graph().as_default(), tf.compat.v1.Session() as sess:
+    model = tfhub.load(vggish_tfhub_export_path)
+    sess.run(tf.compat.v1.global_variables_initializer())
+    def run_model(waveform):
+      embeddings = model(waveform)
+      return sess.run(embeddings)
+    check_model(run_model)
+
+if __name__ == '__main__':
+  main(sys.argv[1:])

research/audioset/vggish/vggish_inference_demo.py

@@ -88,7 +88,7 @@ def main(_):
     num_secs = 5
     freq = 1000
     sr = 44100
-    t = np.linspace(0, num_secs, int(num_secs * sr))
+    t = np.arange(0, num_secs, 1 / sr)
     x = np.sin(2 * np.pi * freq * t)
     # Convert to signed 16-bit samples.
     samples = np.clip(x * 32768, -32768, 32767).astype(np.int16)

research/audioset/vggish/vggish_smoke_test.py

@@ -53,7 +53,7 @@ rel_error = 0.1  # Up to 10%
 num_secs = 3
 freq = 1000
 sr = 44100
-t = np.linspace(0, num_secs, int(num_secs * sr))
+t = np.arange(0, num_secs, 1 / sr)
 x = np.sin(2 * np.pi * freq * t)
 
 # Produce a batch of log mel spectrogram examples.

research/audioset/vggish/vggish_train_demo.py

@@ -94,7 +94,7 @@ def _get_examples_batch():
   # Make a waveform for each class.
   num_seconds = 5
   sr = 44100  # Sampling rate.
-  t = np.linspace(0, num_seconds, int(num_seconds * sr))  # Time axis.
+  t = np.arange(0, num_seconds, 1 / sr)  # Time axis
   # Random sine wave.
   freq = np.random.uniform(100, 1000)
   sine = np.sin(2 * np.pi * freq * t)