tf.keras.models.load_model由于内存不足而失败

Question

我正在尝试从TensorFlow Hub加载模块，将其导出到磁盘，从磁盘重新加载并尝试使用它：

import tensorflow as tf
import tensorflow_hub as hub


def save_module(url, save_path):
    module1 = hub.KerasLayer(url)
    model1 = tf.keras.Sequential(module1)
    print("model1: ", model1(["hello", "world"]))
    tf.keras.models.save_model(module1, save_path)
    module2 = tf.keras.models.load_model(save_path)
    model2 = tf.keras.Sequential(module2)
    print("model2: ", model2(["hello", "world"]))


save_module("https://tfhub.dev/google/universal-sentence-encoder/4", "./saved-module")

由于某些原因，load_module失败的原因是：

2020-05-17 12:52:05.880735: W tensorflow/core/common_runtime/bfc_allocator.cc:434] Allocator (GPU_0_bfc) ran out of memory trying to allocate 32.55MiB (rounded to 34133760)

如果我将保存和加载移动到单独的运行中，则不会耗尽内存。

我的代码有什么问题？我是否忘记在故障之前释放一些资源？如果是这样，我该如何解除分配？