如何将训练数据集拆分为训练、验证和测试数据集?

She*_*ock 2 machine-learning dataset conv-neural-network pytorch

我有一个自定义的图像数据集及其目标。我在 PyTorch 中创建了一个训练数据集。我想把它分成 3 个部分:训练、验证和测试。我该怎么做?

Sha*_*hai 12

一旦你有了“主”数据集,你就可以data.Subset用来拆分它。
这是随机拆分的示例

import torch
from torch.utils import data
import random

master = data.Dataset( ... )  # your "master" dataset
n = len(master)  # how many total elements you have
n_test = int( n * .05 )  # number of test/val elements
n_train = n - 2 * n_test
idx = list(range(n))  # indices to all elements
random.shuffle(idx)  # in-place shuffle the indices to facilitate random splitting
train_idx = idx[:n_train]
val_idx = idx[n_train:(n_train + n_test)]
test_idx = idx[(n_train + n_test):]

train_set = data.Subset(master, train_idx)
val_set = data.Subset(master, val_idx)
test_set = data.Subset(master, test_idx)
Run Code Online (Sandbox Code Playgroud)

这也可以使用data.random_split

train_set, val_set, test_set = data.random_split(master, (n_train, n_val, n_test))
Run Code Online (Sandbox Code Playgroud)