使用Estimator接口通过预训练的张量流对象检测模型进行推理

Question

我正在尝试从Tensorflow Object Detection回购中加载经过预训练的张量流对象检测模型作为tf.estimator.Estimator并使用它进行预测。

我能够使用Estimator.predict()加载模型并运行推理，但是输出是垃圾。加载模型的其他方法，例如作为Predictor，并且运行推断工作正常。

任何帮助正确加载模型（以Estimator调用predict()的方式都将受到赞赏。我当前的代码：

加载并准备图像

def load_image_into_numpy_array(image):
    (im_width, im_height) = image.size
    return np.array(list(image.getdata())).reshape((im_height, im_width, 3)).astype(np.uint8)

image_url = 'https://i.imgur.com/rRHusZq.jpg'

# Load image
response = requests.get(image_url)
image = Image.open(BytesIO(response.content))

# Format original image size
im_size_orig = np.array(list(image.size) + [1])
im_size_orig = np.expand_dims(im_size_orig, axis=0)
im_size_orig = np.int32(im_size_orig)

# Resize image
image = image.resize((np.array(image.size) / 4).astype(int))

# Format image
image_np = load_image_into_numpy_array(image)
image_np_expanded = np.expand_dims(image_np, axis=0)
image_np_expanded = np.float32(image_np_expanded)

# Stick into feature dict
x = {'image': image_np_expanded, 'true_image_shape': im_size_orig}

# Stick into input function
predict_input_fn = tf.estimator.inputs.numpy_input_fn(
    x=x,
    y=None,
    shuffle=False,
    batch_size=128,
    queue_capacity=1000,
    num_epochs=1,
    num_threads=1,
)

旁注：

train_and_eval_dict似乎还包含input_fn用于预测

train_and_eval_dict['predict_input_fn']

但是实际上返回一个tf.estimator.export.ServingInputReceiver，我不确定该怎么做。这可能是我遇到问题的根源，因为在模型实际看到图像之前需要进行大量预处理。

将模型加载为Estimator

从TF模型动物园here下载的模型，用于加载从here改编的模型的代码。

model_dir = './pretrained_models/tensorflow/ssd_mobilenet_v1_coco_2018_01_28/'
pipeline_config_path = os.path.join(model_dir, 'pipeline.config')

config = tf.estimator.RunConfig(model_dir=model_dir)

train_and_eval_dict = model_lib.create_estimator_and_inputs(
    run_config=config,
    hparams=model_hparams.create_hparams(None),
    pipeline_config_path=pipeline_config_path,
    train_steps=None,
    sample_1_of_n_eval_examples=1,
    sample_1_of_n_eval_on_train_examples=(5))

estimator = train_and_eval_dict['estimator']

运行推断

output_dict1 = estimator.predict(predict_input_fn)

这会打印出一些日志消息，其中之一是：

INFO:tensorflow:Restoring parameters from ./pretrained_models/tensorflow/ssd_mobilenet_v1_coco_2018_01_28/model.ckpt

因此，似乎已经加载了预训练的重量。但是结果如下：

加载与Predictor相同的模型

from tensorflow.contrib import predictor

model_dir = './pretrained_models/tensorflow/ssd_mobilenet_v1_coco_2018_01_28'
saved_model_dir = os.path.join(model_dir, 'saved_model')
predict_fn = predictor.from_saved_model(saved_model_dir)

运行推断

output_dict2 = predict_fn({'inputs': image_np_expanded})

结果看起来不错：

Answer 1

当您将模型作为估计量并从检查点文件加载时，这是与ssd模型关联的还原功能。来自ssd_meta_arch.py

def restore_map(self,
                  fine_tune_checkpoint_type='detection',
                  load_all_detection_checkpoint_vars=False):
    """Returns a map of variables to load from a foreign checkpoint.
    See parent class for details.
    Args:
      fine_tune_checkpoint_type: whether to restore from a full detection
        checkpoint (with compatible variable names) or to restore from a
        classification checkpoint for initialization prior to training.
        Valid values: `detection`, `classification`. Default 'detection'.
      load_all_detection_checkpoint_vars: whether to load all variables (when
         `fine_tune_checkpoint_type='detection'`). If False, only variables
         within the appropriate scopes are included. Default False.
    Returns:
      A dict mapping variable names (to load from a checkpoint) to variables in
      the model graph.
    Raises:
      ValueError: if fine_tune_checkpoint_type is neither `classification`
        nor `detection`.
    """
    if fine_tune_checkpoint_type not in ['detection', 'classification']:
      raise ValueError('Not supported fine_tune_checkpoint_type: {}'.format(
          fine_tune_checkpoint_type))

    if fine_tune_checkpoint_type == 'classification':
      return self._feature_extractor.restore_from_classification_checkpoint_fn(
          self._extract_features_scope)

    if fine_tune_checkpoint_type == 'detection':
      variables_to_restore = {}
      for variable in tf.global_variables():
        var_name = variable.op.name
        if load_all_detection_checkpoint_vars:
          variables_to_restore[var_name] = variable
        else:
          if var_name.startswith(self._extract_features_scope):
            variables_to_restore[var_name] = variable

    return variables_to_restore

如您所见，即使配置文件设置了from_detection_checkpoint: True，也只会还原功能提取器作用域中的变量。要恢复所有变量，您必须设置

load_all_detection_checkpoint_vars: True

在配置文件中。

因此，以上情况很明显。当将模型加载为Estimator时，将仅还原来自特征提取器作用域的变量，并且不还原预测器的作用域权重，估计器显然会给出随机预测。

在将模型加载为预测变量时，将加载所有权重，因此预测是合理的。