如何解决加载TensorFlow Hub模型时的ResourceExhaustedError(OOM)问题？

阿华AIGC实验室

2026-5-14

解决TensorFlow Hub NNLM模型加载时的OOM（资源耗尽）错误

遇到这个CPU内存不足的OOM问题很正常——你用的NNLM-en-dim128模型需要加载一个包含97万+词汇的128维嵌入矩阵，这个矩阵大小约480MB，要是你的CPU剩余内存不够容纳它，就会触发这个错误。下面是几个实用的解决办法：

1. 切换到更轻量的预训练嵌入模型

这是最直接的方案，选择维度更小或者词汇表更精简的模型，内存占用会大幅降低。比如：

使用50维的NNLM模型：https://hub.tensorflow.google.cn/google/tf2-preview/nnlm-en-dim50/1
或者超轻量的20维gnews模型：https://hub.tensorflow.google.cn/google/tf2-preview/gnews-swivel-20dim/1

修改后的代码示例：

import tensorflow as tf
import tensorflow_hub as hub
# 替换为轻量模型URL
hub_url = "https://hub.tensorflow.google.cn/google/tf2-preview/gnews-swivel-20dim/1"
embed = hub.KerasLayer(hub_url)
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)

2. 将模型加载到GPU（如果有可用GPU）

如果你的机器有GPU，把模型和张量分配到GPU显存上可以缓解CPU内存压力。在代码开头添加GPU配置：

import tensorflow as tf
import tensorflow_hub as hub

# 配置TensorFlow使用GPU
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # 设置只使用第一个GPU
        tf.config.set_visible_devices(gpus[0], 'GPU')
        print(f"成功使用GPU: {gpus[0]}")
    except RuntimeError as e:
        print(f"GPU配置失败: {e}")

hub_url = "https://hub.tensorflow.google.cn/google/tf2-preview/nnlm-en-dim128/1"
embed = hub.KerasLayer(hub_url)
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)

3. 清理内存并优化TensorFlow配置

在加载模型前清理之前的TensorFlow会话和未使用的内存，同时开启内存优化：

import tensorflow as tf
import tensorflow_hub as hub
import gc

# 清理内存
tf.keras.backend.clear_session()
gc.collect()

# 开启GPU内存增长模式（避免一次性占用过多显存）
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

hub_url = "https://hub.tensorflow.google.cn/google/tf2-preview/nnlm-en-dim128/1"
embed = hub.KerasLayer(hub_url)
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)