代码实现对数据集进行划分
时间: 2023-09-16 21:08:48 浏览: 33
对数据集进行划分是为了在训练模型时,将数据集分成训练集、验证集和测试集。常见的划分比例是70%的数据用于训练,10%的数据用于验证,20%的数据用于测试。
以下是Python代码实现对数据集进行划分的示例:
```python
import os
import random
import shutil
# 数据集路径
data_dir = './data'
# 划分比例
train_ratio = 0.7
val_ratio = 0.1
test_ratio = 0.2
# 创建保存划分后数据集的目录
split_dir = './split_data'
if not os.path.exists(split_dir):
os.makedirs(split_dir)
# 获取数据集中所有文件名
all_files = os.listdir(data_dir)
# 打乱文件名的顺序
random.shuffle(all_files)
# 计算划分后每个数据集的数量
num_all_files = len(all_files)
num_train_files = int(num_all_files * train_ratio)
num_val_files = int(num_all_files * val_ratio)
num_test_files = num_all_files - num_train_files - num_val_files
# 拷贝文件到划分后的数据集目录中
for i, file_name in enumerate(all_files):
if i < num_train_files:
split_file_dir = os.path.join(split_dir, 'train')
elif i < num_train_files + num_val_files:
split_file_dir = os.path.join(split_dir, 'val')
else:
split_file_dir = os.path.join(split_dir, 'test')
if not os.path.exists(split_file_dir):
os.makedirs(split_file_dir)
src_file_path = os.path.join(data_dir, file_name)
dst_file_path = os.path.join(split_file_dir, file_name)
shutil.copyfile(src_file_path, dst_file_path)
```
上述代码中,首先指定了数据集的路径`data_dir`,以及划分比例`train_ratio`、`val_ratio`和`test_ratio`。然后创建了保存划分后数据集的目录`split_dir`。接着获取了数据集中所有文件名,并打乱文件名的顺序。接下来计算了划分后每个数据集的数量,并将文件拷贝到划分后的数据集目录中。最终得到了训练集、验证集和测试集三个数据集。