我正在尝试从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)
如果我将保存和加载移动到单独的运行中,则不会耗尽内存。
我的代码有什么问题?我是否忘记在故障之前释放一些资源?如果是这样,我该如何解除分配?