划分训练集 测试集 验证集 并生成txt文件
时间: 2023-09-20 19:05:11 浏览: 97
划分训练集、测试集和验证集是机器学习中常用的一步。一般情况下,我们需要将数据集划分成三个部分:
- 训练集(training set):用于训练模型。
- 验证集(validation set):用于调整模型的超参数,比如学习率、正则化系数等。
- 测试集(test set):用于评估模型的泛化性能,也就是模型对未见过的数据的预测能力。
下面是一个 Python 代码示例,用于将数据集划分为训练集、验证集和测试集,并生成对应的 txt 文件。
```python
import os
import random
# 数据集路径
data_dir = 'data/'
# 分割比例
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1
# 获取文件列表
files = os.listdir(data_dir)
file_list = []
for file in files:
file_list.append(os.path.join(data_dir, file))
# 打乱文件列表
random.shuffle(file_list)
# 计算分割点
train_point = int(len(file_list) * train_ratio)
val_point = int(len(file_list) * (train_ratio + val_ratio))
# 划分数据集
train_files = file_list[:train_point]
val_files = file_list[train_point:val_point]
test_files = file_list[val_point:]
# 写入 txt 文件
with open('train.txt', 'w') as f:
f.write('\n'.join(train_files))
with open('val.txt', 'w') as f:
f.write('\n'.join(val_files))
with open('test.txt', 'w') as f:
f.write('\n'.join(test_files))
```
上述代码中,我们首先定义了数据集路径、分割比例以及三个空列表,用于存储划分后的训练集、验证集和测试集。然后,我们通过 `os.listdir()` 函数获取数据集中所有文件的路径,将它们存储在一个列表中,并且使用 `random.shuffle()` 函数将列表打乱。
接下来,我们根据分割比例计算出训练集、验证集和测试集的分割点,使用 Python 的切片功能将文件列表划分为三个部分,并将它们分别存储在 `train_files`、`val_files` 和 `test_files` 中。
最后,我们使用 `with open()` 函数分别打开三个 txt 文件,并使用 `f.write()` 函数将文件列表写入相应的文件中。