pytorch划分数据集的函数

PyTorch提供了一些函数来帮助划分数据集，其中常用的有`train_test_split`函数和`Subset`类。下面是它们的简单介绍：

1. `train_test_split`函数：这个函数可以根据一定的比例将数据集划分为训练集和测试集。它可以在`torch.utils.data.dataset.random_split`模块中找到。示例如下：

from torch.utils.data import random_split
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

其中，`dataset`是你的原始数据集，`train_size`和`test_size`是训练集和测试集的大小比例（可以是整数或小数）。

2. `Subset`类：这个类可以根据给定的索引划分数据集。它可以在`torch.utils.data.dataset.subset`模块中找到。示例如下：

from torch.utils.data import Subset
train_dataset = Subset(dataset, train_indices)
test_dataset = Subset(dataset, test_indices)

其中，`dataset`是你的原始数据集，`train_indices`和`test_indices`是训练集和测试集的索引列表。

这些函数和类可以用于任何PyTorch数据集，例如`torchvision.datasets.ImageFolder`或自定义数据集。划分数据集后，你可以使用这些数据集来构建数据加载器（DataLoader）进行训练和测试。