{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Tensorflow Dataset\n\nThis notebook demonstrates processing training data 'on the fly' using a Tensorflow Dataset. This is an efficient way to handle data: \n\n* Files are read and examples are prepared as they are needed by the GPU/TPU device, with buffers to avoid data starvation.\n* Experimenting with different augmentation techniques just involves code changes, rather than reprocessing the dataset.\n* Augmented datasets can be effectively infinite.\n\nA couple caveats are in order, though:\n\n* This implementation processes arbitrary audio using tf.python_function. This is fairly inefficient, due to the Python global interpreter lock. Reprocessing the data to wav files will allow greater efficiency.\n* For the same reason, it's best to implement augmentations, etc, using Tensorflow and numpy. Librosa in particular may require using a tf.python_function, and thus a drop in efficiency."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport tensorflow as tf\nimport tensorflow_datasets.public_api as tfds\nfrom matplotlib import pyplot as plt\nimport soundfile as sf\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\n\nimport io\nimport os\nimport time\n\ncount = 0\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        count += 1\nprint('counted %d files.' % count)\n\nclass BirbsongDatasetBuilder(object):\n    def __init__(self, batch_size, window_size_s=5):\n        self.batch_size = batch_size\n        self.window_size_s = window_size_s\n        self.sample_rate = 32000\n        self.features_rate = 100\n        self.train_file_pattern = '/kaggle/input/birdclef-2021/train_short_audio/*/*.ogg'\n\n    def select_window(self, audio_tensor):\n        # In this example, we just select the first window_size_s.\n        # In practice, you can select a random segment during training, or\n        # use a heuristic to find an interesting segment.\n        return audio_tensor[:self.window_size_s * self.sample_rate]\n\n    def get_species_enum_table(self):\n        # Create a static enum for the species set.\n        species_list = sorted(os.listdir('/kaggle/input/birdclef-2021/train_short_audio'))\n        species_table = tf.lookup.StaticHashTable(\n            tf.lookup.KeyValueTensorInitializer(tf.constant(species_list),\n                                                tf.constant(range(len(species_list)))),\n            default_value=-1)\n        return species_table\n    \n    def augment_time_domain(self, audio, is_train):\n        if not is_train:\n            return audio\n        # Do interesting things here.\n        return audio\n        \n    def augment_features(self, features, is_train):\n        if not is_train:\n            return features\n        # Do interesting things here.\n        return features\n        \n    def process_examples(self, ex0, ex1=None, is_train=True):\n        if ex1 is not None:\n            # Here we combine two examples.\n            labels = tf.stack([ex0['label'], ex1['label']],\n                                   axis=1)\n            labels_enum = tf.stack([self.species_enum_table.lookup(ex0['label']),\n                                    self.species_enum_table.lookup(ex1['label'])],\n                                   axis=1)\n            \n            # This applies a random gain to each member of the batch separately.\n            gain0 = tf.random.uniform([self.batch_size, 1, 1], 0.2, 0.5)\n            merged_audio = gain0 * ex0['audio'] + (1 - gain0) * ex1['audio']\n        else:\n            labels = ex0['label'][:, tf.newaxis]\n            labels_enum = self.species_enum_table.lookup(ex0['label'])[:, tf.newaxis]\n            merged_audio = ex0['audio']\n\n        # This is a good place to apply any extra augmentations to the time\n        # domain audio.\n        merged_audio = self.augment_time_domain(merged_audio, is_train)\n\n        # You can easily replace feature extraction with another representation,\n        # like PCEN.\n        features = self.extract_melspec(merged_audio)\n        features = self.augment_features(features, is_train)\n        combined = {\n            'audio': merged_audio,\n            'label': labels,\n            'label_enum': labels_enum,\n            'features': features,\n            'source': ex0['source'],\n        }\n        return combined\n    \n    def extract_melspec(self, audio):\n        # Create the features your model consumes.\n        # Doing this work early saves time for the model.\n        # There are many options; here's an example of creating a melspectrogram.\n        frame_step = self.sample_rate // self.features_rate\n        frame_length = int(0.1 * self.sample_rate)\n        melspec_depth = 100\n        lower_edge_hz = 100.0\n        upper_edge_hz = 12000.0\n        log_floor = 1e-2\n        logmel_scalar = 0.1\n        \n        # Last dimension needs to be the time dimension.\n        stfts = tf.signal.stft(\n            audio[:, :, 0], frame_length=frame_length, frame_step=frame_step,\n            pad_end=True)\n        magnitude_spectrograms = tf.abs(stfts)\n        num_spectrogram_bins = tf.shape(magnitude_spectrograms)[-1]\n        linear_to_mel_weight_matrix = tf.signal.linear_to_mel_weight_matrix(\n            melspec_depth, num_spectrogram_bins,\n            self.sample_rate, lower_edge_hz, upper_edge_hz)\n        mel_spectrograms = tf.tensordot(magnitude_spectrograms,\n                                        linear_to_mel_weight_matrix, 1)\n        mel_spectrograms.set_shape(magnitude_spectrograms.shape[:-1].concatenate(\n                                   linear_to_mel_weight_matrix.shape[-1:]))\n        logmel = tf.math.log(tf.maximum(mel_spectrograms, log_floor))\n        return logmel\n\n    def _parse_and_trim_audio(self, filename):\n        # In order for the map call to work properly, we need to wrap stateful\n        # python operations with py_function. Otherwise, we get weird repeated\n        # audio in the dataset which doesn't match the labels.\n        # See also:\n        # https://www.tensorflow.org/api_docs/python/tf/data/Dataset#map\n        # https://www.tensorflow.org/api_docs/python/tf/py_function\n        # https://www.tensorflow.org/guide/data#applying_arbitrary_python_logic\n        #\n        # Another alternative is to convert the dataset to wav files, and\n        # use the tf.audio.decode_wav to parse the files.\n        # This won't require using a tf.python_function, so can make better\n        # use of multiple cores/threads.\n        def _soundfile_read(filename):\n            with open(filename.numpy(), 'br') as audio_file:\n                tmp = io.BytesIO(audio_file.read())\n                audio, rate = sf.read(tmp, dtype='float32')\n            return audio\n        [audio,] = tf.py_function(_soundfile_read, [filename], [tf.float32])\n        audio.set_shape([None])\n        audio = tf.reshape(audio, [-1, 1])\n\n        audio = self.select_window(audio)\n        label = tf.strings.split(filename, sep='/')[-2]\n        # TODO: Lookup and include metadata.\n        return {'audio': audio, \n                'label': label,\n                'source': filename}\n\n    def build(self, is_train):\n        self.species_enum_table = self.get_species_enum_table()\n\n        # Build a tensorflow Dataset from the training file pattern.\n        ds = tf.data.Dataset.list_files(self.train_file_pattern)\n\n        # Filter species here, to avoid reading files that you won't use.\n        # ds = ds.filter(lambda x: tf.strings.split(x, '/')[-2] != 'batpig1')\n\n        # Create a train/validation split.\n        if is_train:\n            ds = ds.repeat(-1)\n            ds = ds.shuffle(64000)\n            ds = ds.filter(lambda x: tf.strings.to_hash_bucket(x, 100) != 0)\n        else:\n            ds = ds.filter(lambda x: tf.strings.to_hash_bucket(x, 100) == 0)\n\n        \n        # First, we parse the audio from the ogg files and snip to the same length.\n        ds = ds.map(self._parse_and_trim_audio,\n                    num_parallel_calls=10,\n                    deterministic=False)\n        ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)\n\n        # We typically work on batches of examples. \n        # It's also much more efficient for pre-processing.\n        ds = ds.batch(self.batch_size)\n        # It's common to mix multiple examples to increase the difficulty for the model\n        # during training. If you have additional noise files, you can zip those in, too.\n        # Notice that process_examples has a signature which can handle one or two inputs;\n        # this pattern can extend to handle an additional noise file, as well.\n        if is_train:\n            ds = ds.zip((ds, ds))\n            process_fn = lambda x0, x1: self.process_examples(x0, x1, is_train=True)\n        else:\n            process_fn = lambda x0: self.process_examples(x0, is_train=False)\n\n        # It's good to avoid making too many map calls.\n        # So process_examples does all of the augmentation and example-merging.\n        ds = ds.map(process_fn, \n                    num_parallel_calls=tf.data.experimental.AUTOTUNE)\n        # Prefetch creates a buffer for the model to read batches from.\n        # This means the model (hopefully) is never waiting for preprocessing to finish.\n        ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)\n        return ds\n\ntf.compat.v1.reset_default_graph()\nds_builder = BirbsongDatasetBuilder(batch_size=10)\nds = ds_builder.build(is_train=True)\n\n# Now we have a training dataset.\n# We can draw a batch to demonstrate that it works, but this is never how you want\n# to use it in practice!\n# Usually, you'll do model.fit(ds, ...)\n# Then the model will draw automagically from the dataset as needed,\n# and the dataset will keep its pretech cache filled automagically.\nit = iter(ds)\nstarttime = time.time()\nfor i in range(10):\n    btime = time.time()\n    features = it.next()\n    # print('.', end='')\n    print('\\t batch %02d : %5.3f s' % (i, time.time() - btime))\nprint('\\nelasped per batch : ', (time.time() - starttime) / 10)\nprint('Feature labels : ', features['label'])\nprint('Feature label enums : ', features['label_enum'])\nprint('Feature melspec shape : ', features['features'].shape)\n\nplt.figure(figsize=(15, 5))\nmelspecs = features['features'].numpy()\nfor i in range(10):\n    plt.subplot(2, 5, i + 1)\n    plt.imshow(np.flipud(melspecs[i].T),\n               cmap='Greys', aspect='auto')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import IPython.display as ipd\n\n# Look inside one of the examples.\nb = 9\naudio = features['audio'][b, :, 0].numpy()\nprint(features['source'][b].numpy())\nprint(features['label'][b].numpy())\nipd.Audio(data=audio, rate=32000)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"collapsed":true},"cell_type":"code","source":"import tensorflow_hub as hub\n\nN_SPECIES = 397\nLEARNING_RATE = 1e-2\n\ndef build_model(input_shape):\n    # Features shape is [B, T, D]\n    # Designate input 'names' matching the keys in the dataset output dictionaries.\n    features_input = tf.keras.Input(shape=input_shape, name='features')\n    labels_enum = tf.keras.Input(shape=[None], name='label_enum', dtype=tf.int32)\n\n    activations = tf.expand_dims(features_input, -1)\n    inp_layer = tf.keras.layers.Conv2D(64, 5, (2, 2), activation='relu')\n    activations = inp_layer(activations)\n    activations = tf.keras.layers.Conv2D(64, 5, (2, 2), activation='relu')(activations)\n    activations = tf.keras.layers.Conv2D(64, 3, (2, 2), activation='relu')(activations)\n    activations = tf.keras.layers.Flatten()(activations)\n    activations = tf.keras.layers.Dense(2 * N_SPECIES, activation='relu')(activations)\n    logits = tf.keras.layers.Dense(N_SPECIES)(activations)\n    model = tf.keras.Model(inputs=[features_input, labels_enum], outputs=logits)\n    \n    labels_n_hot = tf.one_hot(labels_enum, depth=N_SPECIES, axis=-1)\n    labels_n_hot = tf.reduce_sum(labels_n_hot, axis=1)\n\n    # Define the loss directly using the features in the dataset outputs.\n    # There's not a great way to pass the labels from the dataset to the model.compile, \n    # so far as I can tell. \n    # Same goes for metrics; define them here, using the input features.\n    bxe = tf.keras.losses.BinaryCrossentropy(from_logits=True)(labels_n_hot, logits)\n    model.add_loss(bxe)\n    return model\n\ntf.compat.v1.reset_default_graph()\nmodel = build_model(\n    input_shape=features['features'].shape[1:])\nmodel.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE))\n\nds_builder = BirbsongDatasetBuilder(batch_size=32)\nds = ds_builder.build(is_train=True)\nprint('starting train loop...')\n# This runs fine, but is slow on a single-CPU notebook.\n# model.fit(ds, epochs=1, batch_size=32, verbose=1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}