pytorch 数据集划分

在PyTorch中，数据集的划分可以通过使用 SubsetRandomSampler 或者 DataLoader 中的参数进行实现。

1. 使用 SubsetRandomSampler

SubsetRandomSampler 可以用来随机划分数据集。首先，我们需要定义一个数据集，并且定义训练集和测试集的索引。然后，我们使用 SubsetRandomSampler 将训练集和测试集的索引传递给 DataLoader，以便在训练和测试期间使用不同的数据集。

import torch
from torch.utils.data import DataLoader, SubsetRandomSampler
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

# 定义数据集
mnist = MNIST(root='data', download=True, transform=ToTensor())

# 定义训练集和测试集的索引
train_indices = list(range(0, 50000))
test_indices = list(range(50000, 60000))

# 定义训练集和测试集的 Sampler
train_sampler = SubsetRandomSampler(train_indices)
test_sampler = SubsetRandomSampler(test_indices)

# 定义 DataLoader
train_loader = DataLoader(mnist, batch_size=64, sampler=train_sampler)
test_loader = DataLoader(mnist, batch_size=64, sampler=test_sampler)

2. 使用 DataLoader 中的参数

在 DataLoader 中，我们可以使用 shuffle 参数来随机化数据集，并且使用 num_workers 参数来设置使用的 CPU 核心数。同时，我们也可以使用 drop_last 参数来控制是否删除最后一个大小小于 batch_size 的 batch。

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

# 定义数据集
mnist = MNIST(root='data', download=True, transform=ToTensor())

# 定义 DataLoader
train_loader = DataLoader(mnist, batch_size=64, shuffle=True, num_workers=4, drop_last=True)
test_loader = DataLoader(mnist, batch_size=64, shuffle=False, num_workers=4, drop_last=True)

使用上述方法可以很方便地划分数据集，并且在训练和测试期间使用不同的数据集。