pytorch入门（1）——pytorch加载数据初认识

Familyism2024-09-16 20:05

环境配置及其安装：

2023最新pytorch安装（超详细版）-CSDN博客

pytorch加载数据初认识

Dataset：创建可被Pytorch使用的数据集

提供一种方式获取数据及其label

Dataloader：向模型传递数据

为网络提供不同的数据形式

代码实战

蚂蚁蜜蜂/练手数据集：链接: https://pan.baidu.com/s/1jZoTmoFzaTLWh4lKBHVbEA 密码: 5suq

train为训练集，val为测试集，可以看到ants里蚂蚁的图片，bees里是蜜蜂的图片

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)
        # 函数是os模块中的一个重要函数，用于获取指定目录下的所有文件和子目录。
        self.image_path = os.listdir(self.path)

    def __getitem__(self, idx):
        img_name = self.image_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.image_path)

root_dir = "hymenoptera_data/train"
ants_label_dir = "ants"
bees_label_dir = "ants"
ants_dataset = MyData(root_dir, ants_label_dir)
bees_dataset = MyData(root_dir, bees_label_dir)
train_dataset = ants_dataset + bees_dataset