如何在Python中高效拆分大.txt文件为训练集与测试集？

针对数GB级的TXT文件，无法全量加载到内存的场景，以下是几种低内存占用的高效拆分方案，替代需要加载全量数据的scikit-learn拆分方法：

方法1：逐行随机抽样（内存占用极低） #

全程仅在内存中保留单行数据，通过随机概率分配每行到训练集或测试集，实现近似比例的拆分。适合对比例精度要求不严格的场景，且仅需遍历文件一次。

import random

# 设置随机种子，确保拆分结果可复现
random.seed(42)
# 设定训练集占比，示例为80%
train_ratio = 0.8

# 同时打开原文件和两个输出文件
with open('large_data.txt', 'r') as src_file, \
     open('train.txt', 'w') as train_file, \
     open('test.txt', 'w') as test_file:
    
    for line in src_file:
        # 可选：跳过空行，避免无效数据
        stripped_line = line.strip()
        if not stripped_line:
            continue
        # 按随机概率分配行到对应文件
        if random.random() < train_ratio:
            train_file.write(line)
        else:
            test_file.write(line)

方法2：精确比例拆分（需两次遍历文件） #

如果需要严格符合设定的拆分比例，先统计文件总行数，再生成随机索引确定测试集行，第二次遍历文件时完成拆分。内存占用依然极低，仅需存储索引集合。

import random

random.seed(42)
train_ratio = 0.8

# 第一步：统计有效行数（跳过空行）
total_valid_lines = 0
with open('large_data.txt', 'r') as f:
    for line in f:
        if line.strip():
            total_valid_lines += 1

# 计算训练集/测试集行数
train_lines = int(total_valid_lines * train_ratio)
test_line_count = total_valid_lines - train_lines
# 生成测试集的行索引（避免重复）
test_indices = set(random.sample(range(total_valid_lines), test_line_count))

# 第二步：执行拆分
current_line_idx = 0
with open('large_data.txt', 'r') as src_file, \
     open('train.txt', 'w') as train_file, \
     open('test.txt', 'w') as test_file:
    
    for line in src_file:
        stripped_line = line.strip()
        if not stripped_line:
            continue
        # 根据索引判断写入目标文件
        if current_line_idx in test_indices:
            test_file.write(line)
        else:
            train_file.write(line)
        current_line_idx += 1

方法3：Shell命令快速拆分（Linux/macOS适用） #

无需编写代码，利用系统命令高效处理，全程不加载全量数据到内存。适合熟悉命令行的用户：

# 可选：先打乱原文件（若数据有序，避免拆分后分布不均）
shuf large_data.txt -o shuffled_data.txt

# 统计总行数
total_lines=$(wc -l < shuffled_data.txt)
# 计算训练集行数（80%比例）
train_lines=$((total_lines * 8 / 10))

# 拆分：前80%到训练集，剩余到测试集
head -n $train_lines shuffled_data.txt > train.txt
tail -n +$((train_lines + 1)) shuffled_data.txt > test.txt

注意事项 #

• 始终避免使用readlines()、pandas.read_csv()（默认加载全量）等会把整个文件读入内存的方法
• 如果文件包含表头（如CSV格式的TXT），需单独处理表头：先读取第一行，同时写入训练集和测试集的开头，再处理后续数据行
• 必须设置随机种子，保证拆分结果可复现，便于后续实验验证

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