Universal Sentence Encoder CPU内存耗尽(OOM)问题排查与解决求助

我来帮你分析这个OOM问题的根源和解决办法：

问题根源分析 #

1. 实际可用内存差异：你本地Ubuntu的3GB内存是实打实的用户可用内存，而CentOS VPS的4GB内存可能被系统服务、后台进程占用了不少，实际留给TensorFlow的内存余量其实比本地还小。
2. 老旧内核的内存管理缺陷：你的VPS用的是Linux 2.6.32这个非常老旧的内核，它的内存碎片处理效率远低于新内核，当需要分配[26667,320]这样的大张量时，即使总内存够，也可能因为内存碎片找不到连续的内存块而触发OOM。
3. 模型加载的内存开销：Universal Sentence Encoder v4本身加载时需要消耗大量内存，加上TensorFlow和Python解释器的基础内存占用，VPS的内存刚好卡在临界值，直接触发了资源耗尽错误。

具体解决方案 #

1. 优化系统内存配置 #

• 增加Swap分区：给VPS添加swap空间可以缓解内存不足的问题，操作步骤如下：

  # 创建2GB的swap文件（可根据需求调整大小）
  sudo fallocate -l 2G /swapfile
  sudo chmod 600 /swapfile  # 设置正确的权限
  sudo mkswap /swapfile
  sudo swapon /swapfile
  

  若要永久生效，将以下内容添加到/etc/fstab文件末尾：

  /swapfile swap swap defaults 0 0
  

• 清理无用进程：用htop或top命令查看占用内存较高的无用进程，用kill命令关闭它们，释放更多内存给TensorFlow。

2. 调整TensorFlow内存分配策略 #

修改代码，限制TensorFlow的内存使用，避免一次性占用过多内存：

import tensorflow_hub as hub
import tensorflow.compat.v1 as tf
tf.disable_eager_execution()

# 配置TensorFlow的内存分配策略
config = tf.ConfigProto()
# 允许内存动态增长，而不是启动时就占用全部可用内存
config.gpu_options.allow_growth = True
# 限制线程数，减少内存开销（CPU环境下有效）
config.intra_op_parallelism_threads = 2
config.inter_op_parallelism_threads = 2

with tf.Session(config=config) as sess:
    embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
    sess.run([tf.global_variables_initializer(), tf.tables_initializer()])
    result = sess.run(embed(["test message"]))
    print(result)

3. 使用轻量版模型 #

如果上述方法还是不行，可以改用Universal Sentence Encoder Lite，它的内存占用远低于标准版，虽然精度略有下降，但大部分场景下足够使用：

import tensorflow_hub as hub
import tensorflow.compat.v1 as tf
tf.disable_eager_execution()

with tf.Session() as sess:
    # 加载轻量版模型
    embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder-lite/2")
    sess.run([tf.global_variables_initializer(), tf.tables_initializer()])
    # 轻量版需要先预处理文本
    preprocessed_text = sess.run(embed.signatures['preprocess'](["test message"]))
    # 生成文本嵌入
    embedding_result = sess.run(embed.signatures['encode'](preprocessed_text))
    print(embedding_result)

4. 升级系统内核（可选） #

如果你的VPS服务商允许升级内核，可以把CentOS的内核升级到3.x及以上版本。新内核的内存碎片管理更高效，能从根本上减少这类OOM问题的发生，但升级需要重启VPS，请注意业务影响。

内容的提问来源于stack exchange，提问作者mmak