tensorflow.saved_model.load() 需要极长的执行时间

tensorflow.saved_model.load() takes extremely long time to execute

我有以下功能可以从 tensorflow model zoo 下载、存储和加载模型:

def load_object_detection_model(model_name: str):

    models = load_model_zoo_list()
    model_url = models[model_name]['url']
    model_filename = models[model_name]['filename']

    pretrained_path = os.path.join(os.path.dirname(__file__), "pretrained_models")
    os.makedirs(pretrained_path, exist_ok=True)

    get_file(fname=model_filename, origin=model_url, cache_dir=pretrained_path, cache_subdir='cptr', extract=True)

    loaded_model = tf.saved_model.load(os.path.join(pretrained_path, 'cptr', model_name, "saved_model"))

    return loaded_model

def load_model_zoo_list():
    """

    :return:
    """

    path = os.path.join(os.path.dirname(__file__), "model_zoo.json")
    with open(path, 'r') as f:
        model_zoo_json = json.load(f)

    return model_zoo_json

model_zoo.json

{
  "ssd_mobilenet_v2_320x320_coco17_tpu-8": {
    "url": "http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_mobilenet_v2_320x320_coco17_tpu-8.tar.gz",
    "filename": "ssd_mobilenet_v2_320x320_coco17_tpu-8.tar.gz"
  }
}

想法是稍后向 json 添加更多模型,目前只是选择 ssd_mobilenet_v2_320x320_coco17_tpu-8 进行测试。

问题如下。 loaded_model = tf.saved_model.load(os.path.join(pretrained_path, 'cptr', model_name, "saved_model")) 行大约需要 25-30 秒才能执行。此时已经下载了模型,saved_model 文件夹的大小约为 32Mb。我还测试了更大的模型,这花费了更长的时间。推理似乎也很慢(与模型动物园页面上列出的速度相比)。

除此之外,该模型似乎有效。

这些模型这么慢的原因是什么?

知道了!在第一次调用模型时,构建了图形,因此第一次调用模型总是很慢。我使用 GPU 在 google colab 上尝试了您的代码:

model = load_object_detection_model("ssd_mobilenet_v2_320x320_coco17_tpu-8")
%%time
a= model(np.random.randint(0, 255, size=(1, 320, 320, 3)).astype("uint8"))

CPU 次:用户 4.32 秒,系统:425 毫秒,总计:4.75 秒
挂墙时间:4.71 秒

%%time
a= model(np.random.randint(0, 255, size=(1, 320, 320, 3)).astype("uint8"))

CPU 次:用户 124 毫秒,系统:18.4 毫秒,总计:143 毫秒
挂墙时间:85.4 毫秒

在文档中,他们说此模型为 22 毫秒,但也许他们有更快的 GPU。