Dataset类
作用：模型的数据集接口
__init__将对象实例化，创建对象时obj = class(..., ...)会立即被调用，需要提供（输入）类中使用到的变量。
__getitem__通过img, label = obj[idx]获取（返回）每一个数据和label
__len__通过len(obj)获取（返回）数据量
点击查看代码
from torch.utils.data import Dataset
import os
from PIL import Image
# 数据集下载地址：https://download.pytorch.org/tutorial/hymenoptera_data.zip
class Mydata(Dataset):
def __init__(self, root_dir, label_dir):
# self为全局变量
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) # 获取path路径下所有子文件的名称
def __getitem__(self, idx):
img_name = self.img_path[idx] # 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 = r'D:\ai-learning\pytorch\hymenoptera_data\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)
train_dataset = ants_dataset + bees_dataset
img, label = train_dataset[0]
img.show()
len(train_dataset)
一个可以参考的通用数据集接口：https://blog.csdn.net/leviopku/article/details/99958182
参考：
[1]https://www.cnblogs.com/seansheep/p/16163159.html
[2]https://www.bilibili.com/video/BV1hE411t7RN?p=7&vd_source=fa1d778abbb911d02be7ac36f2b2e32a