使用tf.data处理文本数据时遭遇严重的速度与内存性能问题

你遇到的核心问题是直接用tf.data.Dataset.from_tensor_slices处理原始字节流导致了极高的内存消耗和性能瓶颈，这也是Colab环境崩溃的核心原因。

问题根源拆解 #

原代码中f.read(f.namelist()[0])读取的是整个text8文件的原始字节串（大小约30MB），当你直接把这个字节串传给from_tensor_slices时：

• TensorFlow会把这个字节串拆分成单个字符的数据集，也就是数百万个单独的字符元素，这完全不是你需要的单词级数据。
• 即便你加上tf.string_split的map操作，这种处理方式也会因为要在极大量的字符元素上重复执行拆分逻辑，导致计算量爆炸，直接耗尽Colab的内存和计算资源。

正确的tf.data处理方式 #

我们需要先把原始文本转换成单词列表，再构建tf.data数据集；或者直接用tf.data的字符串处理API完成从文本到单词的转换，避免在Python层面处理大文本后再转TensorFlow数据集的低效操作。

下面是修正后的完整代码，既保留tf.data的优势，又能和原代码一样高效运行：

from __future__ import print_function
import collections
import math
import numpy as np
import os
import random
import tensorflow as tf
import zipfile
from matplotlib import pylab
from six.moves import range
from six.moves.urllib.request import urlretrieve
from sklearn.manifold import TSNE

url = 'http://mattmahoney.net/dc/'

def maybe_download(filename, expected_bytes):
  """Download a file if not present, and make sure it's the right size."""
  if not os.path.exists(filename):
    filename, _ = urlretrieve(url + filename, filename)
  statinfo = os.stat(filename)
  if statinfo.st_size == expected_bytes:
    print('Found and verified %s' % filename)
  else:
    print(statinfo.st_size)
    raise Exception(
      'Failed to verify ' + filename + '. Can you get to it with a browser?')
  return filename

filename = maybe_download('text8.zip', 31344016)

def read_data_with_tfdata(filename):
  """用tf.data处理zip中的文本，转换成单词数据集"""
  # 从zip文件中读取文本内容
  with zipfile.ZipFile(filename) as f:
    text_content = tf.compat.as_str(f.read(f.namelist()[0]))
  
  # 先把文本按空格拆分成单词列表，再构建数据集
  word_list = text_content.split()
  dataset = tf.data.Dataset.from_tensor_slices(word_list)
  
  # 后续可添加预处理操作（如过滤低频词、生成训练样本）
  return dataset

# 获取单词数据集
word_dataset = read_data_with_tfdata(filename)
# 验证数据集大小（和原代码一致）
print('Data size %d' % len(list(word_dataset.as_numpy_iterator())))

进阶优化：纯tf.data管道（避免Python层面拆分） #

如果想完全用TensorFlow API完成从文件读取到单词拆分的全流程，避免Python处理大文本，可以这样实现：

def read_data_pure_tfdata(filename):
  # 读取zip中的文件内容为字符串张量
  with zipfile.ZipFile(filename) as f:
    text_tensor = tf.convert_to_tensor(tf.compat.as_str(f.read(f.namelist()[0])), dtype=tf.string)
  
  # 按空格拆分字符串为单词，再转成数据集
  words_tensor = tf.strings.split(text_tensor)
  dataset = tf.data.Dataset.from_tensor_slices(words_tensor)
  return dataset

word_dataset = read_data_pure_tfdata(filename)
print('Data size %d' % len(list(word_dataset.as_numpy_iterator())))

关键注意点 #

• 永远不要直接用from_tensor_slices处理原始字节串：这会生成海量单字符元素，直接引发性能灾难。
• 先做高层拆分再构建数据集：无论是在Python层面先拆成单词，还是用TensorFlow的tf.strings.split先得到单词张量，再构建数据集，才能保证数据集的元素是你需要的单词级别。
• tf.data的优势在于批量处理和流水线：后续生成word2vec的训练样本（比如skip-gram的上下文对）时，再用tf.data的window、flat_map等操作，才能真正发挥它的性能优势。

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