torch教程——数据加载和处理

前面训练神经网络，我们使用了torch内置的dataset类作为数据类。这种方法用来做一些demo不错，但是实际应用中需要自定义数据类。

下载数据集

https://download.csdn.net/download/weixin_42708161/14156096
下载后解压此压缩包，并将其存储到"data/faces"的目录中。
这个数据集包含了各种面部姿态，已经手工标注的关键点。
数据集存放在一个csv中，每一行是一个数据。

img_name	point0_x	point0_y	point1_x	point1_y	point2_x	point2_y	…

每个数据集包含68个点。

建立数据集类

数据集类继承自torch.utils.data.Dataset,需要实现三个函数，分别是构造函数(__init__)，__getitem__和__len__。

1. 构造函数(__init__)。初始化数据集类，如存储数据集类的路径等。
2. __getitem__。顾名思义，此函数根据索引返回数据。
3. __len__。此函数返回数据集的尺寸。

#数据加载和处理

import torch
import os
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, Dataset
from torchvision import  transforms, utils
import torch.nn as nn
import torch.optim as optim

#忽略警告
import warnings
warnings.filterwarnings("ignore")

plt.ion()

#数据集类
class FaceLandMarksDataset(Dataset):

    #                  #csv文件的路径 包含所有图片文件的目录 变换
    def __init__(self, csv_file, root_dir, transform = None):
        self.landmark_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

    #返回数据集长度
    def __len__(self):
        return len(self.landmark_frame)

    #根据索引返回某个具体数据：字典（图像和对应的标记点）
    def __getitem__(self, idx):
        #读取图像
        img_name = os.path.join(self.root_dir, self.landmark_frame.iloc[idx, 0])
        image = io.imread(img_name)
        #读取标记点
        landmarks = self.landmark_frame.iloc[idx, 1:]
        landmarks = np.array(landmarks)
        landmarks = np.array(landmarks, dtype=np.single).reshape(-1, 2)
        sample = {'image':image, 'landmark':landmarks}

        if self.transform:
            sample = self.transform(sample)

        return sample

数据变换

数据变换是必要的，尤其是如下情景

1. 把数据类型转换为tensor。
2. 把输入数据转化为同样的大小，这是因为神经网络通常假设输入数据的大小是一样的。
3. 数据增强：深度学习之一的瓶颈就是训练样本不足，通过数据变换可以认为的扩充数据集，常见的方法有裁剪，翻转，旋转等。
   我们把数据变换封装成类，并重写``call()方法，该方法相当于重载运算符()，因此可以像调用函数一样使用此类。之所以封装成类而不是封装成函数，是为了避免每次调用都传递参数。
   必要时，还可以实现构造函数。(__init__())

#定义变换：转换成固定大小
class Resacle(object):

    def __init__(self, out_size:tuple):
        self.out_size = out_size

    def __call__(self, sample):
        image, landmarks = sample["image"], sample["landmark"]
        h , w = image.shape[:2]
        new_h, new_w = self.out_size
        new_h ,new_w = int(new_h), int(new_w)
        image = transform.resize(image, (new_h, new_w))

        landmarks = landmarks * [new_w / w, new_h/h]

        sample = {"image": image, "landmark": landmarks}

        return sample

#转化成Tensor
class ToTensor(object):

    def __call__(self, sample):
        image, landmarks = sample["image"], sample["landmark"]
        #交换颜色轴
        image = image.transpose((2, 0, 1))

        image = torch.from_numpy(image).float()
        landmarks = torch.from_numpy(landmarks).float()

        sample = {"image":image, "landmark":landmarks}

        return sample

注意

这里在把数据转换成tensor时，一定是float类型。

组合变换

将多种类型的变换组合到一起。直接调用函数即可。

composed = transforms.Compose([Resacle((256, 256)), ToTensor()])

在数据集应用变换的例子

dataset = FaceLandMarksDataset(csv_file='data/faces/face_landmarks.csv', root_dir='data/faces')
sample = dataset[0]
print("变换前")
print(type(sample["image"]), sample["image"].shape)

print("变换后")
sample = composed(sample)
print(type(sample["image"]), sample["image"].size())

运行结果

迭代器

torch.utils.data.Dataloader是一个功能丰富的迭代器，其功能包括但不限于
1.批量处理数据
2, 打乱数据
3, 使用多线程并行加载数据

#重新定义数据集
transformed_dataset = FaceLandMarksDataset(csv_file='data/faces/face_landmarks.csv', root_dir='data/faces', transform=composed)

#数据加载器
batch_size = 4
data_loader = DataLoader(transformed_dataset, batch_size = batch_size, shuffle=True)

搭建神经网络并训练

class model(nn.Module):

    def __init__(self):
        super(model, self).__init__()
        self.fc = nn.Linear(256 * 256 * 3, 68 * 2)

    def forward(self, input):
        input = input.view(-1, self.num_flatten(input))
        out = self.fc(input)
        return out

    def num_flatten(self, input):
        size = input.size()[1:]
        result = 1
        for i in size:
            result *= i

        return result


net = model()
loss = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

#开始训练
num_epoches = 2
for epoch in range(num_epoches):
    for i, data in enumerate(data_loader):
        image, landmarks = data["image"], data["landmark"]
        #前向传播
        optimizer.zero_grad()
        out = net(image)
        landmarks = landmarks.view(out.size())
        lossValue = loss(out, landmarks)
        lossValue.backward()
        optimizer.step()

完整代码

#数据加载和处理

import torch
import os
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, Dataset
from torchvision import  transforms, utils
import torch.nn as nn
import torch.optim as optim

#忽略警告
import warnings
warnings.filterwarnings("ignore")

plt.ion()

#数据集类
class FaceLandMarksDataset(Dataset):

    #                  #csv文件的路径 包含所有图片文件的目录 变换
    def __init__(self, csv_file, root_dir, transform = None):
        self.landmark_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

    #返回数据集长度
    def __len__(self):
        return len(self.landmark_frame)

    #根据索引返回某个具体数据：字典（图像和对应的标记点）
    def __getitem__(self, idx):
        #读取图像
        img_name = os.path.join(self.root_dir, self.landmark_frame.iloc[idx, 0])
        image = io.imread(img_name)
        #读取标记点
        landmarks = self.landmark_frame.iloc[idx, 1:]
        landmarks = np.array(landmarks)
        landmarks = np.array(landmarks, dtype=np.single).reshape(-1, 2)
        sample = {'image':image, 'landmark':landmarks}

        if self.transform:
            sample = self.transform(sample)

        return sample

#定义变换：转换成固定大小
class Resacle(object):

    def __init__(self, out_size:tuple):
        self.out_size = out_size

    def __call__(self, sample):
        image, landmarks = sample["image"], sample["landmark"]
        h , w = image.shape[:2]
        new_h, new_w = self.out_size
        new_h ,new_w = int(new_h), int(new_w)
        image = transform.resize(image, (new_h, new_w))

        landmarks = landmarks * [new_w / w, new_h/h]

        sample = {"image": image, "landmark": landmarks}

        return sample

#转化成Tensor
class ToTensor(object):

    def __call__(self, sample):
        image, landmarks = sample["image"], sample["landmark"]
        #交换颜色轴
        image = image.transpose((2, 0, 1))

        image = torch.from_numpy(image).float()
        landmarks = torch.from_numpy(landmarks).float()

        sample = {"image":image, "landmark":landmarks}

        return sample


#组合变换
composed = transforms.Compose([Resacle((256, 256)), ToTensor()])

#重新定义数据集
transformed_dataset = FaceLandMarksDataset(csv_file='data/faces/face_landmarks.csv', root_dir='data/faces', transform=composed)

#数据加载器
batch_size = 4
data_loader = DataLoader(transformed_dataset, batch_size = batch_size, shuffle=True)

class model(nn.Module):

    def __init__(self):
        super(model, self).__init__()
        self.fc = nn.Linear(256 * 256 * 3, 68 * 2)

    def forward(self, input):
        input = input.view(-1, self.num_flatten(input))
        out = self.fc(input)
        return out

    def num_flatten(self, input):
        size = input.size()[1:]
        result = 1
        for i in size:
            result *= i

        return result


net = model()
loss = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

#开始训练
num_epoches = 2
for epoch in range(num_epoches):
    for i, data in enumerate(data_loader):
        image, landmarks = data["image"], data["landmark"]
        #前向传播
        optimizer.zero_grad()
        out = net(image)
        landmarks = landmarks.view(out.size())
        lossValue = loss(out, landmarks)
        lossValue.backward()
        optimizer.step()