import random

# 读取数据集
data = []
with open('data.txt', 'r') as file:
    lines = file.readlines()
    for line in lines:
        line = line.strip().split(',')
        line = [float(x) for x in line]
        data.append(line)

# 分割数据集
random.shuffle(data)
train_data = []
valid_data = []
test_data = []
class_count = [0] * 8
for line in data:
    label = int(line[-1])
    if class_count[label] < 1000:
        train_data.append(line)
        class_count[label] += 1
    elif class_count[label] < 1100:
        valid_data.append(line)
        class_count[label] += 1
    elif class_count[label] < 1200:
        test_data.append(line)
        class_count[label] += 1

# 打乱数据集
random.shuffle(train_data)
random.shuffle(valid_data)
random.shuffle(test_data)

# 保存数据集
with open('train.txt', 'w') as file:
    for line in train_data:
        line = [str(x) for x in line]
        file.write(','.join(line) + '\n')

with open('valid.txt', 'w') as file:
    for line in valid_data:
        line = [str(x) for x in line]
        file.write(','.join(line) + '\n')

with open('test.txt', 'w') as file:
    for line in test_data:
        line = [str(x) for x in line]
        file.write(','.join(line) + '\n')
数据分割:将txt数据集划分成训练集、验证集和测试集

原文地址: https://www.cveoy.top/t/topic/jKwP 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录