models.py

# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model definitions for simple speech recognition.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math

import tensorflow as tf


def prepare_model_settings(label_count, sample_rate, clip_duration_ms,
                           window_size_ms, window_stride_ms,
                           dct_coefficient_count,
                           first_filter_height, first_filter_width, first_filter_count,
                           first_convolution_xavier_or_not, first_convolution_stride, first_convolution_padding,
                           first_maxpool_or_not, first_maxpool_stride, first_maxpool_padding,

                           second_filter_height, second_filter_width, second_filter_count,
                           second_convolution_xavier_or_not, second_convolution_stride, second_convolution_padding,
                           second_maxpool_or_not, second_maxpool_stride, second_maxpool_padding,

                           third_filter_height, third_filter_width, third_filter_count,
                           third_convolution_xavier_or_not, third_convolution_stride, third_convolution_padding):
    """Calculates common settings needed for all models.
    Args:
      label_count: How many classes are to be recognized.
      sample_rate: Number of audio samples per second.
      clip_duration_ms: Length of each audio clip to be analyzed.
      window_size_ms: Duration of frequency analysis window.
      window_stride_ms: How far to move in time between frequency windows.
      dct_coefficient_count: Number of frequency bins to use for analysis.
    Returns:
      Dictionary containing common settings.
    """
    desired_samples = int(sample_rate * clip_duration_ms / 1000)
    window_size_samples = int(sample_rate * window_size_ms / 1000)
    window_stride_samples = int(sample_rate * window_stride_ms / 1000)
    length_minus_window = (desired_samples - window_size_samples)
    if length_minus_window < 0:
        spectrogram_length = 0
    else:
        spectrogram_length = 1 + int(length_minus_window / window_stride_samples)
    fingerprint_size = dct_coefficient_count * spectrogram_length
    return {
        'desired_samples': desired_samples,
        'window_size_samples': window_size_samples,
        'window_stride_samples': window_stride_samples,
        'spectrogram_length': spectrogram_length,
        'dct_coefficient_count': dct_coefficient_count,
        'fingerprint_size': fingerprint_size,
        'label_count': label_count,
        'sample_rate': sample_rate,
        'first_filter_height': first_filter_height,
        'first_filter_width': first_filter_width,
        'first_filter_count': first_filter_count,
        'first_convolution_xavier_or_not': first_convolution_xavier_or_not,
        'first_convolution_stride': first_convolution_stride,
        'first_convolution_padding': first_convolution_padding,
        'first_maxpool_or_not': first_maxpool_or_not,
        'first_maxpool_stride': first_maxpool_stride,
        'first_maxpool_padding': first_maxpool_padding,

        'second_filter_height': second_filter_height,
        'second_filter_width': second_filter_width,
        'second_filter_count': second_filter_count,
        'second_convolution_xavier_or_not': second_convolution_xavier_or_not,
        'second_convolution_stride': second_convolution_stride,
        'second_convolution_padding': second_convolution_padding,
        'second_maxpool_or_not': second_maxpool_or_not,
        'second_maxpool_stride': second_maxpool_stride,
        'second_maxpool_padding': second_maxpool_padding,

        'third_filter_height': third_filter_height,
        'third_filter_width': third_filter_width,
        'third_filter_count': third_filter_count,
        'third_convolution_xavier_or_not': third_convolution_xavier_or_not,
        'third_convolution_stride': third_convolution_stride,
        'third_convolution_padding': third_convolution_padding,


    }


def create_model(fingerprint_input, model_settings, model_architecture,
                 is_training, runtime_settings=None):
    """Builds a model of the requested architecture compatible with the settings.
    There are many possible ways of deriving predictions from a spectrogram
    input, so this function provides an abstract interface for creating different
    kinds of models in a black-box way. You need to pass in a TensorFlow node as
    the 'fingerprint' input, and this should output a batch of 1D features that
    describe the audio. Typically this will be derived from a spectrogram that's
    been run through an MFCC, but in theory it can be any feature vector of the
    size specified in model_settings['fingerprint_size'].
    The function will build the graph it needs in the current TensorFlow graph,
    and return the tensorflow output that will contain the 'logits' input to the
    softmax prediction process. If training flag is on, it will also return a
    placeholder node that can be used to control the dropout amount.
    See the implementations below for the possible model architectures that can be
    requested.
    Args:
      fingerprint_input: TensorFlow node that will output audio feature vectors.
      model_settings: Dictionary of information about the model.
      model_architecture: String specifying which kind of model to create.
      is_training: Whether the model is going to be used for training.
      runtime_settings: Dictionary of information about the runtime.
    Returns:
      TensorFlow node outputting logits results, and optionally a dropout
      placeholder.
    Raises:
      Exception: If the architecture type isn't recognized.
    """
    if model_architecture == 'single_fc':
        return create_single_fc_model(fingerprint_input, model_settings,
                                      is_training)
    elif model_architecture == 'conv':
        return create_conv_model(fingerprint_input, model_settings, is_training)
    elif model_architecture == 'low_latency_conv':
        return create_low_latency_conv_model(fingerprint_input, model_settings,
                                             is_training)
    elif model_architecture == 'low_latency_svdf':
        return create_low_latency_svdf_model(fingerprint_input, model_settings,
                                             is_training, runtime_settings)
    else:
        raise Exception('model_architecture argument "' + model_architecture +
                        '" not recognized, should be one of "single_fc", "conv",' +
                        ' "low_latency_conv, or "low_latency_svdf"')


def load_variables_from_checkpoint(sess, start_checkpoint):
    """Utility function to centralize checkpoint restoration.
    Args:
      sess: TensorFlow session.
      start_checkpoint: Path to saved checkpoint on disk.
    """
    saver = tf.train.Saver(tf.global_variables())
    saver.restore(sess, start_checkpoint)


def create_single_fc_model(fingerprint_input, model_settings, is_training):
    """Builds a model with a single hidden fully-connected layer.
    This is a very simple model with just one matmul and bias layer. As you'd
    expect, it doesn't produce very accurate results, but it is very fast and
    simple, so it's useful for sanity testing.
    Here's the layout of the graph:
    (fingerprint_input)
            v
        [MatMul]<-(weights)
            v
        [BiasAdd]<-(bias)
            v
    Args:
      fingerprint_input: TensorFlow node that will output audio feature vectors.
      model_settings: Dictionary of information about the model.
      is_training: Whether the model is going to be used for training.
    Returns:
      TensorFlow node outputting logits results, and optionally a dropout
      placeholder.
    """
    if is_training:
        dropout_prob = tf.placeholder(tf.float32, name='dropout_prob')
    fingerprint_size = model_settings['fingerprint_size']
    label_count = model_settings['label_count']
    weights = tf.Variable(
        tf.truncated_normal([fingerprint_size, label_count], stddev=0.001))
    bias = tf.Variable(tf.zeros([label_count]))
    logits = tf.matmul(fingerprint_input, weights) + bias
    if is_training:
        return logits, dropout_prob
    else:
        return logits


def create_conv_model(fingerprint_input, model_settings, is_training):
    """Builds a standard convolutional model.
    This is roughly the network labeled as 'cnn-trad-fpool3' in the
    'Convolutional Neural Networks for Small-footprint Keyword Spotting' paper:
    http://www.isca-speech.org/archive/interspeech_2015/papers/i15_1478.pdf
    Here's the layout of the graph:
    (fingerprint_input)
            v
        [Conv2D]<-(weights)
            v
        [BiasAdd]<-(bias)
            v
          [Relu]
            v
        [MaxPool]
            v
        [Conv2D]<-(weights)
            v
        [BiasAdd]<-(bias)
            v
          [Relu]
            v
        [MaxPool]
            v
        [MatMul]<-(weights)
            v
        [BiasAdd]<-(bias)
            v
    This produces fairly good quality results, but can involve a large number of
    weight parameters and computations. For a cheaper alternative from the same
    paper with slightly less accuracy, see 'low_latency_conv' below.
    During training, dropout nodes are introduced after each relu, controlled by a
    placeholder.
    Args:
      fingerprint_input: TensorFlow node that will output audio feature vectors.
      model_settings: Dictionary of information about the model.
      is_training: Whether the model is going to be used for training.
    Returns:
      TensorFlow node outputting logits results, and optionally a dropout
      placeholder.
    """
    if is_training:
        dropout_prob = tf.placeholder(tf.float32, name='dropout_prob')
    input_frequency_size = model_settings['dct_coefficient_count']
    input_time_size = model_settings['spectrogram_length']
    fingerprint_4d = tf.reshape(fingerprint_input,
                                [-1, input_time_size, input_frequency_size, 1])
    print(fingerprint_4d)
    first_filter_width = model_settings['first_filter_width']
    first_filter_height = model_settings['first_filter_height']
    first_filter_count = model_settings['first_filter_count']

    first_bias = tf.Variable(tf.zeros([first_filter_count]))

    """
    if model_settings['first_convolution_xavier_or_not']:
        first_weights = tf.get_variable(
            shape=[first_filter_height, first_filter_width, 1, first_filter_count],
            initializer=tf.contrib.layers.xavier_initializer(),name=first_xavier)
    else:
    """
    first_weights = tf.Variable(
            tf.truncated_normal(
                [first_filter_height, first_filter_width, 1, first_filter_count],
                stddev=0.01))


    cst = model_settings["first_convolution_stride"]
    cst.insert(0,1)
    cst.insert(3,1)
    if model_settings['first_convolution_padding']:
        first_conv = tf.nn.conv2d(fingerprint_4d, first_weights, cst,
                                  'VALID') + first_bias
    else:
        first_conv = tf.nn.conv2d(fingerprint_4d, first_weights, cst,
                                  'SAME')
    first_relu = tf.nn.relu(first_conv)
    if is_training:
        first_dropout = tf.nn.dropout(first_relu, dropout_prob)
    else:
        first_dropout = first_relu

    mst = model_settings["first_maxpool_stride"]
    mst.insert(0,1)
    mst.insert(3,1)
    if model_settings["first_maxpool_padding"]:
        max_pool = tf.nn.max_pool(first_dropout, mst, mst, 'VALID')
    else:
        max_pool = tf.nn.max_pool(first_dropout, mst, mst, 'SAME')
    

    second_filter_width = model_settings["second_filter_width"]
    second_filter_height = model_settings["second_filter_height"]
    second_filter_count = model_settings["second_filter_count"]

    second_bias = tf.Variable(tf.zeros([second_filter_count]))
    """
    if model_settings["second_convolution_xavier_or_not"]:
        second_weights = tf.get_variable(
            shape=[second_filter_height, second_filter_width, first_filter_count, second_filter_count],
            initializer=tf.contrib.layers.xavier_initializer(), name=second_xavier)

    else:
    """
    second_weights = tf.Variable(
            tf.truncated_normal(
                [
                    second_filter_height, second_filter_width, first_filter_count, second_filter_count
                ],
                stddev=0.01))
    cst = model_settings["second_convolution_stride"]
    cst.insert(0,1)
    cst.insert(3,1)
    if model_settings["first_maxpool_or_not"]:
        if model_settings["second_convolution_padding"]:
            second_conv = tf.nn.conv2d(max_pool, second_weights, cst,
                                       'VALID') + second_bias
        else:
            second_conv = tf.nn.conv2d(max_pool, second_weights, cst,
                                       'SAME') + second_bias
            
    else:
        if model_settings["second_convolution_padding"]:
            second_conv = tf.nn.conv2d(first_dropout, second_weights, cst,
                                       'VALID') + second_bias
        else:
            second_conv = tf.nn.conv2d(first_dropout, second_weights, cst,
                                        'SAME') + second_bias

    second_relu = tf.nn.relu(second_conv)
    if is_training:
        second_dropout = tf.nn.dropout(second_relu, dropout_prob)
    else:
        second_dropout = second_relu

    mst = model_settings["second_maxpool_stride"]
    mst.insert(0,1)
    mst.insert(3,1)
    if model_settings["second_maxpool_or_not"]:
        max_pool = tf.nn.max_pool(second_dropout, mst, mst, 'VALID')
    else:
        max_pool = tf.nn.max_pool(second_dropout, mst, mst, 'SAME')


    third_filter_width = model_settings["third_filter_width"]
    third_filter_height = model_settings["third_filter_height"]
    third_filter_count = model_settings["third_filter_count"]

    third_bias = tf.Variable(tf.zeros([third_filter_count]))
    """
    if model_settings["third_convolution_xavier_or_not"]:
        third_weights = tf.get_variable(
            shape=[third_filter_height, third_filter_width, second_filter_count, third_filter_count],
            initializer=tf.contrib.layers.xavier_initializer(), name=third_xavier)

    else:
    """
    third_weights = tf.Variable(
            tf.truncated_normal(
                [
                    third_filter_height, third_filter_width, second_filter_count, third_filter_count,
                ],
                stddev=0.01))

    cst = model_settings["third_convolution_stride"]
    cst.insert(0,1)
    cst.insert(3,1)
    if model_settings["second_maxpool_or_not"]:
        if model_settings["third_convolution_padding"]:
            third_conv = tf.nn.conv2d(max_pool, third_weights, cst,
                                      'VALID') + third_bias
        else:
            third_conv = tf.nn.conv2d(max_pool, third_weights, cst,
                                      'SAME') + third_bias
    else:
        if model_settings["third_convolution_padding"]:
            third_conv = tf.nn.conv2d(second_dropout, third_weights, cst,
                                      'VALID') + third_bias
        else:
            third_conv = tf.nn.conv2d(second_dropout, third_weights, cst,
                                      'SAME') + third_bias
    
    third_relu = tf.nn.relu(third_conv)
    if is_training:
        third_dropout = tf.nn.dropout(third_relu, dropout_prob)
    else:
        third_dropout = third_relu

    third_conv_shape = third_dropout.get_shape()
    third_conv_output_width = third_conv_shape[2]
    third_conv_output_height = third_conv_shape[1]
    third_conv_element_count = int(
        third_conv_output_width * third_conv_output_height *
        third_filter_count)
    flattened_third_conv = tf.reshape(third_dropout,
                                      [-1, third_conv_element_count])
    label_count = model_settings['label_count']
    final_fc_weights = tf.Variable(
        tf.truncated_normal(
            [third_conv_element_count, label_count], stddev=0.01))
    final_fc_bias = tf.Variable(tf.zeros([label_count]))
    final_fc = tf.matmul(flattened_third_conv, final_fc_weights) + final_fc_bias
    

    if is_training:
        return final_fc, dropout_prob
    else:
        return final_fc