在前两篇文章中我们了解了pytorch的一些基础概念以及模型训练的整体框架。今天我们继续介绍pytorch的一些常用包。
传送门:
Dataset&DataLoader
对于数据处理上,我们希望样本数据集代码和模型训练代码分开,这样便于获得更强的可读性和模块化。因此,pytorch提供了两个数据元素:
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
Dataset存储样本示例以及其对应的标签;DataLoader包装了数据集的可迭代对象,方便访问样本,可以认为是一个数据加载器。
Dataset是一种抽象类,所有继承它的子类都应该包含__len__方法和__getitem__方法,前者表示数据集大小,后者支持实例化后的样本整数索引。一般情况下,继承此类的对象也都包含初始化方法__init__.
下面举一个例子:
import os
import pandas as pd
from torchvision.io import read_image
class CustomImageDataset(Dataset):
def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
self.img_labels = pd.read_csv(annotations_file)
self.img_dir = img_dir
self.transform = transform
self.target_transform = target_transform
def __len__(self):
return len(self.img_labels)
def __getitem__(self, idx):
img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
image = read_image(img_path)
label = self.img_labels.iloc[idx, 1]
if self.transform:
image = self.transform(image)
if self.target_transform:
label = self.target_transform(label)
return image, label
img_dir是存储数据(这个例子应该是图片)路径,annotations_file是存储图片对应标签的CSV文件。然后,方法__len__返回的是数据集所有图片的个数;__getitem__方法是在给定样本的索引后,加载样本数据以及其对应的标签。两个tranform是做一些必要的转换,这里先不用管。
在训练模型的时候,一般只取批量样本进行训练,因此,DataLoader实现这种“小批量”训练过程:
from torch.utils.data import DataLoader
DataLoader(