def _build_keras_model()

in tfx/experimental/templates/taxi/models/keras_model/model.py [0:0]

• 66 lines of code
• 11 McCabe index (conditional complexity)

def _build_keras_model(hidden_units, learning_rate):
  """Creates a DNN Keras model for classifying taxi data.

  Args:
    hidden_units: [int], the layer sizes of the DNN (input layer first).
    learning_rate: [float], learning rate of the Adam optimizer.

  Returns:
    A keras Model.
  """
  deep_input = {
      colname: tf.keras.layers.Input(name=colname, shape=(1,), dtype=tf.float32)
      for colname in features.transformed_names(features.DENSE_FLOAT_FEATURE_KEYS)
  }
  wide_vocab_input = {
      colname: tf.keras.layers.Input(name=colname, shape=(1,), dtype='int32')
      for colname in features.transformed_names(features.VOCAB_FEATURE_KEYS)
  }
  wide_bucket_input = {
      colname: tf.keras.layers.Input(name=colname, shape=(1,), dtype='int32')
      for colname in features.transformed_names(features.BUCKET_FEATURE_KEYS)
  }
  wide_categorical_input = {
      colname: tf.keras.layers.Input(name=colname, shape=(1,), dtype='int32')
      for colname in features.transformed_names(features.CATEGORICAL_FEATURE_KEYS)
  }
  input_layers = {
      **deep_input,
      **wide_vocab_input,
      **wide_bucket_input,
      **wide_categorical_input,
  }

  deep = tf.keras.layers.concatenate(
      [tf.keras.layers.Normalization()(layer) for layer in deep_input.values()]
  )
  for numnodes in (hidden_units or [100, 70, 50, 25]):
    deep = tf.keras.layers.Dense(numnodes)(deep)

  wide_layers = []
  for key in features.transformed_names(features.VOCAB_FEATURE_KEYS):
    wide_layers.append(
        tf.keras.layers.CategoryEncoding(num_tokens=features.VOCAB_SIZE + features.OOV_SIZE)(
            input_layers[key]
        )
    )
  for key, num_tokens in zip(
      features.transformed_names(features.BUCKET_FEATURE_KEYS),
      features.BUCKET_FEATURE_BUCKET_COUNT,
  ):
    wide_layers.append(
        tf.keras.layers.CategoryEncoding(num_tokens=num_tokens)(
                input_layers[key]
        )
    )
  for key, num_tokens in zip(
      features.transformed_names(features.CATEGORICAL_FEATURE_KEYS),
      features.CATEGORICAL_FEATURE_MAX_VALUES,
  ):
    wide_layers.append(
        tf.keras.layers.CategoryEncoding(num_tokens=num_tokens)(
            input_layers[key]
        )
    )
  wide = tf.keras.layers.concatenate(wide_layers)

  output = tf.keras.layers.Dense(1, activation='sigmoid')(
      tf.keras.layers.concatenate([deep, wide])
  )
  output = tf.keras.layers.Reshape((1,))(output)

  model = tf.keras.Model(input_layers, output)
  model.compile(
      loss='binary_crossentropy',
      optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
      metrics=[tf.keras.metrics.BinaryAccuracy()],
  )
  model.summary(print_fn=logging.info)
  return model