使用 Estimator 接口通過預訓練的 tensorflow 對象檢測模型進行推理

Question

我正在嘗試從Tensorflow 對象檢測存儲庫中加載一個預訓練的 tensorflow 對象檢測模型作為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 Model Zoo 下載在這里，加載模型的代碼從這里改編。

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 ，只會恢復特征提取器范圍中的變量，而不會恢復預測器的范圍權重，估計器顯然會給出隨機預測。

當加載模型作為預測器時，所有權重都被加載，因此預測是合理的。