Pytorch模型训练实用教程学习笔记：一、数据加载和transforms方法总结

前言

最近在重温Pytorch基础，然而Pytorch官方文档的各种API是根据字母排列的，并不适合学习阅读。
于是在gayhub上找到了这样一份教程《Pytorch模型训练实用教程》，写得不错，特此根据它来再学习一下Pytorch。
仓库地址：https://github.com/TingsongYu/PyTorch_Tutorial

数据集转换

首先练习对数据集的处理方式。
这里采用的是cifar-10数据集，从官网下载下来的格式长这样：

data_batch_1-5是训练集，test_batch是测试集。
这种形式不利于直观阅读，因此利用pickle来对其进行转换，转换成png格式。
另附cifar-10数据集备份：https://pan.baidu.com/s/1uxQ7RGjLChe99fpiotM7jw?pwd=8888

转换代码

# coding:utf-8
""" 将cifar10的data_batch_12345 转换成 png格式的图片 每个类别单独存放在一个文件夹，文件夹名称为0-9 """
from imageio import imwrite
import numpy as np
import os
import pickle

data_dir = os.path.join("..", "..", "Data", "cifar-10-batches-py")
train_o_dir = os.path.join("..", "..", "Data", "cifar-10-png", "raw_train")
test_o_dir = os.path.join("..", "..", "Data", "cifar-10-png", "raw_test")


# 解压缩，返回解压后的字典
def unpickle(file):
    with open(file, 'rb') as fo:
        dict_ = pickle.load(fo, encoding='bytes')
    return dict_


def my_mkdir(my_dir):
    if not os.path.isdir(my_dir):
        os.makedirs(my_dir)


if __name__ == '__main__':
    # 生成训练集图片
    for j in range(1, 6):
        data_path = os.path.join(data_dir, "data_batch_" + str(j))  # data_batch_12345
        train_data = unpickle(data_path)
        print(data_path + " is loading...")

        for i in range(0, 10000):
            img = np.reshape(train_data[b'data'][i], (3, 32, 32))
            img = img.transpose(1, 2, 0) # （channels,imagesize,imagesize）转换成（imagesize,imagesize,channels)

            label_num = str(train_data[b'labels'][i])
            o_dir = os.path.join(train_o_dir, label_num)
            my_mkdir(o_dir)

            img_name = label_num + '_' + str(i + (j - 1) * 10000) + '.png'
            img_path = os.path.join(o_dir, img_name)
            imwrite(img_path, img)
        print(data_path + " loaded.")

    print("test_batch is loading...")

    # 生成测试集图片
    test_data_path = os.path.join(data_dir, "test_batch")
    test_data = unpickle(test_data_path)
    for i in range(0, 10000):
        img = np.reshape(test_data[b'data'][i], (3, 32, 32))
        img = img.transpose(1, 2, 0)

        label_num = str(test_data[b'labels'][i])
        o_dir = os.path.join(test_o_dir, label_num)
        my_mkdir(o_dir)

        img_name = label_num + '_' + str(i) + '.png'
        img_path = os.path.join(o_dir, img_name)
        imwrite(img_path, img)

    print("test_batch loaded.")

转换后的数据集长这样：

注：cifar-10共有10个类别，每张图片大小为32x32像素。

数据集划分

下面对数据集划分，这里只是为了演示学习，因此仅对原本的测试集数据进行划分，划分比例为8：1：1。
代码：

# coding: utf-8
""" 将原始数据集进行划分成训练集、验证集和测试集 """

import os
import glob
import random
import shutil

dataset_dir = os.path.join("..", "..", "Data", "cifar-10-png", "raw_test")
train_dir = os.path.join("..", "..", "Data", "train")
valid_dir = os.path.join("..", "..", "Data", "valid")
test_dir = os.path.join("..", "..", "Data", "test")

train_per = 0.8
valid_per = 0.1
test_per = 0.1


def makedir(new_dir):
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)


if __name__ == '__main__':

    for root, dirs, files in os.walk(dataset_dir):
        for sDir in dirs:
            imgs_list = glob.glob(os.path.join(root, sDir, '*.png'))  # glob匹配路径，匹配所有png格式图片
            random.seed(666)
            random.shuffle(imgs_list)
            imgs_num = len(imgs_list)

            train_point = int(imgs_num * train_per)
            valid_point = int(imgs_num * (train_per + valid_per))

            for i in range(imgs_num):
                if i < train_point:
                    out_dir = os.path.join(train_dir, sDir)
                elif i < valid_point:
                    out_dir = os.path.join(valid_dir, sDir)
                else:
                    out_dir = os.path.join(test_dir, sDir)

                makedir(out_dir)
                out_path = os.path.join(out_dir, os.path.split(imgs_list[i])[-1])
                shutil.copy(imgs_list[i], out_path)

            print('Class:{}, train:{}, valid:{}, test:{}'.format(sDir, train_point, valid_point-train_point, imgs_num-valid_point))

划分好的数据如图所示：

数据集加载文件

通常来说，数据加载都是通过txt文件进行路径读取，在我之前的博文【目标检测】YOLOv5跑通VOC2007数据集(修复版)也实现过这一效果，这里不作赘述。

代码：

# coding:utf-8
import os
''' 为数据集生成对应的txt文件 '''

train_txt_path = os.path.join("..", "..", "Data", "train.txt")
train_dir = os.path.join("..", "..", "Data", "train")

valid_txt_path = os.path.join("..", "..", "Data", "valid.txt")
valid_dir = os.path.join("..", "..", "Data", "valid")


def gen_txt(txt_path, img_dir):
    f = open(txt_path, 'w')
    
    for root, s_dirs, _ in os.walk(img_dir, topdown=True):  # 获取 train文件下各文件夹名称
        for sub_dir in s_dirs:
            i_dir = os.path.join(root, sub_dir)             # 获取各类的文件夹 绝对路径
            img_list = os.listdir(i_dir)                    # 获取类别文件夹下所有png图片的路径
            for i in range(len(img_list)):
                if not img_list[i].endswith('png'):         # 若不是png文件，跳过
                    continue
                label = img_list[i].split('_')[0]
                img_path = os.path.join(i_dir, img_list[i])
                line = img_path + ' ' + label + '\n'
                f.write(line)
    f.close()


if __name__ == '__main__':
    gen_txt(train_txt_path, train_dir)
    gen_txt(valid_txt_path, valid_dir)

生成结果：

构建Dataset

数据加载通常使用Pytorch提供的DataLoader，在此之前，需要构建自己的数据集类，在数据集类中，可以包含transform一些数据处理方式。

from PIL import Image
from torch.utils.data import Dataset


class MyDataset(Dataset):
    def __init__(self, txt_path, transform=None, target_transform=None):
        fh = open(txt_path, 'r')
        imgs = []
        for line in fh:
            line = line.rstrip()
            words = line.split()
            imgs.append((words[0], int(words[1])))

        self.imgs = imgs        # 最主要就是要生成这个list， 然后DataLoader中给index，通过getitem读取图片数据
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        fn, label = self.imgs[index]
        img = Image.open(fn).convert('RGB')     # 像素值 0~255，在transfrom.totensor会除以255，使像素值变成 0~1

        if self.transform is not None:
            img = self.transform(img)   # 在这里做transform，转为tensor等等

        return img, label

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

注：在DataLoader中，会调用__getitem__方法，需要返回的是data+label的形式。

数据标准化

数据标准化(Normalize)是非常常见的数据处理方式，在Pytorch中的调用示例：

normMean = [0.4948052, 0.48568845, 0.44682974]
normStd = [0.24580306, 0.24236229, 0.2603115]
normTransform = transforms.Normalize(normMean, normStd)

注：这里的均值和标准差是需要自定义的。

下面这段程序就是随机读取CNum张图片，来计算三通道的均值和标准差。

# coding: utf-8

import numpy as np
import cv2
import random
import os

""" 随机挑选CNum张图片，进行按通道计算均值mean和标准差std 先将像素从0～255归一化至 0-1 再计算 """


train_txt_path = os.path.join("..", "..", "Data/train.txt")

CNum = 2000     # 挑选多少图片进行计算

img_h, img_w = 32, 32
imgs = np.zeros([img_w, img_h, 3, 1])
means, stdevs = [], []

with open(train_txt_path, 'r') as f:
    lines = f.readlines()
    random.shuffle(lines)   # shuffle , 随机挑选图片

    for i in range(CNum):
        img_path = lines[i].rstrip().split()[0]

        img = cv2.imread(img_path)
        img = cv2.resize(img, (img_h, img_w))

        img = img[:, :, :, np.newaxis]
        imgs = np.concatenate((imgs, img), axis=3)
        print(i)

imgs = imgs.astype(np.float32)/255.


for i in range(3):
    pixels = imgs[:,:,i,:].ravel()  # 拉成一行
    means.append(np.mean(pixels))
    stdevs.append(np.std(pixels))

means.reverse() # BGR --> RGB
stdevs.reverse()

print("normMean = {}".format(means))
print("normStd = {}".format(stdevs))
print('transforms.Normalize(normMean = {}, normStd = {})'.format(means, stdevs))

transforms方法汇总

对于数据处理，pytorch专门提供的transforms函数，该函数有下列一些方法可以使用。

裁剪——Crop

中心裁剪：transforms.CenterCrop

功能：依据给定的 size 从中心裁剪
参数：
size- (sequence or int)，若为 sequence,则为(h,w)，若为 int，则(size,size)

随机裁剪：transforms.RandomCrop

功能：依据给定的 size 随机裁剪
参数：
size- (sequence or int)，若为 sequence,则为(h,w)，若为 int，则(size,size)

padding-(sequence or int, optional)，此参数是设置填充多少个 pixel。
当为 int 时，图像上下左右均填充 int 个，例如 padding=4，则上下左右均填充 4 个 pixel，若为 32x32，则会变成 40x40。当为 sequence 时，若有 2 个数，则第一个数表示左右扩充多少，第二个数表示上下的。当有 4 个数时，则为左，上，右，下。

fill- (int or tuple) 填充的值是什么（仅当填充模式为 constant 时有用）。int 时，各通道均填充该值，当长度为 3 的 tuple 时，表示 RGB 通道需要填充的值。

padding_mode- 填充模式，这里提供了 4 种填充模式，1.constant，常量。2.edge 按照图片边缘的像素值来填充。3.reflect。 4. symmetric。

随机长宽比裁剪：transforms.RandomResizedCrop

功能：随机大小，随机长宽比裁剪原始图片，最后将图片 resize 到设定好的 size
参数：
size- 输出的分辨率
scale- 随机 crop 的大小区间，如 scale=(0.08, 1.0)，表示随机 crop 出来的图片会在的 0.08倍至 1 倍之间。
ratio- 随机长宽比设置
interpolation- 插值的方法，默认为双线性插值(PIL.Image.BILINEAR)

上下左右中心裁剪：transforms.FiveCrop

功能：对图片进行上下左右以及中心裁剪，获得 5 张图片，返回一个 4D-tensor
参数：
size- (sequence or int)，若为 sequence,则为(h,w)，若为 int，则(size,size)

上下左右中心裁剪后翻转，transforms.TenCrop

功能：对图片进行上下左右以及中心裁剪，然后全部翻转（水平或者垂直），获得 10 张图
片，返回一个 4D-tensor。
参数：
size- (sequence or int)，若为 sequence,则为(h,w)，若为 int，则(size,size)
vertical_flip (bool) - 是否垂直翻转，默认为 flase，即默认为水平翻转

翻转和旋转——Flip and Rotations

依概率 p 水平翻转：transforms.RandomHorizontalFlip(p=0.5)

功能：依据概率 p 对 PIL 图片进行水平翻转
参数：
p- 概率，默认值为 0.5

依概率 p 垂直翻转：transforms.RandomVerticalFlip(p=0.5)

功能：依据概率 p 对 PIL 图片进行垂直翻转
参数：
p- 概率，默认值为 0.5

随机旋转：transforms.RandomRotation

功能：依 degrees 随机旋转一定角度
参数：
degress- (sequence or float or int) ，若为单个数，如 30，则表示在（-30，+30）之间随机旋转，若为 sequence，如(30，60)，则表示在 30-60 度之间随机旋转

图像变换

图像缩放：transforms.Resize

功能：重置图像分辨率
参数：
size- If size is an int, if height > width, then image will be rescaled to (size * height / width, size)，所以建议 size 设定为 h*w
interpolation- 插值方法选择，默认为 PIL.Image.BILINEAR

标准化：transforms.Normalize

class torchvision.transforms.Normalize(mean, std)
功能：对数据按通道进行标准化，即先减均值，再除以标准差，注意是 h * w * c

转为 tensor，并归一化至[0-1]：transforms.ToTensor

功能：将 PIL Image 或者 ndarray 转换为 tensor，并且归一化至[0-1]
注意事项：归一化至[0-1]是直接除以 255，若自己的 ndarray 数据尺度有变化，则需要自行
修改。

填充：transforms.Pad

功能：对图像进行填充
参数：
padding-(sequence or int, optional)，此参数是设置填充多少个 pixel。
当为 int 时，图像上下左右均填充 int 个，例如 padding=4，则上下左右均填充 4 个 pixel，若为 32x32，则会变成 40x40。
fill- (int or tuple) 填充的值是什么
padding_mode- 填充模式，这里提供了 4 种填充模式，1.constant，常量。2.edge 按照图片边缘的像素值来填充。3.reflect 4. symmetric

修改亮度、对比度和饱和度：transforms.ColorJitter

功能：修改修改亮度、对比度和饱和度

转灰度图：transforms.Grayscale

功能：将图片转换为灰度图
参数：
num_output_channels- (int) ，当为 1 时，正常的灰度图，当为 3 时， 3 channel with r == g == b

线性变换：transforms.LinearTransformation()

功能：对矩阵做线性变化

仿射变换：transforms.RandomAffine

功能：仿射变换

依概率 p 转为灰度图：transforms.RandomGrayscale

功能：依概率 p 将图片转换为灰度图，若通道数为 3，则 3 channel with r == g == b

将数据转换为 PILImage：transforms.ToPILImage

功能：将 tensor 或者 ndarray 的数据转换为 PIL Image 类型数据
参数：
mode- 为 None 时，为 1 通道， mode=3 通道默认转换为 RGB，4 通道默认转换为 RGBA

transforms操作

transforms.RandomChoice(transforms)

功能：从给定的一系列 transforms 中选一个进行操作

transforms.RandomApply(transforms, p=0.5)

功能：给一个 transform 加上概率，依概率进行操作

transforms.RandomOrder

功能：将 transforms 中的操作随机打乱

使用示例：
例如，想对数据进行缩放、随机裁剪、归一化和标准化，可以这样进行设置：

# 数据预处理设置
normMean = [0.4948052, 0.48568845, 0.44682974]
normStd = [0.24580306, 0.24236229, 0.2603115]
normTransform = transforms.Normalize(normMean, normStd)
trainTransform = transforms.Compose([
    transforms.Resize(32),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    normTransform
])
train_data = MyDataset(txt_path=train_txt_path, transform=trainTransform)