PyTorch学习笔记1:Dataset和Dataloader
1. 加载数据集
1.1 Dataset
提供一种方式获取数据及其标签(label)
- 如何获取每一个数据及其label
- 告诉我们总共有多少的数据
代码示例:
这里使用的数据集为蚂蚁蜜蜂数据集:
数据集下载地址
密码:5suq
定义Dataset类
from torch.utils.data import Dataset # PyTorch的数据集模块
from PIL import Image # 图像加载模块
import os # 路径操作模块
# 继承Dataset类,重写__getitem__(self, index), __len__(self)方法
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, index):
img_name = self.img_path[index] # 获取图像文件名
img_item_path = os.path.join(self.path, img_name) # 获取图像完整路径
img = Image.open(img_item_path) # 加载图像
label = self.label_dir # 标签是子目录名称
return img, label # 返回图像和标签
# 获数据集大小函数
def __len__(self):
return len(self.img_path) # 图像数量
实例化类并使用:
# 实例化并使用
ants_dataset = MyData('data/hymenoptera_data/train', 'ants')
img, label = ants_dataset[1]
print('数据集大小为:', len(ants_dataset))
print("标签为:", label)
img.show()
1.2 Dataloader
为后面的网络提供不同的数据形式(迭代器)
采样并以迭代的形式提供数据
分batch、打乱之类的操作:
Dataloader常用参数介绍
参数 | 数据类型 | 解释 |
---|---|---|
dataset | Dataset | 加载数据的数据集 |
batchsize | int | 每批次加载多少个样本(默认值:1) |
shuffle | bool | 是否打乱 |
num_workers | int | 是否多进程读取,默认是0表示主进程,-1表示所有 |
drop_last | bool | 当样本数不能被batchsize整除时, 是否舍弃最后一批数据1 |
为了理解drop_last,我们需要搞清楚epoch,iteration和batch_size:
epoch:所有训练样本都经过了模型一次训练,称为一个epoch
Iteration:一批样本输入到模型中,称为一个Iteration
batch_size: 一批样本(bath)的大小, 决定一个Epoch有多少个Iteration ↩︎