Pytorch加载数据集-小土堆

1.Dataset 和 Dataloader

 

Dataset主要是提供一种方式去获取数据及其label。

 

Dataloader主要作用是为后面的网络提供不同的数据形式。

2.Dataset的使用

蚂蚁蜜蜂分类数据集和下载连接https://download.pytorch.org/tutorial/hymenoptera_data.zip

from torch.utils.data import Dataset
from PIL import Image
import cv2
import os
class MyData(Dataset):
    def __init__(self,root_dir,label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir,self.label_dir)
        self.img_path = os.listdir(self.path)


    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)
        img = Image.open(img_item_path)
        label = self.label_dir
        return img,label
    def __len__(self):
        return len(self.img_path)

root_dir = "hymenoptera_data/train"
ants_label_dir = "ants"
bees_label_dir = "bees"
bees_dataset = MyData(root_dir,bees_label_dir)
ants_dataset = MyData(root_dir,ants_label_dir)
train_dataset = ants_dataset + bees_dataset  #将数据集合并