export_model.py 文件源码

def build_prediction_graph(self, serialized_examples):
    model_input_raw, labels_batch = (
        self.reader.prepare_serialized_examples(serialized_examples))

    model_input = model_input_raw

    with tf.variable_scope("tower"):
      result = self.model.create_model(
          model_input,
          num_classes=self.reader.num_classes,
          labels=labels_batch,
          is_training=False)

      for variable in slim.get_model_variables():
        tf.summary.histogram(variable.op.name, variable)

      predictions = result["predictions"]

      prediction, index = tf.nn.top_k(predictions, 1)
    return prediction, index