溫馨提示×
您好，登錄后才能下訂單哦！
點擊 登錄注冊 即表示同意《億速云用戶服務條款》
您好，登錄后才能下訂單哦！
這篇文章將為大家詳細講解有關如何在PyTorch中重寫或改寫Dataset并載入Dataloader，小編覺得挺實用的，因此分享給大家做個參考，希望大家閱讀完這篇文章后可以有所收獲。
前言
眾所周知，Dataset和Dataloder是pytorch中進行數據載入的部件。必須將數據載入后，再進行深度學習模型的訓練。在pytorch的一些案例教學中，常使用torchvision.datasets自帶的MNIST、CIFAR-10數據集，一般流程為：
# 下載并存放數據集 train_dataset = torchvision.datasets.CIFAR10(root="數據集存放位置",download=True) # load數據 train_loader = torch.utils.data.DataLoader(dataset=train_dataset)
但是，在我們自己的模型訓練中，需要使用非官方自制的數據集。這時應該怎么辦呢？
我們可以通過改寫torch.utils.data.Dataset中的__getitem__和__len__來載入我們自己的數據集。__getitem__獲取數據集中的數據，__len__獲取整個數據集的長度（即個數）。
改寫
采用pytorch官網案例中提供的一個臉部landmark數據集。數據集中含有存放landmark的csv文件，但是我們在這篇文章中不使用（其實也可以隨便下載一些圖片作數據集來實驗）。
import os import torch from skimage import io, transform import numpy as np import matplotlib.pyplot as plt from torch.utils.data import Dataset, DataLoader from torchvision import transforms, utils plt.ion() # interactive mode
torch.utils.data.Dataset是一個抽象類，我們自己的數據集需要繼承Dataset，然后改寫上述兩個函數：
class ImageLoader(Dataset):
def __init__(self, file_path, transform=None):
super(ImageLoader,self).__init__()
self.file_path = file_path
self.transform = transform # 對輸入圖像進行預處理，這里并沒有做，預設為None
self.image_names = os.listdir(self.file_path) # 文件名的列表
def __getitem__(self,idx):
image = self.image_names[idx]
image = io.imread(os.path.join(self.file_path,image))
# if self.transform:
# image= self.transform(image)
return image
def __len__(self):
return len(self.image_names)
# 設置自己存放的數據集位置，并plot展示
imageloader = ImageLoader(file_path="D:\\Projects\\datasets\\faces\\")
# imageloader.__len__() # 輸出數據集長度（個數），應為71
# print(imageloader.__getitem__(0)) # 以數據形式展示
plt.imshow(imageloader.__getitem__(0)) # 以圖像形式展示
plt.show()得到的圖片輸出：
得到的數據輸出，：
array([[[ 66, 59, 53],
[ 66, 59, 53],
[ 66, 59, 53],
...,
[ 59, 54, 48],
[ 59, 54, 48],
[ 59, 54, 48]],
...,
[153, 141, 129],
[158, 146, 134],
[158, 146, 134]]], dtype=uint8)上面看到dytpe=uint8，實際進行訓練的時候，常常需要更改成float的數據類型。可以使用：
# 直接改成pytorch中的tensor下的float格式 # 也可以用numpy的改成普通的float格式 to_float= torch.from_numpy(imageloader.__getitem__(0)).float()
改寫完成后，直接使用train_loader =torch.utils.data.DataLoader(dataset=imageloader)載入到Dataloader中，就可以使用了。
下面的代碼可以試著運行一下，產生的是一模一樣的圖片結果。
train_loader = torch.utils.data.DataLoader(dataset=imageloader) train_loader.dataset[0] plt.imshow(train_loader.dataset[0]) plt.show()
關于如何在PyTorch中重寫或改寫Dataset并載入Dataloader就分享到這里了，希望以上內容可以對大家有一定的幫助，可以學到更多知識。如果覺得文章不錯，可以把它分享出去讓更多的人看到。
免責聲明：本站發布的內容（圖片、視頻和文字）以原創、轉載和分享為主，文章觀點不代表本網站立場，如果涉及侵權請聯系站長郵箱：is@yisu.com進行舉報，并提供相關證據，一經查實，將立刻刪除涉嫌侵權內容。
