Scikit 学习 train_test_split 到 Pytorch Dataloader

Question

我有一个用于二进制 classification 的数据集，其 PNG 标题如下面的附件所示，标题中的第一个 0 或 1 决定了它的 class。它们在一个名为“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)

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

Answer 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

Scikit 学习 train_test_split 到 Pytorch Dataloader

Scikit learn train_test_split into Pytorch Dataloader

python

training-data

scikit-learn

pytorch