1 關於Pytorch內置的Dataset
我們在《torch.utils.data.DataLoader與迭代器轉換》中介紹了如何使用Pytorch內置的數據集進行論文實現，如torchvision.datasets。下面是加載內置訓練數據集的常見操作：
from torchvision.datasets import FashionMNIST
from torchvision.transforms import Compose, ToTensor, Normalize
RAW_DATA_PATH = './rawdata'
transform = Compose(
[ToTensor(),
Normalize((0.1307,), (0.3081,))
]
)
train_data = FashionMNIST(
root=RAW_DATA_PATH,
download=True,
train=True,
transform=transform
)
這里的train_data做為dataset對象，它擁有許多熟悉，我們可以通過以下方法獲取樣本數據的分類類別集合、樣本的特征維度、樣本的標簽集合等信息。
classes = train_data.classes
num_features = train_data.data[0].shape[0]
train_labels = train_data.targets
print(classes)
print(num_features)
print(train_labels)
輸出如下:
['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
28
tensor([9, 0, 0, ..., 3, 0, 5])
但是，我們常常會在訓練集的基礎上拆分出驗證集（或者只用部分數據來進行訓練），而我們在后續的代碼中常常會將拆分后的數據集也默認為Dataset對象，那么我們如何保證拆分后的對象還能夠具備Dataset類應有的屬性和功能，從而做到代碼的一致性呢？
2 自定義Subset類
關於數據集拆分，我們想到的第一個方法是使用torch.utils.data.random_split對dataset進行划分，下面我們假設划分10000個樣本做為訓練集，其余樣本做為驗證集：
from torch.utils.data import random_split
k = 10000
train_data, valid_data = random_split(train_data, [k, len(train_data)-k])
注意我們如果打印train_data和valid_data的類型，可以看到顯示：
<class 'torch.utils.data.dataset.Subset'>
已經不再是torchvision.datasets.mnist.FashionMNIST對象，而是一個所謂的Subset對象！此時Subset對象雖然仍然還存有data屬性，但是內置的target和classes屬性已經不復存在，比如如果我們強行訪問valid_data的target屬性：
valid_target = valid_data.target
就會報如下錯誤：
'Subset' object has no attribute 'target'
為了解決這個問題，這里有一個trick，那就是以繼承SubSet類的方式的方式定義一個新的CustomSubSet類，使新類在保持SubSet類的基本屬性的基礎上，擁有和原本數據集類相似的屬性，如targets和classes等:
from torch.utils.data import Subset
class CustomSubset(Subset):
'''A custom subset class'''
def __init__(self, dataset, indices):
super().__init__(dataset, indices)
self.targets = dataset.targets # 保留targets屬性
self.classes = dataset.classes # 保留classes屬性
def __getitem__(self, idx): #同時支持索引訪問操作
x, y = self.dataset[self.indices[idx]]
return x, y
def __len__(self): # 同時支持取長度操作
return len(self.indices)
然后就引出了第二種划分方法，即通過初始化CustomSubset對象的方式直接對數據集進行划分（這里為了簡化省略了shuffle的步驟）：
import numpy as np
from copy import deepcopy
origin_data = deepcopy(train_data)
train_data = CustomSubset(origin_data, np.arange(k))
valid_data = CustomSubset(origin_data, np.arange(k, len(origin_data))-k)
注意，CustomSubset類的初始化方法的第二個參數indices為樣本索引，我們可以通過np.arange()的方法來創建。
然后，我們再訪問valid_data對應的classes和targes屬性：
print(valid_data.classes)
print(valid_data.targets)
此時，我們發現可以成功訪問這些屬性了：
['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
tensor([9, 0, 0, ..., 3, 0, 5])
當然，CustomSubset的作用並不只是添加數據集的屬性，我們還可以自定義一些數據預處理操作。我們將類的結構修改如下：
class CustomSubset(Subset):
'''A custom subset class with customizable data transformation'''
def __init__(self, dataset, indices, subset_transform=None):
super().__init__(dataset, indices)
self.targets = dataset.targets
self.classes = dataset.classes
self.subset_transform = subset_transform
def __getitem__(self, idx):
x, y = self.dataset[self.indices[idx]]
if self.subset_transform:
x = self.subset_transform(x)
return x, y
def __len__(self):
return len(self.indices)
我們可以在使用樣本前設置好數據預處理算子:
from torchvision import transforms
valid_data.subset_transform = transforms.Compose(\
[transforms.RandomRotation((180,180))])
這樣，我們再像下列這樣用索引訪問取出數據集樣本時，就會自動調用算子完成預處理操作:
print(valid_data[0])
打印結果縮略如下：
(tensor([[[-0.4242, -0.4242, -0.4242, ......-0.4242, -0.4242, -0.4242, -0.4242, -0.4242]]]), 9)
3 自定義Dataset類
推廣到更一般的情況，我們可以直接自定義數據集。操作幾乎跟上面所說的完全一致，比如在這里我們可以直接繼承torch.utils.data.Dataset類然后自定義一個CustomDataset類就行。如下所示：
from torch.utils.data import Dataset
class CustomDataset(Dataset):
"""An abstract Dataset class wrapped around Pytorch Dataset class.
"""
def __init__(self, dataset, indices):
self.dataset = dataset
self.indices = [int(i) for i in indices]
self.targets = dataset.targets # 保留targets屬性
self.classes = dataset.classes # 保留classes屬性
def __len__(self):
return len(self.indices)
def __getitem__(self, item):
x, y = self.dataset[self.indices[item]]
return x, y
事實上，自定義Dataset類的使用遠不限於訓練集和測試集的拆分（我們這里僅僅是用訓練集和測試集的拆分來舉個例子），還可以用於許多CV、NLP的數據集自定義。所有自定義的數據集在重寫了__getitem__方法和__len__方法后，那么我們就可以使用torch.utils.data.DataLoader（參見《torch.utils.data.DataLoader與迭代器轉換》）將自定義的數據集像普通數據集一樣進行迭代訓練。對我們自定義Dataset類的迭代測試如下：
import numpy as np
from torch.utils.data import DataLoader
from torchvision.datasets import FashionMNIST
from torchvision.transforms import Compose, ToTensor, Normalize
if __name__ == "__main__":
RAW_DATA_PATH = './rawdata'
transform = Compose(
[ToTensor(),
Normalize((0.1307,), (0.3081,))
]
)
train_data = FashionMNIST(
root=RAW_DATA_PATH,
download=True,
train=True,
transform=transform
)
train_dataset = CustomDataset(train_data, np.arange(100, 200))
# train_dataloader
train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
for batch_idx, (images, labels) in enumerate(train_dataloader):
print("Iteration %d: " % batch_idx, images.shape, labels.shape)
可以看到，我們設置索引為\([100,200)\)內的樣本為訓練樣本，然后采用batch size大小為\(64\)進行迭代，可以看到迭代結果如下：
Iteration 0: torch.Size([64, 1, 28, 28]) torch.Size([64])
Iteration 1: torch.Size([36, 1, 28, 28]) torch.Size([36])
可見我們自定義的數據集類工作正確。