Scikit 学习 train_test_split 成 Pytorch Dataloader答案

【问题标题】：Scikit learn train_test_split into Pytorch DataloaderScikit 学习 train_test_split 成 Pytorch Dataloader
【发布时间】：2021-11-14 21:32:49
【问题描述】：

我有一个带有 PNG 的二进制分类数据集，标题如下面的附件所示，其中标题中的第一个 0 或 1 决定了它的类别。它们位于一个名为“annotation_class”的文件夹中，我有一个小脚本将它们分开：

import cv2,glob
import numpy as np
from sklearn.model_selection import train_test_split

filelist = glob.glob('annotation_class'+'/*.png')
size_row, size_col = 256, 256
X,y = [],[]

for name in filelist:
        img = cv2.imread(name) 
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 
        img = cv2.resize(img,(size_row, size_col))
        X.append(img)
        y.append(int(name.split('\\')[-1].split('_')[1]))


x_train, x_test, y_train, y_test= train_test_split(X, y, test_size=0.2, train_size=0.8, random_state=4)

返回的都是列表。我在这个项目中使用 Pytorch，并且想制作一个自定义数据集来使用 Dataloader，但我不确定在使用 train_test_split 之后如何最好地包含这些。我应该完全放弃它并使用其他东西吗？我想最终得到两个 DataLoader 用于训练和测试。

【问题讨论】：

标签： python scikit-learn pytorch training-data

【解决方案1】：

您不必重写。您可以在 PyTorch 数据集中重用核心数据加载逻辑

import cv2,glob
import numpy as np
from sklearn.model_selection import train_test_split

from torch.utils.data import Dataset

class MyCoolDataset(Dataset):

    def __init__(self, dir, train=True):
        filelist = glob.glob(dir + '/*.png')
        ...
        # all your data loading logic using cv2, glob ..
        x_train, x_test, y_train, y_test = train_test_split(X, y, test_size=0.2, train_size=0.8, random_state=4)
        
        # two modes - train and test
        if train:
            self.x_data, self.y_data = x_train, y_train
        else:
            self.x_data, self.y_data = x_test, y_test
    
    def __getitem__(self, i):
        return self.x_data[i], self.y_data[i]

然后像往常一样使用DataLoader

dl = DataLoader(MyCoolDataset(...), batch_size=...)
for X, Y in dl:
    pass

【讨论】：