pytorch中dataloader自定義數據集

前言

在深度學習中我們需要使用自己的數據集做訓練，因此需要將自定義的數據和標簽加載到pytorch里面的dataloader里，也就是自實現一個dataloader。

數據集處理

以花卉識別項目為例，我們分別做出圖片的訓練集和測試集，訓練集的標簽和測試集的標簽

flower_data/
├── train_filelist/
│   ├── image_0001.jpg
│   └── ...
├── val_filelist/
│   ├── image_1001.jpg
│   └── ...
├── train.txt  # 格式：文件名 標簽
└── val.txt

?數據目錄的組織方式如上所示。

首先看圖片的處理。圖片只要做好編號放在同一個文件夾里就好了。

再看標簽的處理。標簽處理我們自己規定了一種形式，就是圖像文件的名稱+空格+分類標簽。

可以看到前面第一列數據是圖像名稱，第二列數據是圖像的分組，同樣的數字為一組。比如分組為0的圖像就是同一種花朵。

自定義dataset

源碼

import os.path
import numpy as np
import torch
from PIL import Image  # 從PIL庫導入Image類
from torch.utils.data import Datasetclass FlowerDataSet(Dataset):"""花朵分類任務數據集類，繼承自torch的Dataset類"""def __init__(self, root_dir, ann_file, transform=None):"""初始化數據集實例Args:root_dir (str): 數據集根目錄路徑ann_file (str): 標注文件路徑transform (callable, optional): 數據預處理變換函數"""self.ann_file = ann_fileself.root_dir = root_dir# 加載圖片路徑與標簽的映射字典 {文件名: 標簽}self.image_label = self.load_annotations()# 構建完整圖片路徑列表 [root_dir/文件名1, ...]self.image = [os.path.join(self.root_dir, img) for img in list(self.image_label.keys())]# 構建標簽列表 [標簽1, 標簽2, ...]self.label = [lbl for lbl in list(self.image_label.values())]  # 重命名為lbl避免與導入的label沖突self.transform = transformdef __len__(self):"""返回數據集樣本數量"""return len(self.image)def __getitem__(self, index):"""獲取單個樣本數據Args:index (int): 樣本索引Returns:tuple: (預處理后的圖像數據, 對應的標簽)"""# 打開圖片文件image = Image.open(self.image[index])# 獲取對應標簽label = self.label[index]# 應用數據預處理if self.transform:image = self.transform(image)# 將標簽轉換為torch張量label = torch.from_numpy(np.array(label))return image, labeldef load_annotations(self):"""加載標注文件，解析圖片文件名和標簽的映射關系Returns:dict: {圖片文件名: 對應標簽} 的字典"""data_infos = {}with open(self.ann_file) as f:# 讀取所有行并分割，每行格式應為 "文件名 標簽"samples = [x.strip().split(' ') for x in f.readlines()]for filename, label in samples:# 將標簽轉換為int64類型的numpy數組data_infos[filename] = np.array(label, dtype=np.int64)return data_infos

解析

1、將標簽數據進行讀取，組成一個哈希表，哈希表的鍵是圖像的文件名稱，哈希表的值是分組標簽。

    def load_annotations(self):"""加載標注文件，解析圖片文件名和標簽的映射關系Returns:dict: {圖片文件名: 對應標簽} 的字典"""data_infos = {}with open(self.ann_file) as f:# 讀取所有行并分割，每行格式應為 "文件名 標簽"samples = [x.strip().split(' ') for x in f.readlines()]for filename, label in samples:# 將標簽轉換為int64類型的numpy數組data_infos[filename] = np.array(label, dtype=np.int64)return data_infos

上面的代碼里，在錄入標簽的時候使用數組進行記錄，這是為了兼容多標簽的場景。如果不考慮兼容問題，僅考慮在單標簽場景下的簡單實現，可以用下面的代碼：

def load_annotations(self):data_infos = {}with open(self.ann_file) as f:for line in f:filename, label = line.strip().split()  # 直接解包data_infos[filename] = int(label)        # 存為 Python 整數return data_infos# 在 __getitem__ 中直接轉為張量
label = torch.tensor(self.labels[index], dtype=torch.long)

2、遍歷哈希表，將文件名和標簽分別存在兩個數組里。這里注意，為了方便后面dataloader按照batch去讀取圖片，這里要將圖片的全路徑加到文件名里。

        # 構建完整圖片路徑列表 [root_dir/文件名1, ...]self.image = [os.path.join(self.root_dir, img) for img in list(self.image_label.keys())]# 構建標簽列表 [標簽1, 標簽2, ...]self.label = [lbl for lbl in list(self.image_label.values())]  # 重命名為lbl避免與導入的label沖突

3、在dataloader向顯卡/cpu加載數據的時候會調用getitem方法。比如一個batch里有64個數據，dataloader就會調用64次該方法，將64組圖片和標簽全部獲取后交給運算單元去處理。

    def __getitem__(self, index):"""獲取單個樣本數據Args:index (int): 樣本索引Returns:tuple: (預處理后的圖像數據, 對應的標簽)"""# 打開圖片文件image = Image.open(self.image[index])# 獲取對應標簽label = self.label[index]# 應用數據預處理if self.transform:image = self.transform(image)# 將標簽轉換為torch張量label = torch.from_numpy(np.array(label))return image, label

測試dataloader

import os
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import DataLoader
from torchvision import transforms
from dataloader import FlowerDataSet  # 假設你的數據集類在dataloader.py中def denormalize(image_tensor):"""將歸一化的圖像張量轉換為可顯示的格式"""mean = np.array([0.485, 0.456, 0.406])std = np.array([0.229, 0.224, 0.225])image = image_tensor.numpy().transpose((1, 2, 0))  # 轉換維度順序image = std * image + mean  # 反歸一化image = np.clip(image, 0, 1)  # 限制像素值范圍return imagedef test_dataloader():# 定義數據預處理data_transforms = {'train': transforms.Compose([transforms.Resize(64),transforms.RandomRotation(45),transforms.CenterCrop(64),transforms.RandomHorizontalFlip(p=0.5),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),'valid': transforms.Compose([transforms.Resize(64),transforms.CenterCrop(64),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}# 檢查文件路徑是否存在print("[1/5] 檢查文件路徑...")required_files = {'train_txt': './flower_data/train.txt','val_txt': './flower_data/val.txt','train_dir': './flower_data/train_filelist','val_dir': './flower_data/val_filelist'}for name, path in required_files.items():if not os.path.exists(path):print(f"? 文件/目錄不存在: {path}")returnprint(f"? {name}: {path} 存在")# 初始化數據集print("\n[2/5] 加載數據集...")try:train_dataset = FlowerDataSet(root_dir=required_files['train_dir'],ann_file=required_files['train_txt'],transform=data_transforms['train'])val_dataset = FlowerDataSet(root_dir=required_files['val_dir'],ann_file=required_files['val_txt'],transform=data_transforms['valid'])print("? 數據集加載成功")except Exception as e:print(f"? 數據集加載失敗: {str(e)}")return# 打印數據集信息print("\n[3/5] 數據集統計:")print(f"訓練集樣本數: {len(train_dataset)}")print(f"驗證集樣本數: {len(val_dataset)}")# 檢查單個樣本print("\n[4/5] 檢查單個樣本:")sample_idx = 0try:img, label = train_dataset[sample_idx]print(f"圖像張量形狀: {img.shape} (應接近 torch.Size([3, 64, 64]))")print(f"標簽類型: {type(label)} (應為 torch.Tensor)")print(f"標簽值: {label.item()} (應為整數)")except Exception as e:print(f"? 樣本檢查失敗: {str(e)}")# 可視化樣本print("\n[5/5] 可視化訓練集樣本...")try:plt.figure(figsize=(8, 8))img_show = denormalize(img)plt.imshow(img_show)plt.title(f"Label: {label.item()}")plt.axis('off')plt.show()except Exception as e:print(f"? 可視化失敗: {str(e)}")# 檢查DataLoaderprint("\n[附加] 檢查DataLoader:")train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)for loader, name in [(train_loader, '訓練集'), (val_loader, '驗證集')]:print(f"\n{name} DataLoader測試:")try:batch = next(iter(loader))images, labels = batchprint(f"批次圖像形狀: {images.shape} (應接近 [batch, 3, 64, 64])")print(f"批次標簽示例: {labels[:5].numpy()}")print(f"像素值范圍: [{images.min():.3f}, {images.max():.3f}]")except Exception as e:print(f"? {name} DataLoader錯誤: {str(e)}")if __name__ == '__main__':test_dataloader()

在測試代碼中，分別測試了文件路徑，dataset是否正常創建，dataset樣本數量，dataset樣本格式，dataset數據可視化，dataloader數據樣式。

在打印日志的時候需要注意，dataset和dataloader里面的變量都是張量形式的，所以需要轉換成python標量再打印。比如從dataset里取出的標簽label是一個一維張量，需要通過label.item()進行轉換。

?在遍歷的時候為了簡化代碼，將兩個dataloader放在同一個循環語句中處理，并且通過增加name變量來區分兩個dataloader。

for loader, name in [(train_loader, '訓練集'), (val_loader, '驗證集')]: