pytorch学习（随手记）

使用 PyTorch 的 Dataset 类来创建一个自定义的数据集类 MyData，用于加载图像和对应的标签。

首先提前下载好数据集添加到目录

1. 初始化 (__init__ 方法):

       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)

2. 获取单个数据项 (__getitem__ 方法):

       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

完整代码展示： 

from torch.utils.data import Dataset
from PIL import Image
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="dataset/train"
ants_label_dir = "ants"
bees_label_dir = "bees"
ants_dataset = MyData(root_dir , ants_label_dir)
bees_dataset = MyData(root_dir , bees_label_dir)