如何使用pytorch的Dataset, 来定义自己的Dataset

Dataset与DataLoader的关系


  1. Dataset: 构建一个数据集,其中含有所有的数据样本
  2. DataLoader:将构建好的Dataset,通过shuffle、划分batch、多线程num_workers运行的方式,加载到可训练的迭代容器。
python 复制代码
import torch
from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    """创建自己的数据集"""
    def __init__(self):
        """初始化构建数据集所需要的参数"""
        pass

    def __getitem__(self, index):
        """来获取数据集中样本的索引"""
        pass

    def __len__(self):
        """获取数据集中的样本个数"""
        pass

# 实例化自定义的数据集
dataset = MyDataset()
# 将自定义的数据集加载到可训练的迭代容器
train_loader = DataLoader(dataset=dataset,  # 自定义的数据集
                          batch_size=32,  # 数据集中小批量的大小
                          shuffle=True,  # 是否要打乱数据集中样本的次序
                          num_workers=2)  # 是否要并行

实战1:CSV数据集(结构化数据集)

python 复制代码
import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    """创建自己的数据集"""
    def __init__(self, filepath):
        """初始化构建数据集所需要的参数"""
        xy = np.loadtxt(filepath, delimiter=',', dtype=np.float32)
        self.len = xy.shape[0]  # 查看数据集中样本的个数
        self.x_data = torch.from_numpy(xy[:, :-1])
        self.y_data = torch.from_numpy(xy[:, [-1]])
        print("数据已准备好......")

    def __getitem__(self, index):
        """为了支持下标操作, 即索引dataset[index]:来获取数据集中样本的索引"""
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        """为了使用len(dataset):获取数据集中的样本个数"""
        return self.len

file = "D:\\BaiduNetdiskDownload\\Dataset_Dataload\\diabetes1.csv"

""" 1.使用 MyDataset类 构建自己的dataset """
mydataset = MyDataset(file)
""" 2.使用 DataLoader 构建train_loader """
train_loader = DataLoader(dataset=mydataset,
                          batch_size=32,
                          shuffle=True,
                          num_workers=0)

class MyModel(torch.nn.Module):
    """定义自己的模型"""
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(8, 6)
        self.linear2 = torch.nn.Linear(6, 4)
        self.linear3 = torch.nn.Linear(4, 1)
        self.sigmooid = torch.nn.Sigmoid()

    def forward(self, x):
        x = self.sigmooid(self.linear1(x))
        x = self.sigmooid(self.linear2(x))
        x = self.sigmooid(self.linear3(x))
        return x

# 实例化模型
model = MyModel()

# 定义损失函数
criterion = torch.nn.BCELoss(size_average=True)
# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)


if __name__ == "__main__":
    for epoch in range(10):
        for i, data in enumerate(train_loader, 0):
            # 1. 准备数据
            inputs, labels = data

            # 2. 前向传播
            y_pred= model(inputs)
            loss = criterion(y_pred, labels)
            print(epoch, i, loss.item())

            # 3. 反向传播
            optimizer.zero_grad()
            loss.backward()

            # 4. 梯度更新
            optimizer.step()

实战2:图片数据集

├── flower_data

---├── flower_photos(解压的数据集文件夹,3670个样本)

---├── train(生成的训练集,3306个样本)

---└── val(生成的验证集,364个样本)

主函数文件main.py
python 复制代码
import os

import torch
from torchvision import transforms

from my_dataset import MyDataSet
from utils import read_split_data, plot_data_loader_image

root = "../data/flower_data/flower_photos"  # 数据集所在根目录


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    train_images_path, train_images_label, val_images_path, val_images_label = read_split_data(root)

    data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
        "val": transforms.Compose([transforms.Resize(256),
                                   transforms.CenterCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}

    train_data_set = MyDataSet(images_path=train_images_path,
                               images_class=train_images_label,
                               transform=data_transform["train"])

    batch_size = 8
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    print('Using {} dataloader workers'.format(nw))
    train_loader = torch.utils.data.DataLoader(train_data_set,
                                               batch_size=batch_size,
                                               shuffle=True,
                                               num_workers=nw,
                                               collate_fn=train_data_set.collate_fn)

    # plot_data_loader_image(train_loader)

    for epoch in range(100):
	    for step, data in enumerate(train_loader):
	        images, labels = data
	        # 然后在进行相应的训练操作即可


if __name__ == '__main__':
    main()
自定义数据集文件my_dataset.py
python 复制代码
from PIL import Image
import torch
from torch.utils.data import Dataset


class MyDataSet(Dataset):
    """自定义数据集"""

    def __init__(self, images_path: list, images_class: list, transform=None):
        self.images_path = images_path
        self.images_class = images_class
        self.transform = transform

    def __len__(self):
        return len(self.images_path)

    def __getitem__(self, item):
        img = Image.open(self.images_path[item])
        # RGB为彩色图片,L为灰度图片
        if img.mode != 'RGB':
            raise ValueError("image: {} isn't RGB mode.".format(self.images_path[item]))
        label = self.images_class[item]

        if self.transform is not None:
            img = self.transform(img)

        return img, label

    @staticmethod
    def collate_fn(batch):
        # 官方实现的default_collate可以参考
        # https://github.com/pytorch/pytorch/blob/67b7e751e6b5931a9f45274653f4f653a4e6cdf6/torch/utils/data/_utils/collate.py
        images, labels = tuple(zip(*batch))

        images = torch.stack(images, dim=0)
        labels = torch.as_tensor(labels)
        return images, labels
功能文件utils.py(训练集、验证集的划分与可视化)
python 复制代码
import os
import json
import pickle
import random

import matplotlib.pyplot as plt


def read_split_data(root: str, val_rate: float = 0.2):
    random.seed(0)  # 保证随机结果可复现
    assert os.path.exists(root), "dataset root: {} does not exist.".format(root)  # 判断路径是否存在

    # 遍历文件夹,一个文件夹对应一个类别
    flower_class = [cla for cla in os.listdir(root) if os.path.isdir(os.path.join(root, cla))]
    # 排序,保证顺序一致
    flower_class.sort()
    # 生成类别名称以及对应的数字索引: 字典{'花名':0,'花名':1,···}
    class_indices = dict((k, v) for v, k in enumerate(flower_class))
    json_str = json.dumps(dict((val, key) for key, val in class_indices.items()), indent=4)  # 将花名与对应的序号分行保存
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)

    train_images_path = []  # 存储训练集的所有图片路径
    train_images_label = []  # 存储训练集图片对应索引信息
    val_images_path = []  # 存储验证集的所有图片路径
    val_images_label = []  # 存储验证集图片对应索引信息
    every_class_num = []  # 存储每个类别的样本总数
    supported = [".jpg", ".JPG", ".png", ".PNG"]  # 支持的文件后缀类型
    # 遍历每个文件夹下的文件
    for cla in flower_class:
        cla_path = os.path.join(root, cla)
        # 遍历获取supported支持的所有文件路径
        images = [os.path.join(root, cla, i) for i in os.listdir(cla_path)
                  if os.path.splitext(i)[-1] in supported]
        # 获取该类别对应的索引
        image_class = class_indices[cla]
        # 记录该类别的样本数量
        every_class_num.append(len(images))
        # 按比例随机采样验证样本
        val_path = random.sample(images, k=int(len(images) * val_rate))

        for img_path in images:
            if img_path in val_path:  # 如果该路径在采样的验证集样本中则存入验证集
                val_images_path.append(img_path)
                val_images_label.append(image_class)
            else:  # 否则存入训练集
                train_images_path.append(img_path)
                train_images_label.append(image_class)

    print("{} images were found in the dataset.".format(sum(every_class_num)))
    print("{} images for training.".format(len(train_images_path)))
    print("{} images for validation.".format(len(val_images_path)))

    plot_image = True
    if plot_image:
        # 绘制每种类别个数柱状图
        plt.bar(range(len(flower_class)), every_class_num, align='center')
        # 将横坐标0,1,2,3,4替换为相应的类别名称
        plt.xticks(range(len(flower_class)), flower_class)
        # 在柱状图上添加数值标签
        for i, v in enumerate(every_class_num):
            plt.text(x=i, y=v + 5, s=str(v), ha='center')
        # 设置x坐标
        plt.xlabel('image class')
        # 设置y坐标
        plt.ylabel('number of images')
        # 设置柱状图的标题
        plt.title('flower class distribution')
        plt.show()

    return train_images_path, train_images_label, val_images_path, val_images_label


def plot_data_loader_image(data_loader):
    batch_size = data_loader.batch_size
    plot_num = min(batch_size, 4)

    json_path = './class_indices.json'
    assert os.path.exists(json_path), json_path + " does not exist."
    json_file = open(json_path, 'r')
    class_indices = json.load(json_file)

    for data in data_loader:
        images, labels = data
        for i in range(plot_num):
            # [C, H, W] -> [H, W, C]
            img = images[i].numpy().transpose(1, 2, 0)
            # 反Normalize操作
            img = (img * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]) * 255
            label = labels[i].item()
            plt.subplot(1, plot_num, i+1)
            plt.xlabel(class_indices[str(label)])
            plt.xticks([])  # 去掉x轴的刻度
            plt.yticks([])  # 去掉y轴的刻度
            plt.imshow(img.astype('uint8'))
        plt.show()


def write_pickle(list_info: list, file_name: str):
    with open(file_name, 'wb') as f:
        pickle.dump(list_info, f)


def read_pickle(file_name: str) -> list:
    with open(file_name, 'rb') as f:
        info_list = pickle.load(f)
        return info_list
相关推荐
开出南方的花2 分钟前
DeepSeek模型架构及优化内容
人工智能·pytorch·深度学习·机器学习·架构·nlp·attention
千层冷面5 分钟前
Java Stream 全面解析
java·windows·python
0wioiw06 分钟前
Android逆向(Python和JS通信)
android·javascript·python
goomind28 分钟前
YOLOv11实战海洋动物图像识别
pytorch·深度学习·yolo·计算机视觉·cnn·海洋动物
杀死这个程序员40 分钟前
conda介绍及常用命令举例
人工智能·conda
三皮仔41 分钟前
安装WPS后,导致python调用Excel.Application异常,解决办法
python·excel·wps
说私域1 小时前
互联网企业线上业务拓展与开源AI智能名片2+1链动模式S2B2C商城小程序的创新应用
人工智能·小程序·开源
java_python源码1 小时前
[含文档+PPT+源码等]精品基于Python实现的django个性化健康餐计划订制系统
开发语言·前端·python
Orange_sparkle1 小时前
大模型基本原理(四)——如何武装ChatGPT
人工智能·语言模型·chatgpt
时间很奇妙!1 小时前
NLP & Word Embeddings
人工智能·自然语言处理·word