快捷導(dǎo)航

PyTorch?Dataset與DataLoader使用超詳細(xì)講解

更新時(shí)間：2022年10月13日 14:09:20 作者：竹清蘭香

用于處理數(shù)據(jù)樣本的代碼可能會(huì)變得凌亂且難以維護(hù)；理想情況下，我們希望數(shù)據(jù)集代碼與模型訓(xùn)練代碼解耦，以獲得更好的可讀性和模塊化。PyTorch提供的torch.utils.data.DataLoader和torch.utils.data.Dataset允許你使用預(yù)下載的數(shù)據(jù)集或自己制作的數(shù)據(jù)

一、Dataset

Dataset 類提供一種方式去獲取數(shù)據(jù)及其標(biāo)簽

主要有兩個(gè)目的：

獲取每一個(gè)數(shù)據(jù)及其標(biāo)簽
獲取數(shù)據(jù)的總量大小

1. 在控制臺(tái)進(jìn)行操作

Hymenoptera （膜翅目昆蟲(chóng)）數(shù)據(jù)集下載地址：

鏈接: https://pan.baidu.com/s/1XKwXsAtE2yzZW2IsvBDpnw?pwd=8a5t

提取碼: 8a5t

這是一個(gè)螞蟻蜜蜂二分類的數(shù)據(jù)集，通常數(shù)據(jù)集有以下三種組織形式（上面的數(shù)據(jù)集屬于第一種）：

不同的類別以文件夾的形式存在，文件夾中是該類別的圖片
圖片與標(biāo)簽分別存儲(chǔ)，圖片在一個(gè)文件夾下，label信息在另一個(gè)文件夾下
label直接寫在圖片名稱里

①獲取圖片的基本信息

在Pycharm 中，點(diǎn)擊下方的PythonConsole進(jìn)入控制臺(tái)進(jìn)行操作（通過(guò)控制臺(tái)可以看到變量的詳細(xì)信息）

首先加載圖片，逐行輸入下方代碼：

from PIL import Image
img_path = "./dataset/hymenoptera_data/train/ants/0013035.jpg"
img = Image.open(img_path)

此時(shí)我們就可以在右側(cè)看到相關(guān)變量的信息：

點(diǎn)擊img變量，可以查看圖片的詳細(xì)信息。通過(guò)控制臺(tái)執(zhí)行程序能夠直觀地獲取后續(xù)操作所需的數(shù)據(jù)：

最后可以通過(guò)img.show()打開(kāi)圖片查看：

②獲取文件的基本信息

同樣還是在控制臺(tái)逐行輸入以下代碼：

dir_path = "dataset/hymenoptera_data/train/ants"
import os
img_path_list = os.listdir(dir_path)
img_path_list[0]

我們就可以獲取到文件夾下的文件名稱，由于是使用控制臺(tái)，我們還可以在右側(cè)查看列表的詳細(xì)信息：

因此在控制臺(tái)操作是有很大的優(yōu)點(diǎn)的，我們可以在控制臺(tái)逐行執(zhí)行已經(jīng)編寫好的文件中的語(yǔ)句，通過(guò)查看右側(cè)變量的值來(lái)判斷程序?qū)懙氖欠裼袉?wèn)題

2. 編寫一個(gè)繼承Dataset 的類加載數(shù)據(jù)

下面的代碼也可以在控制臺(tái)運(yùn)行（可以多行復(fù)制粘貼）來(lái)檢驗(yàn)程序是否有誤

①定義 MyData類

導(dǎo)入所需頭文件：

from torch.utils.data import Dataset
from PIL import Image
import os

定義MyData類：

__init__：初始化函數(shù)
__getitem__：返回指定下標(biāo)的圖片和標(biāo)簽
__len__：返回?cái)?shù)據(jù)集的大小

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)

其中os.path.join()可以實(shí)現(xiàn)多個(gè)路徑的合并且不出錯(cuò)

②創(chuàng)建類的實(shí)例并調(diào)用

創(chuàng)建 MyData 類的實(shí)例：

if __name__ == "__main__":
    root_dir = "../dataset/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)

調(diào)用類中寫好的函數(shù)：

    img, label = ants_dataset.__getitem__(3)
    print(ants_dataset.__len__(), label)
    img.show()

同時(shí)我們也可以通過(guò)下面這種方式用已有的數(shù)據(jù)集來(lái)創(chuàng)造數(shù)據(jù)集：

train_dataset = ants_dataset + bees_dataset

二、DataLoader

DataLoader 類是為后面的網(wǎng)絡(luò)提供不同的數(shù)據(jù)形式
DataLoader 會(huì)根據(jù)batch_size的值對(duì)數(shù)據(jù)進(jìn)行打包
導(dǎo)入所需的包

import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

加載數(shù)據(jù)：

test_data = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.ToTensor())
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=True)

測(cè)試：

img, target = test_data[0]
print(img.shape)
print(target)

進(jìn)行日志記錄，開(kāi)始訓(xùn)練：

writer = SummaryWriter("dataloader")
for epoch in range(2):
    step = 0
    for data in test_loader:
        imgs, targets = data
        print(imgs.shape)
        print(targets)
        writer.add_images("Epoch: {}".format(epoch), imgs, step)
        step = step + 1
writer.close()

到此這篇關(guān)于PyTorch Dataset與DataLoader使用超詳細(xì)講解的文章就介紹到這了,更多相關(guān)PyTorch Dataset與DataLoader內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: