code: def small_cnn_fn(features, labels, mode):
      """Setup a small CNN for image classification."""
      input_layer = tf.reshape(features["x"], [-1, 32, 32, 3])
      for _ in range(3):
          y = tf.keras.layers.Conv2D(32, (3, 3), activation="relu")(
              input_layer
          )
          y = tf.keras.layers.MaxPool2D()(y)
  
      y = tf.keras.layers.Flatten()(y)
      y = tf.keras.layers.Dense(64, activation="relu")(y)
      logits = tf.keras.layers.Dense(10)(y)
  
      if mode != tf_estimator.ModeKeys.PREDICT:
          vector_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
              labels=labels, logits=logits
          )
          scalar_loss = tf.reduce_mean(input_tensor=vector_loss)
  
      # Configure the training op (for TRAIN mode).
      if mode == tf_estimator.ModeKeys.TRAIN:
          optimizer = tf.train.MomentumOptimizer(
              learning_rate=FLAGS.learning_rate, momentum=0.9
          )
          global_step = tf.compat.v1.train.get_global_step()
          train_op = optimizer.minimize(
              loss=scalar_loss, global_step=global_step
          )
          return tf_estimator.EstimatorSpec(
              mode=mode, loss=scalar_loss, train_op=train_op
          )
  
      # Add evaluation metrics (for EVAL mode).
      elif mode == tf_estimator.ModeKeys.EVAL:
          eval_metric_ops = {
              "accuracy": tf.metrics.accuracy(
                  labels=labels,
                  predictions=tf.argmax(input=logits, axis=1),
              )
          }
          return tf_estimator.EstimatorSpec(
              mode=mode,
              loss=scalar_loss,
              eval_metric_ops=eval_metric_ops,
          )
  
      # Output the prediction probability (for PREDICT mode).
      elif mode == tf_estimator.ModeKeys.PREDICT:
          predictions = tf.nn.softmax(logits)
          return tf_estimator.EstimatorSpec(
              mode=mode, predictions=predictions
          )
  
file_path: privacy/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/tf_estimator_evaluation_example.py
imports: from absl import app
  from absl import flags
  from absl import logging
  import numpy
  import tensorflow
  from tensorflow import estimator
  from tensorflow.compat.v1 import estimator
  from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.data_structures import AttackType
  from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.data_structures import get_flattened_attack_metrics
  from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.data_structures import SlicingSpec
  from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.tf_estimator_evaluation import MembershipInferenceTrainingHook
  from tensorflow_privacy.privacy.privacy_tests.membership_inference_attack.tf_estimator_evaluation import run_attack_on_tf_estimator_model
label_desc: Use image node_id 1 for calling a global function with example usage: small_cnn_fn(features, labels, mode) and returns: tf_estimator, tf_estimator, tf_estimator