pytorch之分类问题图片数据加载

pytorch中加载数据的顺序是：
①创建一个dataset对象
②创建一个dataloader对象
③循环dataloader对象，将data,label拿到模型中去训练

1.创建datasets格式的数据

这里使用的是ImageFolder，ImageFolder是一个通用的数据加载器，数据集中的数据以以下方式组织

root/dog/xxx.png
root/dog/xxy.png
root/dog/xxz.png

root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png

ImageFolder假设所有的文件按文件夹保存，每个文件夹下存储同一个类别的图片，文件夹名为类名，其构造函数如下：

ImageFolder(root, transform=None, target_transform=None, loader=default_loader)

- root：在root指定的路径下寻找图片
- transform：对PIL Image进行的转换操作，transform的输入是使用loader读取图片的返回对象
- target_transform：对label的转换
- loader：给定路径后如何读取图片，默认读取为RGB格式的PIL Image对象

label是按照文件夹名顺序排序后存成字典，即{类名:类序号(从0开始)}，一般来说最好直接将文件夹命名为从0开始的数字，这样会和ImageFolder实际的label一致

ImageFolder有如下成员变量:
- self.classes - 用一个list保存 类名
- self.class_to_idx - 类名对应的 索引
- self.imgs - 保存(img-path, class) tuple的list

返回值为一个DataSet对象，一维是img，二维是lable

from torchvision import transforms as T
import matplotlib.pyplot as plt
from torchvision.datasets import ImageFolder

normalize = T.Normalize(mean=[0.4, 0.4, 0.4], std=[0.2, 0.2, 0.2])
transform  = T.Compose([
	T.Resized(224,224),
	T.ToTensor(),
	normalize
])


dataset = ImageFolder('data/dogcat_2/',transform=transform)

# cat文件夹的图片对应label 0，dog对应1
print(dataset.class_to_idx)

# 所有图片的路径和对应的label
print(dataset.imgs)

2.产生batch数据

Dataset只负责数据的抽象，一次调用getitem只返回一个样本。前面提到过，在训练神经网络时，最好是对一个batch的数据进行操作，同时还需要对数据进行shuffle和并行加速等。对此，PyTorch提供了DataLoader帮助我们实现这些功能

DataLoader的函数定义如下

ataLoader(dataset, batch_size=1, shuffle=False, sampler=None,num_workers=0, collate_fn=default_collate, pin_memory=False,drop_last=False)

• dataset：加载的数据集(Dataset对象)
• batch_size：batch size
• shuffle:：是否将数据打乱
• sampler： 样本抽样（集成学习时可能会用到）
• num_workers：使用多进程加载的进程数，0代表不使用多进程
• collate_fn： 如何将多个样本数据拼接成一个batch，一般使用默认的拼接方式即可
• pin_memory：是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一些
• drop_last：dataset中的数据个数可能不是batch_size的整数倍，drop_last为True会将多出来不足一个batch的数据丢弃

3.训练

def train():
	# 创建模型并搬到显卡上
	model = ResNet18(5).to(device)

	# 创建优化器
	optimizer = optim.Adam(model.parameters(), lr=lr)

	# 创建损失函数
	criterion = nn.CrossEntropyLoss()

	 # 开始训练
	 for epoch in range(epochs):
	 	for step, (x, y) in enumerate(train_loader):

			x, y = x.to(device), y.to(device) # 搬到显卡上

			# 前向传播，并计算损失
			logits = model(x)
			loss = criterion(logits, y)
			
			# 梯度清零，反向传播，更新参数
			optimizer.zero_grad()
			loss.backward()
			optimizer.step()

			
	#模型保存
	torch.save(model.state_dict(), 'model.mdl')