Pytorch如何划分训练集和测试集

如果我们使用scikit-learn库，可以使用如下方法划分训练集和测试集：

from sklearn.model_selection import train_test_split

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2)

如果我们不使用该库，在pytorch中还有一种可以划分训练接和测试集的方法。如下：

import numpy as np
from torch.utils.data import DataLoader, SubsetRandomSampler

# 首先产生数据索引的乱序排列
shuffled_indices = np.random.permutation(len(data))
train_idx = shuffled_indices[:int(0.8*len(data))]
val_idx = shulled_indices[int(0.8*len(data)):]

train_loader = DataLoader(data, batch_size = 32, drop_last = True, sampler = SubserRandomSampler(train_idx))
val_loader = DataLoader(data, batch_size = 32, drop_last=False, sampler=SubsetRandomSampler(val_idx))

Pytorch中也有一个类似于Numpy中random.permutation的函数：torch.randperm。

如果我们并不想使用所有的数据进行训练，需要从中随机挑选一部分，比如从中挑选1000个样本进行训练，则可以使用如下方法：

import random

train_idx = random.sample(len(data), 1000)
train_loader = DataLoader(data, batch_size = 32, drop_last = True, sampler = SubserRandomSampler(train_idx))

这里的SuberRandomSampler是一种子集随机采样方式，传入的是data的随机索引。更多的sampler使用方法可以参考docs

References:

Pytorch Docs: torch.utils.data