model/tf_models.py

# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Tensorflow models used in Bad CL Detector service."""

from __future__ import print_function
import functools

import tensorflow as tf

from lib import feature_preprocessor

# pylint: disable=pointless-statement
# pylint: disable=no-value-for-parameter

# a small number to add to prediction result for taking log.
LOG_RESIDUAL = 1e-15


def DoubleWrap(function):
  """A decorator decorator.

  This allows decorator to be used without parentheses if no arguments are
  provided. All arguments must be optional.

  ref: https://danijar.com/structuring-your-tensorflow-models/

  Args:
    function: a function to wrap.

  Returns:
    a decorator to wrap functions.
  """
  @functools.wraps(function)
  def Decorator(*args, **kwargs):
    if (len(args) == 1) and (len(kwargs) is 0) and (callable(args[0])):
      return function(args[0])
    else:
      return lambda wrapee: function(wrapee, *args, **kwargs)
  return Decorator


@DoubleWrap
def DefineScope(function, scope=None, *args, **kwargs):
  """A decorator for functions that define TensorFlow operations.

  The wrapped function will only be executed once. Subsequent calls to it will
  directly return the result so that operations are added to the graph only
  once. The operations added by the function live within a tf.variable_scope().
  If this decorator is used with arguments, they will be forwarded to the
  variable scope. The scope name defaults to the name of the wrapped function.

  ref: https://danijar.com/structuring-your-tensorflow-models/

  Args:
    function: a tensorflow operation function to wrap.
    scope: string, the name of the scope of variables in the function.
    *args: optional arguments to pass to function.
    **kwargs: optional arguments to pass to function.

  Returns:
    a decorator to wrap tensorflow operation function.
  """
  attribute = '_cache_' + function.__name__
  name = scope or function.__name__
  @property
  @functools.wraps(function)
  def Decorator(self):
    if not hasattr(self, attribute):
      with tf.variable_scope(name, *args, **kwargs):
        setattr(self, attribute, function(self))
    return getattr(self, attribute)
  return Decorator


class BadCLExistenceModel(object):
  """Model to calculate probability of having a bad cl for a failed build."""
  # size of first hidden layer.
  HIDDEN_LAYER_1_SIZE = 16
  # size of second hidden layer.
  HIDDEN_LAYER_2_SIZE = 8
  # number of output classes, should be 2 for 'bad_cl' and 'non bad_cl'.
  OUTPUT_CLASSES = 2
  # learning rate for ML pipeline.
  LEARNING_RATE = 0.0001
  L2_LOSS_COEFFICIENT = 0.01

  def __init__(self, features, labels, keep_prob):
    """Initialize a FailureTypeModel.

    Args:
      features: stage features as produced by StageFeaturePreprocessor.
      labels: stage labels as produced by StageFeaturePreprocessor.
      keep_prob: a float of (0, 1]. Keep probability for the dropout layer.
    """
    self.features = features
    self.labels = labels
    self.keep_prob = keep_prob
    # these statements are necessary in tensorflow model. Refer to ref in
    # define_scope()
    self.Predict
    self.Optimize
    self.GetError

  @DefineScope(initializer=tf.contrib.slim.xavier_initializer())
  def Predict(self):
    """Calculates the probabilities of CQ failure reasons being bad cl.

    Returns:
      a numpy array of floats in range [0, 1]. Each float represents the
      probabilities of the build failure reason is 'bad cl'.
    """
    nodes = self.features
    nodes = tf.contrib.slim.fully_connected(nodes, self.HIDDEN_LAYER_1_SIZE)
    nodes = tf.contrib.slim.dropout(nodes, self.keep_prob)
    nodes = tf.contrib.slim.fully_connected(nodes, self.HIDDEN_LAYER_2_SIZE)
    nodes = tf.contrib.slim.dropout(nodes, self.keep_prob)
    probabilities = tf.contrib.slim.fully_connected(nodes,
                                                    self.OUTPUT_CLASSES,
                                                    tf.nn.softmax)

    return probabilities

  @DefineScope
  def Optimize(self):
    """Returns a tensflow optimizer operation function.

    Returns:
      an optimizer for tensflow model.
    """
    # weights and biases in the graph defined above. Elements 0, 2, 4 are
    # weights while 1, 3, 5 are biases. Apply L2 regularization to weights.
    weights_biases = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)

    one_hot = tf.one_hot(self.labels, self.OUTPUT_CLASSES)

    loss = (tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
        logits=self.Predict, labels=one_hot)) +
            self.L2_LOSS_COEFFICIENT*tf.nn.l2_loss(weights_biases[0]) +
            self.L2_LOSS_COEFFICIENT*tf.nn.l2_loss(weights_biases[2]) +
            self.L2_LOSS_COEFFICIENT*tf.nn.l2_loss(weights_biases[4]))

    optimizer = tf.train.RMSPropOptimizer(self.LEARNING_RATE)
    return optimizer.minimize(loss)

  @DefineScope
  def GetError(self):
    """Calculates the training error in current model.

    Returns:
      a float in range [0, 1], which is the percentage of wrongly classified
      build counts over total build counts in the CQ.
    """
    one_hot = tf.one_hot(self.labels, 2)
    mistakes = tf.not_equal(tf.argmax(one_hot, 1),
                            tf.argmax(self.Predict, 1))
    error = tf.reduce_mean(tf.cast(mistakes, tf.float32), name='error')
    return error


class BadCLDetectionModel(object):
  """Model to calculate probability of a CL to be a bad one."""
  LEARNING_RATE = 0.005
  L2_LOSS_COEFFICIENT = 0.05
  # total class count: bad cl or other cl
  OUTPUT_CLASSES = 2
  # hidden layer 1 size
  HIDDEN_LAYER_1_SIZE = 16
  # hidden layer 2 size
  HIDDEN_LAYER_2_SIZE = 8
  # parameters for RNN of exception features
  EXCEPTION_EMBEDDING_SIZE = 100
  EXCEPTION_OUTPUT_SIZE = 32
  EXCEPTION_LSTM_SIZE = 64
  EXCEPTION_LSTM_LAYERS = 1
  # parameters for RNN of commit features
  COMMIT_EMBEDDING_SIZE = 200
  COMMIT_OUTPUT_SIZE = 64
  COMMIT_LSTM_SIZE = 128
  COMMIT_LSTM_LAYERS = 1
  # RNN vocabulary sizes for exception and commit, plus one for None input.
  EXCEPTION_RNN_SIZE = len(
      feature_preprocessor.CLFeaturePreprocessor.GetWord2IntMapping(
          None, 'exception_rnn', False)) + 1
  COMMIT_RNN_SIZE = len(
      feature_preprocessor.CLFeaturePreprocessor.GetWord2IntMapping(
          None, 'commit_rnn', False)) + 1

  def __init__(self, nn_features, exception_rnn_input, commit_rnn_input, labels,
               keep_prob):
    self.nn_features = nn_features
    self.exception_rnn_input = exception_rnn_input
    self.commit_rnn_input = commit_rnn_input
    self.labels = labels
    self.keep_prob = keep_prob
    self.Predict
    self.Optimize
    self.GetError

  @DefineScope(initializer=tf.contrib.slim.xavier_initializer())
  def Predict(self):
    """Calculates the probabilities of CLs being bad.

    Returns:
      a numpy array of floats in range [0, 1]. Each float represents the chance
      of a CL being bad.
    """
    with tf.variable_scope('exception_rnn'):
      exception_lstm_cell = self.GetLSTMCellStack(self.EXCEPTION_LSTM_SIZE,
                                                  self.EXCEPTION_LSTM_LAYERS)

      exception_rnn_input = self.GetEmbeddedInput(self.exception_rnn_input,
                                                  self.EXCEPTION_RNN_SIZE,
                                                  self.EXCEPTION_EMBEDDING_SIZE)

      exception_rnn_outputs = self.GetRNNOutputs(exception_lstm_cell,
                                                 exception_rnn_input,
                                                 self.EXCEPTION_OUTPUT_SIZE)

    with tf.variable_scope('commit_rnn'):
      commit_lstm_cell = self.GetLSTMCellStack(self.COMMIT_LSTM_SIZE,
                                               self.COMMIT_LSTM_LAYERS)

      commit_rnn_input = self.GetEmbeddedInput(self.commit_rnn_input,
                                               self.COMMIT_RNN_SIZE,
                                               self.COMMIT_EMBEDDING_SIZE)

      commit_rnn_outputs = self.GetRNNOutputs(commit_lstm_cell,
                                              commit_rnn_input,
                                              self.COMMIT_OUTPUT_SIZE)

    rnn_outputs = tf.concat([exception_rnn_outputs, commit_rnn_outputs], 1)

    # get final output from mixure NN model.
    logits = self.GetNNOuputs(rnn_outputs)

    return logits

  def GetLSTMCellStack(self, cell_size, number_of_layers):
    """Get a lstm cell stack with dropout layers.

    Args:
      cell_size: an int, size of the LSTM cell.
      number_of_layers: an int, number of LSTM cell layers.

    Returns:
     A tuple of (stacked LSTM cells).
    """
    lstm_cell = tf.contrib.rnn.BasicLSTMCell(cell_size)
    drop = tf.contrib.rnn.DropoutWrapper(lstm_cell,
                                         output_keep_prob=self.keep_prob)
    # Stack up multiple LSTM cells for deep learning
    lstm_cell_stack = tf.contrib.rnn.MultiRNNCell([drop] * number_of_layers)

    return lstm_cell_stack

  def GetRNNOutputs(self, lstm_cell, embedded_input, output_size):
    """Get RNN features for given embedded input and LSTM cell.

    Args:
      lstm_cell: a stack of BasicLSTMCell.
      embedded_input: a numpy array of embedded text inputs.
      output_size: an int, number of output nodes of current RNN.

    Returns:
      fully connected layer of size output_size from RNN.
    """
    raw_outputs, _ = tf.nn.dynamic_rnn(lstm_cell, embedded_input,
                                       dtype=tf.float32)

    # connect the last layer of RNN to RNN_OUT_SIZE output nodes.
    rnn_outputs = tf.contrib.layers.fully_connected(raw_outputs[:, -1],
                                                    output_size,
                                                    activation_fn=tf.sigmoid)
    return rnn_outputs

  def GetNNOuputs(self, rnn_outputs):
    """Get output logics from the mixture regular and recurrent NN.

    Args:
      rnn_outputs: output features from the RNNs.

    Returns:
      final logits of classification results.
    """
    weights, biases = self.GetWeightsAndBiases(
        rnn_outputs.get_shape().as_list()[1])
    # Combine regular nn features and RNN outputs
    layer_1_inputs = tf.concat([self.nn_features, rnn_outputs], 1)
    # First hidden layer with RELU activation
    layer_1 = tf.add(tf.matmul(layer_1_inputs, weights['hidden_layer_1']),
                     biases['hidden_layer_1'])
    layer_1 = tf.nn.relu(layer_1)
    # add a dropout layer
    layer_1 = tf.nn.dropout(layer_1, self.keep_prob)
    # add a second layer
    layer_2 = tf.add(tf.matmul(layer_1, weights['hidden_layer_2']),
                     biases['hidden_layer_2'])
    # Combine RNN and NN layers to produce output logits
    logits = tf.add(tf.matmul(layer_2, weights['output']), biases['output'])

    return logits

  def GetEmbeddedInput(self, rnn_input, vocabulary_size, embedding_size):
    """Get embedded input of the text ids as inputs for RNN.

    Args:
      rnn_input: a numpy array of word id integers.
      vocabulary_size: an integer, total size of vaocabulary of given input.
      embedding_size: an integer, output embedding dimension of input.

    Returns:
      an embedded input ready for RNN.
    """
    embedding = tf.Variable(tf.random_uniform((vocabulary_size,
                                               embedding_size), -1, 1))

    return tf.nn.embedding_lookup(embedding, rnn_input)

  def GetWeightsAndBiases(self, rnn_outputs_dim):
    """Get weights and biases.

    Args:
      rnn_outputs_dim: an integer, the dimension of all RNN outputs.

    Returns:
      weights and biases dictionaries, to be used in the NN.
    """
    # regular features dimension + RNN output dimension
    input_layer_size = (self.nn_features.get_shape().as_list()[1] +
                        rnn_outputs_dim)

    with tf.variable_scope('defining_weights'):
      layer_1_weights = tf.get_variable(
          'layer_1_weights', [input_layer_size, self.HIDDEN_LAYER_1_SIZE])

      layer_2_weights = tf.get_variable(
          'layer_2_weights', [self.HIDDEN_LAYER_1_SIZE,
                              self.HIDDEN_LAYER_2_SIZE])

      out_weights = tf.get_variable(
          'out_weights', [self.HIDDEN_LAYER_2_SIZE, self.OUTPUT_CLASSES])

    weights = {
        'hidden_layer_1': layer_1_weights,
        'hidden_layer_2': layer_2_weights,
        'output': out_weights}

    biases = {
        'hidden_layer_1': tf.Variable(tf.random_normal(
            [self.HIDDEN_LAYER_1_SIZE])),
        'hidden_layer_2': tf.Variable(tf.random_normal(
            [self.HIDDEN_LAYER_2_SIZE])),
        'output': tf.Variable(tf.random_normal([self.OUTPUT_CLASSES]))}

    return weights, biases

  @DefineScope('Predict')
  def Optimize(self):
    """Returns a tensflow optimizer operation function.

    Returns:
      an optimizer for tensflow model.
    """
    with tf.variable_scope('defining_weights', reuse=True):
      layer_1_weights = tf.get_variable('layer_1_weights')
      layer_2_weights = tf.get_variable('layer_2_weights')
      out_weights = tf.get_variable('out_weights')

    one_hot = tf.one_hot(self.labels, self.OUTPUT_CLASSES)

    loss = (tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
        logits=self.Predict, labels=one_hot)) +
            self.L2_LOSS_COEFFICIENT*tf.nn.l2_loss(layer_1_weights) +
            self.L2_LOSS_COEFFICIENT*tf.nn.l2_loss(layer_2_weights) +
            self.L2_LOSS_COEFFICIENT*tf.nn.l2_loss(out_weights))

    optimizer = tf.train.RMSPropOptimizer(self.LEARNING_RATE)

    return optimizer.minimize(loss)

  @DefineScope
  def GetError(self):
    """Calculates the training error in current model.

    Returns:
      a float in range [0, 1], which is the percentage of wrongly classified CL
      counts over total CL counts.
    """
    one_hot = tf.one_hot(self.labels, self.OUTPUT_CLASSES)
    mistakes = tf.not_equal(tf.argmax(one_hot, 1),
                            tf.argmax(self.Predict, 1))
    return tf.reduce_mean(tf.cast(mistakes, tf.float32))