如何在不写入文件系统的情况下从 Google 存储桶恢复 Tensorflow 模型?

问题描述:

我有一个 2gb 的 Tensorflow 模型,我想将它添加到我在 App Engine 上的 Flask 项目中,但我似乎找不到任何说明我正在尝试做的事情是可能的文档.

I have a 2gb Tensorflow model that I'd like to add to a Flask project I have on App Engine but I can't seem to find any documentation stating what I'm trying to do is possible.

由于 App Engine 不允许写入文件系统,我将模型的文件存储在 Google Bucket 中并尝试从那里恢复模型.这些是那里的文件:

Since App Engine doesn't allow writing to the file system, I'm storing my model's files in a Google Bucket and attempting to restore the model from there. These are the files there:

  • model.ckpt.data-00000-of-00001
  • model.ckpt.index
  • model.ckpt.meta
  • 检查站

在本地工作,我可以使用

Working locally, I can just use

with tf.Session() as sess:
    logger.info("Importing model into TF")
    saver = tf.train.import_meta_graph('model.ckpt.meta')
    saver.restore(sess, model.ckpt)

使用 Flask 的 @before_first_request 将模型加载到内存中.

Where the model is loaded into memory using Flask's @before_first_request.

一旦它在 App Engine 上,我以为我可以这样做:

Once it's on App Engine, I assumed I could to this:

blob = bucket.get_blob('blob_name')
filename = os.path.join(model_dir, blob.name)
blob.download_to_filename(filename)

然后做同样的恢复.但 App Engine 不允许.

Then do the same restore. But App Engine won't allow it.

有没有办法将这些文件流式传输到 Tensorflow 的恢复功能中,这样文件就不必写入文件系统?

Is there a way to stream these files into Tensorflow's restore functions so the files don't have to be written to the file system?

经过 Dan Cornilescu 的一些提示并深入研究后,我发现 Tensorflow 使用名为 ParseFromStringMetaGraphDef/code>,所以这是我最终做的:

After some tips from Dan Cornilescu and digging into it I found that Tensorflow builds the MetaGraphDef with a function called ParseFromString, so here's what I ended up doing:

from google.cloud import storage
from tensorflow import MetaGraphDef

client = storage.Client()
bucket = client.get_bucket(Config.MODEL_BUCKET)
blob = bucket.get_blob('model.ckpt.meta')
model_graph = blob.download_as_string()

mgd = MetaGraphDef()
mgd.ParseFromString(model_graph)

with tf.Session() as sess:
    saver = tf.train.import_meta_graph(mgd)