如何在 Python 中有效地将大型 .txt 文件拆分为训练集和测试集？

Question

我有一个非常大的 .txt 文件（几 GB），我需要将其拆分为机器学习项目的训练集和测试集。由于内存限制，通常的将整个文件读入内存然后分割的方法是不可行的。我正在寻找一种在不超载内存的情况下有效分割文件的方法。

我尝试使用 scikit-learn 进行分割，但它将整个文件加载到内存中，这会导致性能问题，并且不适合我的大型数据集。

Answer 1

我记得当我有一个用于自动驾驶汽车中的物体检测的 12 GB 图像数据集时，我必须将其分成

Train

和

Test

数据集

我使用了下面的代码 只需在 Pycharm 中添加您的路径即可。

import os
import random

def split_large_file(input_file_path, train_file_path, test_file_path, test_ratio=0.2, buffer_size=100000):
    train_file = open(train_file_path, 'w')
    test_file = open(test_file_path, 'w')

    with open(input_file_path, 'r') as input_file:
        buffer = []
        for i, line in enumerate(input_file):
            buffer.append(line)
            if len(buffer) >= buffer_size:
                for buf_line in buffer:
                    if random.random() < test_ratio:
                        test_file.write(buf_line)
                    else:
                        train_file.write(buf_line)
                buffer = []
        
        for buf_line in buffer:
            if random.random() < test_ratio:
                test_file.write(buf_line)
            else:
                train_file.write(buf_line)

    train_file.close()
    test_file.close()

# you can see the below 
input_file_path = 'large_input_file.txt'
train_file_path = 'train_set.txt'
test_file_path = 'test_set.txt'
split_large_file(input_file_path, train_file_path, test_file_path)

如果您遇到任何错误，请告诉我