Dataloader数据集制作

新的存储方式

上节中，我们介绍了图像分类模型，所用的数据集是不同命名的文件夹存储各类花朵的图片。然而由于这样的存储方式应用很少，所以在实战中不算实用。本次我们介绍另外一种较为常见且比较简单的数据存储方式，将训练集和测试集图片放在两个文件价值中，然后建立两个.txt文件存储图片名称和其对应的分类：

这样的存储方式更常见，因为其在实践中更适合制作。今天我们就来介绍这样的数据集如何读入我们的模型。
那么在正式开始今天的学习之前，我们还是先把用得到的包全部引用进来：

import os
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
import torch.optim as optim
import torchvision
from torchvision import transforms, models, datasets
import imageio
import time
import warnings
warnings.filterwarnings("ignore")
import random
import sys
import copy
import json
from PIL import Image

读取txt文件中的路径和标签

先给大家介绍一种按字典模式读取到.txt文件的方法，比如我们有一个txt文档存储着这样的内容：

我们想把第一个字符串当做是字典的键，后面的内容都读取成自己的值，那么就可以这样操作：

data_infos = {}
with open('test.txt') as f:  # 打开路径下的文件
    # strip()可以避免读到换行符
    # split(' ')可以指定以空格为分割符对数据进行切分
    samples = [x.strip().split(' ') for x in f.readlines()]
for i,j,k in samples:
    data_infos[i] = np.array((j,k), dtype=np.int64)
print(data_infos)
# 运行结果为：{'ABC': array([ 1, 11], dtype=int64), 'def': array([ 2, 22], dtype=int64), 'cbg': array([ 3, 34], dtype=int64)}

把上述的代码包装成一个函数，就可以完成读取数据的第一步了：

def load_annotations(ann_file): # 传入读取数据的路径
    data_infos = {} # 定义空字典存储数据
    with open(ann_file) as f: # 打开路径下的文件
        # strip()可以避免读进换行符
        # split(' ')可以指定以空格为分割符对数据进行切分
        samples = [x.strip().split(' ') for x in f.readlines()] # f.readlines()指定计算机按行读取数据文本的内容
        for filename, gt_label in samples:
            data_infos[filename] = np.array(gt_label, dtype=np.int64)
    return data_infos

当我们传入一个正确的路径后，就可以拿到正确的字典了：

img_label = load_annotations('./flower_data/new_store/train.txt')
print(img_label)

以下展示部分显示内容：

当然，打包成dict的数据集并不能直接使用，我们需要把它存成list的格式。这是因为dataloader不能处理dict类型的数据。使用dataloader处理数据可以很大程度的加快处理数据的速度，因此我们还是屈服的好，list就list吧：

image_name = list(img_label.keys())
label = list(img_label.values())

相信这两句代码大家都已经很熟悉了，这里就不在介绍啦~

读取完整的数据路径

读取好文件名和分类名，接下来就要把图片的完整路径传递给计算机，以便计算机可以正确的读取到图像文件：

data_dir = './flower_data/new_store/'
train_dir = data_dir + '/train_filelist'
valid_dir = data_dir + '/val_filelist'
# 把图片具体路径填成图像名的前缀
image_path = [os.path.join(train_dir,img) for img in image_name]
# 查看
print(image_path)

不必理会方向不同的斜杠，这不影响计算机读取该路径下的内容：

import cv2
img= cv2.imread(image_path[0])
# 调整颜色通道
img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB) 
plt.imshow(img)
plt.show()

来展示一下：

看懂了吧~

数据预处理

预处理的方法这里就不给大家过多介绍了，数据预处理的方法和上篇中介绍的一样，这里给大家附上代码：

data_transforms = {
    'train': 
        transforms.Compose([
        transforms.Resize(64),
        transforms.RandomRotation(45),
        transforms.CenterCrop(64),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomVerticalFlip(p=0.5),
        transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),
        transforms.RandomGrayscale(p=0.025)
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'valid': 
        transforms.Compose([
        transforms.Resize(64),
        transforms.CenterCrop(64),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

需要注意的是，如果我们的数据量足够庞大的话，是可以不进行预处理的，我们在制作数据集的时候可以自行决定是否进行预处理，以及如何进行预处理。

整理任务，制作模型

下面我们要做一件事，把上述任务打包成一个模型的子类。话不多说直接上代码：

from torch.utils.data import Dataset, DataLoader
class FlowerDataset(Dataset):
    def __init__(self, root_dir, ann_file, transform=None):
        # 构造函数必须设置，其用于存储图片路径及对应标签
        # 三个参数分别代表了训练集或验证集的路径、
        # 训练集或验证集标记文本路径以及如何进行预处理
        self.ann_file = ann_file
        self.root_dir = root_dir
        self.img_label = self.load_annotations() # 做好图像标签对应的字典
        # 给图片加好路径
        self.img = [os.path.join(self.root_dir,img) for img in list(self.img_label.keys())]
        # 取到标签列表
        self.label = [label for label in list(self.img_label.values())]
        self.transform = transform
 
    def __len__(self):
        return len(self.img)
        
    # 该函数每次只能打包一个数据
    def __getitem__(self, idx): # idx代表一个索引，通过idx可以找到该值对应的图片和分类
    	# 取到图像对应的完整路径信息
        image = Image.open(self.img[idx])
        # 取到下标idx对应的标签信息
        label = self.label[idx]
        if self.transform: # 判断是否需要预处理，默认为不需要
            image = self.transform(image)
        # 把标签加载成torch格式
        label = torch.from_numpy(np.array(label)) # 因为不支持list直接向torch转换，所以需要借助array过渡
        return image, label
        
    # 读取图像名称和标签信息
    def load_annotations(self):
        data_infos = {}
        with open(self.ann_file) as f:
            samples = [x.strip().split(' ') for x in f.readlines()]
            for filename, gt_label in samples:
                data_infos[filename] = np.array(gt_label, dtype=np.int64)
        return data_infos

实例化我们的dataloader

至此，我们已经完成了准备工作，接下来需要做的就是加载好测试集和验证集了，这里的方法也和上次一样：

# 加载训练集
train_dataset = FlowerDataset(root_dir=train_dir, ann_file = './flower_data/new_store/train.txt', transform=data_transforms['train'])
                             # 图像文件路径         标注文件路径                                        预处理方式
# 加载验证集
val_dataset = FlowerDataset(root_dir=valid_dir, ann_file = './flower_data/val.txt', transform=data_transforms['valid'])

# 分别打包：
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=True)

检测实例化是否正确

在开始检验之前，我们先给大家详细介绍两个方法，分别为squeeze()和permute()。squeeze可以帮我们去除掉数组中没有实际意义的维度，观察这段代码：

a=[[1,2,3,4]]
a=np.array(a)
print(a)
print(a.squeeze()) 
print(a.squeeze()) # 如果array没有多余的维度则压缩之后没有变化
# 输出为：[1,2,3,4]
        # [1,2,3,4]

permute()方法可以调换数组的维度，在实战中我们如果获得了一个3×64×64的矩阵想要转化为64×64×3的矩阵，就需要用到这个函数，示例如下：

b=[[[1,2,3],[4,5,6]],[[7,8,9],[10,11,12]],[[13,14,15],[16,17,18]],[[19,20,21],[22,23,24]]]
c=torch.from_numpy(np.array(b))
print(c) # 规格：4*2*3
d=c.permute((1, 2, 0))
print(d) # 规格：2*3*4
e=c.permute((2, 1, 0)).numpy() # 转换成ndarray格式，因为计算机无法直接
                               # 对tensor格式的内容进行绘图
                               # 如果使用GPU去跑，还需要先转换成CPU再转换成ndarray
print(e) # 规格：3*2*4
print(type(d),type(e)) # 查看类型

我们来看看结果：

开胃菜吃完了，咱们上主食了：

image, label = iter(train_loader).next() # 取一个batch数据
print(label,image.shape)
sample = image[0].squeeze() # image取出第0标签的图片
                            # 但是取出格式可能为1*3*64*64
                            # squeeze可以压缩掉多余的维度（此处为第一个维度）
sample = sample.permute((1, 2, 0)).numpy() # 变换维度，将通道数滞后，整理成64*64*3的维度
# 一定要注意维度正确，切记切记
sample *= [0.229, 0.224, 0.225]
sample += [0.485, 0.456, 0.406]
plt.imshow(sample)
plt.show()
print('Label is: {}'.format(label[0].numpy()))

看看输出结果：

这样看来我们确实设置正确了，下面就该跑模型了~
如果小伙伴对矩阵的乘法和加法有疑问，可以运行以下代码查看计算机是如何执行的：

b=[[[1,2],[4,5]],[[7,8],[10,11]],[[13,14],[16,17]]]
b=np.array(b)
print(b)
print()
c=torch.from_numpy(np.array(b)).permute((1, 2, 0)).numpy()
d=c*[2,1,3]
print(d)
print()
e=c+[1,2,3]
print(e)

代码整合

今天我们要讲的内容是如何定义自己的数据集，因此到这里本节的内容就已经全部结束了，但是在最后，我还是把全部的代码整合给大家，当然这次就不再做迁移学习了：

import os
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
import torch.optim as optim
import torchvision
from torchvision import transforms, models, datasets
import imageio
import time
import warnings
warnings.filterwarnings("ignore")
import random
import sys
import copy
import json
from PIL import Image
from torch.utils.data import Dataset, DataLoader
data_dir = './flower_data/new_store/'
train_dir = data_dir + '/train_filelist'
valid_dir = data_dir + '/val_filelist'
data_transforms = {
    'train':
        transforms.Compose([
            transforms.Resize(64),
            transforms.RandomRotation(45),
            transforms.CenterCrop(64),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomVerticalFlip(p=0.5),
            transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),
            transforms.RandomGrayscale(p=0.025),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    'valid':
        transforms.Compose([
            transforms.Resize(64),
            transforms.CenterCrop(64),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
}


class FlowerDataset(Dataset):
    def __init__(self, root_dir, ann_file, transform=None):
        self.ann_file = ann_file
        self.root_dir = root_dir
        self.img_label = self.load_annotations()
        self.img = [os.path.join(self.root_dir, img) for img in list(self.img_label.keys())]
        # 取到标签列表
        self.label = [label for label in list(self.img_label.values())]
        self.transform = transform

    def __len__(self):
        return len(self.img)

    def __getitem__(self, idx):
        image = Image.open(self.img[idx])
        label = self.label[idx]
        if self.transform:
            image = self.transform(image)
        label = torch.from_numpy(np.array(label))
        return image, label

    def load_annotations(self):
        data_infos = {}
        with open(self.ann_file) as f:
            samples = [x.strip().split(' ') for x in f.readlines()]
            for filename, gt_label in samples:
                data_infos[filename] = np.array(gt_label, dtype=np.int64)
        return data_infos


train_dataset = FlowerDataset(root_dir=train_dir, ann_file='./flower_data/new_store/train.txt',
                              transform=data_transforms['train'])
val_dataset = FlowerDataset(root_dir=valid_dir, ann_file='./flower_data/new_store/val.txt', transform=data_transforms['valid'])

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=True)

dataloaders = {'train': train_loader, 'valid': val_loader}
# 是否用GPU训练
train_on_gpu = torch.cuda.is_available()

if not train_on_gpu:
    print('CUDA is not available.  Training on CPU ...')
else:
    print('CUDA is available!  Training on GPU ...')

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_ft = models.resnet18()
# 是否用人家训练好的特征来做
feature_extract = True
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Sequential(nn.Linear(num_ftrs, 102))
input_size = 64
# 优化器设置
optimizer_ft = optim.Adam(model_ft.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)  # 学习率每7个epoch衰减成原来的1/10
criterion = nn.CrossEntropyLoss()


def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, filename='best.pt'):
    # 记录时间
    since = time.time()
    # 记录最好一次的准确率，防止迭代次数过多之后出现准确率下降的情况
    best_acc = 0
    # 把模型放到CPU或者GPU
    model.to(device)
    # 训练过程中打印一堆损失和指标
    # 保存训练集、验证集的准确率及损失
    val_acc_history = []
    train_acc_history = []
    train_losses = []
    valid_losses = []
    # 取到字典结构中保存的学习率
    LRs = [optimizer.param_groups[0]['lr']]
    # 验证集的效果比之前的好才会更新
    # 拷贝模型当前的权重参数，后面会不断地用更好的参数覆盖
    best_model_wts = copy.deepcopy(model.state_dict())
    # 一个个epoch来遍历
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # 训练和验证
        for phase in ['train', 'valid']:
            if phase == 'train':
                model.train()  # 训练
            else:
                model.eval()  # 验证

            running_loss = 0.0
            running_corrects = 0

            # 把数据都取个遍
            for inputs, labels in dataloaders[phase]:
                # 把数据和标签放到CPU或GPU
                inputs = inputs.to(device)  # 这里不会修改inputs和labels的实质内容
                labels = labels.to(device)

                # 清零
                optimizer.zero_grad()
                outputs = model(inputs)  # 计算损失
                loss = criterion(outputs, labels)
                _, preds = torch.max(outputs, 1)  # 拿到预测值（概率最大的分类标签）
                # 训练阶段更新权重
                if phase == 'train':
                    # 反向传播
                    loss.backward()
                    # 只有训练的时候计算和更新梯度
                    # 参数更新
                    optimizer.step()

                # 计算损失
                # 一轮完整的训练会进行52次参数的修改
                # 这里计算的损失是每一次修改前的损失
                # 由于每一次更新参数后损失和预测准确率都会有所变化
                # 所以后续还要求均值进行比较
                running_loss += loss.item() * inputs.size(0)  # 0表示batch那个维度
                running_corrects += torch.sum(preds == labels.data)  # 预测结果最大的和真实值是否一致

            # 求平均
            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

            time_elapsed = time.time() - since  # 一个epoch我浪费了多少时间
            print('Time elapsed {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

            # 查看当前模型训练结果是否为最优，是则保存
            # 这部分会不断被覆盖，最终只会留下一个结果
            if phase == 'valid' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())  # 用当前的参数作为最好的参数进行保存
                state = {
                    'state_dict': best_model_wts,  # 字典里key就是各层的名字，值就是训练好的权重
                    'best_acc': best_acc,
                    'optimizer': optimizer.state_dict(),
                }
                torch.save(state, filename)
            # 将所有信息都进行保存，不会进行覆盖操作
            if phase == 'valid':
                val_acc_history.append(epoch_acc)
                valid_losses.append(epoch_loss)
            if phase == 'train':
                train_acc_history.append(epoch_acc)
                train_losses.append(epoch_loss)

        print('Optimizer learning rate : {:.7f}'.format(optimizer.param_groups[0]['lr']))
        LRs.append(optimizer.param_groups[0]['lr'])  # 保存当前学习率
        print()
        scheduler.step()  # 学习率衰减，累加到设定值的时候会自动衰减

    time_elapsed = time.time() - since  # 计算整体运行所花时间
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # 训练完后用最好的一次当做模型最终的结果,等着一会测试
    model.load_state_dict(best_model_wts)
    return model, val_acc_history, train_acc_history, valid_losses, train_losses, LRs


model_ft, val_acc_history, train_acc_history, valid_losses, train_losses, LRs = train_model(model_ft, dataloaders,
                                                                                            criterion, optimizer_ft,
                                                                                            num_epochs=20,
                                                                                            filename='best.pth')

# 验证一下预测结果：
with open('./flower_data/new_store/cat_to_name.json', 'r') as f: # json文件存放的是字典内容，编号为键，名字为值
    cat_to_name = json.load(f)
dataiter = iter(dataloaders['valid']) # 制作成迭代器
images, labels = dataiter.next() # 每次取下一组的64个验证集图片数据
model_ft.eval()

if train_on_gpu:
    output = model_ft(images.cuda())
else:
    output = model_ft(images)
def im_convert(tensor):
    """ 展示数据"""
    image = tensor.to("cpu").clone().detach()
    image = image.numpy().squeeze() 
    image = image.transpose(1,2,0) # 把颜色通道数放在最后
    image = image * np.array((0.229, 0.224, 0.225)) + np.array((0.485, 0.456, 0.406)) # 去标准化，尽量更接近原图
    image = image.clip(0, 1)
    return image

fig=plt.figure(figsize=(20, 20))
columns =4
rows = 2

_, preds_tensor = torch.max(output, 1)
# GPU和CPU转换格式的方法有所不同
preds = np.squeeze(preds_tensor.numpy()) if not train_on_gpu else np.squeeze(preds_tensor.cpu().numpy())

for idx in range (columns*rows):
    ax = fig.add_subplot(rows, columns, idx+1, xticks=[], yticks=[])
    plt.imshow(im_convert(images[idx]))
    ax.set_title("{} ({})".format(cat_to_name[str(preds[idx])], cat_to_name[str(labels[idx].item())]),
                 # 如果是绿色字，证明预测正确，红色字则代表预测错误
                 color=("green" if cat_to_name[str(preds[idx])]==cat_to_name[str(labels[idx].item())] else "red"))
plt.show()

这样一来我们就可以得到和前篇类似的书出结果了：

这也算是对上节内容的最终整理了，感谢看到最后的小伙伴，我们下期再见~