帮我写一个按8:1:1分配数据集的脚本
时间: 2024-03-02 14:48:40 浏览: 30
好的,这里是一个简单的 Python 脚本,可以按照 8:1:1 的比例将数据集分成训练集、验证集和测试集:
```python
import os
import random
import shutil
# 设置随机种子,确保每次运行结果一致
random.seed(42)
# 设置数据集路径和目标路径
data_dir = "/path/to/data"
train_dir = "/path/to/train"
valid_dir = "/path/to/valid"
test_dir = "/path/to/test"
# 获取所有数据文件的路径和标签
data_files = []
for label in os.listdir(data_dir):
label_dir = os.path.join(data_dir, label)
for file_name in os.listdir(label_dir):
file_path = os.path.join(label_dir, file_name)
data_files.append((file_path, label))
# 随机打乱数据
random.shuffle(data_files)
# 计算数据集大小和分配比例
data_size = len(data_files)
train_size = int(data_size * 0.8)
valid_size = int(data_size * 0.1)
test_size = data_size - train_size - valid_size
# 按分配比例拷贝数据到目标路径
for i, (file_path, label) in enumerate(data_files):
if i < train_size:
dest_dir = train_dir
elif i < train_size + valid_size:
dest_dir = valid_dir
else:
dest_dir = test_dir
dest_label_dir = os.path.join(dest_dir, label)
os.makedirs(dest_label_dir, exist_ok=True)
dest_file_path = os.path.join(dest_label_dir, os.path.basename(file_path))
shutil.copy(file_path, dest_file_path)
```
这个脚本假设数据集的目录结构如下:
```
data/
├── label1/
│ ├── file1.txt
│ ├── file2.txt
│ └── ...
├── label2/
│ ├── file1.txt
│ ├── file2.txt
│ └── ...
└── ...
```
它会将数据集按照 8:1:1 的比例拷贝到训练集、验证集和测试集的目录下,目录结构如下:
```
train/
├── label1/
│ ├── file1.txt
│ ├── file2.txt
│ └── ...
├── label2/
│ ├── file1.txt
│ ├── file2.txt
│ └── ...
└── ...
valid/
├── label1/
│ ├── file1.txt
│ ├── file2.txt
│ └── ...
├── label2/
│ ├── file1.txt
│ ├── file2.txt
│ └── ...
└── ...
test/
├── label1/
│ ├── file1.txt
│ ├── file2.txt
│ └── ...
├── label2/
│ ├── file1.txt
│ ├── file2.txt
│ └── ...
└── ...
```