CV-01-AlexNet学习笔记

0、背景知识

在以往的数据分析任务中，人做特征工程，人做分类；机器学习中，人做特征工程，SVM等分类算法做分类；深度学习中，CNN做特征工程，Softmax或SVM做分类。

ImageNet数据集中的图片主要是人文类。

Benchmark：具有商业价值的模型。

0.1、ILSVRC大规模图像识别挑战赛

地址：http://image-net.org/challenges/LSVRC/2012/index

ILSVRC挑战赛的数据集为从ImageNet中抽取出来的一部分，包括1000个类别，训练数据包括120万张图片，1000个类别，验证和测试集一共15万张图片。

 

公式d的含义为预测标签与真实标签相等时为0，min（d）的含义为对5个预测标签分别计算d并从中找出最小值，也就是说只要5个预测结果中有一个正确则为0 ，最终错误率为求平均

 

  

 

一些常用数据集：

 0.2、ImageNet数据集

地址：http://www.image-net.org/

数据集根据WordNet对类别进行分组 ，使得数据集的语义信息更合理。

1、数据预处理

原始图像尺寸：尺寸不统一

图片处理方式：对于给定的图像首先将最短边缩放为256，另一边等比例缩放，然后从中心裁剪出256*256的图像。对图像的每个像素做了减去均值的操作。

2、网络结构

多通道卷积基础知识：

1. 对于多通道图像+多卷积核做卷积，计算方式如下：

 

         如图所示，输入有3个通道，同时有2个卷积核。对于每个卷积核，先在输入3个通道分别作卷积，再将3个通道结果加起来得到卷积输出。所以对于某个卷积层，无论输入图像有多少个通道，输出图像通道数总是等于卷积核数量！

对多通道图像做1x1卷积，其实就是将输入图像于每个通道乘以卷积系数后加在一起，即相当于把原图像中本来各个独立的通道“联通”在了一起

 

总体包括8个可学习的层：5个卷积层、3个全连接层。提出了一些新的技巧或方法。

• LRN层：conv1和conv2之后（Relu之后）
• Max-pooling层：conv1、conv2、conv5之后
• Relu应用在所有的卷积层和全连接层的输出
• Dropout应用在前两个全连接层（Dropout仅能应用于FC，通过以一定的概率将神经元的输出设置为0来实现，这样的话这个神经元不会参与前向传播和反向传播的过程。测试时使用全部神经元，输出乘以dropout的概率来补偿这一过程）

Feature map尺寸计算：

• 卷积时输入图像周围padding a个像素，则输出尺寸为[(输入尺寸+2*a-卷积核尺寸)/卷积滑动步长]+1
• 卷积方式为valid(滑动窗口时超出feature ma则这部分全部舍弃)：[(输入尺寸-卷积核尺寸)/卷积滑动步长]+1
• 卷积方式为same（超出时超出部分pad 0并正常卷积）：[(输入尺寸 - 卷积核尺寸+（卷积核尺寸-步长）)/卷积滑动步长]+1
• 池化操作：[(输入尺寸-池化尺寸)/滑动步长]+1
• 该网络中卷积操作无说明步长时默认为1

AlexNet网络结构细节

AlexNet网络权重参数细节

即下表中的parameters数量的计算，卷积和的参数量加上bias的数量

注：计算第一个全连接层的权重参数数量时，由于上一层是卷积层，则需要把卷积层输出的feature map压缩为一维tensor（也就是6*6*256），其中的每一个元素都可以看成是一个神经元需要与全连接层的每一个神经元相连，因此全连接层有一个神经元就需要6*6*256个权重参数，共有4096个神经元，权重参数的数量即为两者的乘积。第一个全连接层参数占了整个网络参数的一多半，这也是全连接层的一个缺点

AlexNet网络pytorch构建过程

在pytorch的实现中，是以输入图片为224*224搭建的，因此第一个卷积层有2像素的padding，才能得到55*55的特征图。

(224-11+2*2)/4 +1 = 55，其中红色的部分进行了向下取整，这是pytorch的计算特性

AdaptiveAvgPool2d 是为了在输入图片不是224*224时自适应的将feature map池化到6*6继续后续全连接层的运算。 

2.1、Relu激活函数

之前的激活函数比如sigmoid或者tanh具有饱满的非线性，训练起来比非饱满的非线性激活函数收敛慢（下图展示的实验中Relu作为激活函数比tanh快了大约6倍）

Relu激活函数的优点：1.加快网络训练；2.可以防止梯度消失（正半轴梯度均为1）；3.使网络具有稀疏性（负半轴输出值为0，也就是可以使得网络中有些神经元的激活结果为0）

 

2.2局部响应归一化LRN

局部响应归一化简单理解为将卷积核i（通道）位于(x,y)处的激活值值除以周围n个通道（左右各一半）该位置激活值的平方和的运算，如果周围通道该位置的激活值大，那么总的平方和数据就很大，那么i通道该位置的激活值经过LRN就会变小。

其中的超参数是由验证集确定的

idea：LRN模拟了大脑中的侧抑制机制，但大脑中侧抑制考虑的是周围神经元的活动，而此处使用相邻通道似乎是不work的原因，理论上是想让网络的不同feature map学习不同的特征，但应用于不同通道的同一位置可能就导致了不同feature map学习了不同位置的特征（如果有一个通道a位置的响应很高，那么它的相邻通道该位置的激活就很低），并没有起到想让不同的神经元学习不同特征的功能，而卷积神经网络学习到的特征又具有位置不变性。

2.3、Overlapping Pooling

传统的池化层的kernel size z 与 步长 s 相等或者s>z时，此时的池化是没有重叠的；作者使用了s=2，z=3来进行重叠池化，这一方法使得top-1和top-5错误率分别下降了0.4%和0.3%，作者还强调降低了过拟合的可能性

 

3、训练细节

batch-size=128，momentum=0.9，weight decay = 0.05

权重初始化：

• 均值为0，标准差0.01的高斯分布初始化每一层的w
• conv2，conv4，conv5以及全连接的隐藏层bias初始化为1，其他初始化为0.（提供relu正的输入加速了早期的学习）
• 初始学习率0.01，验证错误率不动时学习率除以10

 

4、如何减轻过拟合

4.1、数据增强

本文使用了两种数据增强方法：

• 对原始256*256的图像水平翻转，然后从原始图片和翻转之后的图片中随机裁剪224*224的像素块作为网络的输入。这使得训练图像的数量扩大了2048倍:（256-224）^2 *2

• 测试时从图像的上下左右和中心分别裁剪出五张图用来预测结果（代码实现时再水平或垂直翻转得到10张）

• stack之后得到 shape=[B,10,C,H,W] 的数据 

• 修改RGB三通道像素值的强度：具体来说，对ImageNet训练数据集进行PCA获得特征值和特征向量乘以随机倍数，然后加到原始图片的各通道上

4.2、Dropout

• 训练时以概率P随机使一些神经元的权重为0
• 测试时使用全部的神经元，并且神经元的输出值需要乘以P 

5.1、实验结果及历史意义

5.2、一些有趣的现象

这种不同GPU的卷积核偏好不同类型特征发生在每次运行期间，并且独立于任何特定的随机权重初始化。

相似的图片的在最后一个隐藏层输出的特征向量具有较小的欧式距离（图中第一个为测试图片，剩下5个为欧式距离最小的训练集图片），因此这个现象可以启发我们用卷积神经网络处理过的高级特征来做图像检索、图像聚类和图像编码等任务。

6、论文总结

本文关键点：

• 大量带标签的数据——ImageNet数据集的出现
• 高性能计算资源——GPU
• 合理的算法模型——CNN

创新点：

• 提出Relu激活函数
• 利用LRN提升网络泛化能力
• 采用overlapping pooling提升准确率
• 采用随机裁剪和色彩扰动进行数据增强
• 采用Dropout减轻过拟合

启发点：

 

• 卷积核学习到频率、方向和颜色特征
• 相似图片具有相近的高级特征
• 图像检索可基于高级特征，效果应该优于基于原始图像

7、代码实现

几个关键函数：

7.1、alexnet_inference.py

 alexnet预训练模型下载地址：https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth

# -*- coding: utf-8 -*-

import os
os.environ['NLS_LANG'] = 'SIMPLIFIED CHINESE_CHINA.UTF8'
import time
import json
import torch.nn as nn
import torch
import torchvision.transforms as transforms
from PIL import Image
from matplotlib import pyplot as plt
import torchvision.models as models
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def img_transform(img_rgb, transform=None):
    """
    将数据转换为模型读取的形式
    :param img_rgb: PIL Image
    :param transform: torchvision.transform
    :return: tensor
    """

    if transform is None:
        raise ValueError("找不到transform！必须有transform对img进行处理")

    img_t = transform(img_rgb)
    return img_t


def load_class_names(p_clsnames, p_clsnames_cn):
    """
    加载标签名
    :param p_clsnames:
    :param p_clsnames_cn:
    :return:
    """
    with open(p_clsnames, "r") as f:
        class_names = json.load(f)
    with open(p_clsnames_cn, encoding='UTF-8') as f:  # 设置文件对象
        class_names_cn = f.readlines()
    return class_names, class_names_cn


def get_model(path_state_dict, vis_model=False):
    """
    创建模型，加载参数
    :param path_state_dict:
    :return:
    """
    model = models.alexnet()
    pretrained_state_dict = torch.load(path_state_dict)
    model.load_state_dict(pretrained_state_dict)  # 加载预训练的权重
    model.eval()  # 标记是测试模型

    if vis_model:
        from torchsummary import summary
        summary(model, input_size=(3, 224, 224), device="cpu")  # 生成虚拟数据进行forward，记录每个网络层特征图的大小

    model.to(device)
    return model


def process_img(path_img):

    # hard code
    norm_mean = [0.485, 0.456, 0.406]  # 数据统计自ImageNet
    norm_std = [0.229, 0.224, 0.225]
    inference_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop((224, 224)),
        transforms.ToTensor(),  # 将数据变为tensor形式，且值域变到[0,1]
        transforms.Normalize(norm_mean, norm_std),
    ])

    # path --> img
    img_rgb = Image.open(path_img).convert('RGB')

    # img --> tensor
    img_tensor = img_transform(img_rgb, inference_transform)  # 数据预处理
    img_tensor.unsqueeze_(0)        # chw --> bchw
    img_tensor = img_tensor.to(device)

    return img_tensor, img_rgb


if __name__ == "__main__":

    # config
    path_state_dict = os.path.join(BASE_DIR, "..", "data", "alexnet-owt-4df8aa71.pth")
    # path_img = os.path.join(BASE_DIR, "..", "data", "Golden Retriever from baidu.jpg")
    path_img = os.path.join(BASE_DIR, "..", "data", "tiger cat.jpg")
    path_classnames = os.path.join(BASE_DIR, "..", "data", "imagenet1000.json")
    path_classnames_cn = os.path.join(BASE_DIR, "..", "data", "imagenet_classnames.txt")

    # load class names
    cls_n, cls_n_cn = load_class_names(path_classnames, path_classnames_cn)

    # 1/5 load img
    img_tensor, img_rgb = process_img(path_img)

    # 2/5 load model
    alexnet_model = get_model(path_state_dict, True)

    # 3/5 inference  tensor --> vector
    with torch.no_grad():  # 由于是预测，因此不需要梯度反向传播，可以节约时间和内存
        time_tic = time.time()
        outputs = alexnet_model(img_tensor)
        time_toc = time.time()

    # 4/5 index to class names
    _, pred_int = torch.max(outputs.data, 1)  # 会返回最大的值和索引，因此我们只获取索引即可，topk同理
    _, top5_idx = torch.topk(outputs.data, 5, dim=1)

    pred_idx = int(pred_int.cpu().numpy())
    pred_str, pred_cn = cls_n[pred_idx], cls_n_cn[pred_idx]
    print("img: {} is: {}\n{}".format(os.path.basename(path_img), pred_str, pred_cn))
    print("time consuming:{:.2f}s".format(time_toc - time_tic))

    # 5/5 visualization
    plt.imshow(img_rgb)
    plt.title("predict:{}".format(pred_str))
    top5_num = top5_idx.cpu().numpy().squeeze()
    text_str = [cls_n[t] for t in top5_num]
    for idx in range(len(top5_num)):
        plt.text(5, 15+idx*30, "top {}:{}".format(idx+1, text_str[idx]), bbox=dict(fc='yellow'))
    plt.show()

 7.2、卷积核可视化

得到结果文件后打开终端到结果文件目录使用tensorboard --logdir=./ 命令来查看可视化结果

第一个卷积可视化结果

第二个卷积可视化结果

图片经过第一个卷积层后的结果（64个卷积核）

# -*- coding: utf-8 -*-

import os
import torch
import torch.nn as nn
from PIL import Image
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
import torchvision.utils as vutils
import torchvision.models as models
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

if __name__ == "__main__":

    log_dir = os.path.join(BASE_DIR, "..", "results")
    # ----------------------------------- kernel visualization -----------------------------------
    writer = SummaryWriter(log_dir=log_dir, filename_suffix="_kernels")

    # m1
    # alexnet = models.alexnet(pretrained=True)

    # m2
    path_state_dict = os.path.join(BASE_DIR, "..", "data", "alexnet-owt-4df8aa71.pth")
    alexnet = models.alexnet()
    pretrained_state_dict = torch.load(path_state_dict)
    alexnet.load_state_dict(pretrained_state_dict)

    kernel_num = -1
    vis_max = 1  # 可视化的最大卷积数, 1表示可视化 0 和 1 卷积层
    for sub_module in alexnet.modules():
        if not isinstance(sub_module, nn.Conv2d):  # 循环遍历出卷积层，非卷积层跳过
            continue
        kernel_num += 1
        if kernel_num > vis_max:  # 如果当前卷积层大于要可视化的卷积层则退出循环
            break

        kernels = sub_module.weight
        c_out, c_int, k_h, k_w = tuple(kernels.shape)  # 获得输入输出通道数，卷积核宽和高

        # 拆分channel
        for o_idx in range(c_out):
            kernel_idx = kernels[o_idx, :, :, :].unsqueeze(1)  # 获得(3, h, w), 但是make_grid需要 BCHW，这里拓展C维度变为（3， 1， h, w）
            kernel_grid = vutils.make_grid(kernel_idx, normalize=True, scale_each=True, nrow=c_int)
            writer.add_image('{}_Convlayer_split_in_channel'.format(kernel_num), kernel_grid, global_step=o_idx)  #将图片写入tensorboard文件

        kernel_all = kernels.view(-1, 3, k_h, k_w)  # 3, h, w 可视化RGB三通道的卷积核，view的功能类似numpy中的resize
        kernel_grid = vutils.make_grid(kernel_all, normalize=True, scale_each=True, nrow=8)  # c, h, w
        writer.add_image('{}_all'.format(kernel_num), kernel_grid, global_step=620)

        print("{}_convlayer shape:{}".format(kernel_num, tuple(kernels.shape)))

    # ----------------------------------- feature map visualization -----------------------------------
    writer = SummaryWriter(log_dir=log_dir, filename_suffix="_feature map")

    # 数据
    path_img = os.path.join(BASE_DIR, "..", "data", "tiger cat.jpg")  # your path to image
    normMean = [0.49139968, 0.48215827, 0.44653124]
    normStd = [0.24703233, 0.24348505, 0.26158768]
    norm_transform = transforms.Normalize(normMean, normStd)
    img_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        norm_transform
    ])

    img_pil = Image.open(path_img).convert('RGB')  # 读取图片，转化为RGB模式
    img_tensor = img_transforms(img_pil)
    img_tensor.unsqueeze_(0)  # chw --> bchw

    # 模型
    # alexnet = models.alexnet(pretrained=True)

    # forward
    convlayer1 = alexnet.features[0]  # 取出第一个卷积层
    fmap_1 = convlayer1(img_tensor)

    # 预处理
    fmap_1.transpose_(0, 1)  # bchw=(1, 64, 55, 55) --> (64, 1, 55, 55) make_grid接收的数据为BCHW形式，B为图片数，C为图片通道数
    fmap_1_grid = vutils.make_grid(fmap_1, normalize=True, scale_each=True, nrow=8)

    writer.add_image('feature map in conv1', fmap_1_grid, global_step=620)
    writer.close()

7.3、模型训练

# -*- coding: utf-8 -*-
"""
# @file name  : train_alexnet.py
# @author     : TingsongYu https://github.com/TingsongYu
# @date       : 2020-02-14
# @brief      : alexnet traning
"""

import os
import numpy as np
import torch.nn as nn
import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torch.optim as optim
from matplotlib import pyplot as plt
import torchvision.models as models
from tools.my_dataset import CatDogDataset

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def get_model(path_state_dict, vis_model=False):
    """
    创建模型，加载参数
    :param path_state_dict:
    :return:
    """
    model = models.alexnet()
    pretrained_state_dict = torch.load(path_state_dict)
    model.load_state_dict(pretrained_state_dict)

    if vis_model:
        from torchsummary import summary
        summary(model, input_size=(3, 224, 224), device="cpu")

    model.to(device)
    return model


if __name__ == "__main__":

    # config
    data_dir = os.path.join(BASE_DIR, "..", "data", "train")
    path_state_dict = os.path.join(BASE_DIR, "..", "data", "alexnet-owt-4df8aa71.pth")
    num_classes = 2

    MAX_EPOCH = 3       # 可自行修改
    BATCH_SIZE = 128    # 可自行修改
    LR = 0.001          # 可自行修改
    log_interval = 1    # 多少个iteration打印一次结果
    val_interval = 1    # 多少个epoch 执行一次验证集
    classes = 2
    start_epoch = -1
    lr_decay_step = 1   # 可自行修改

    # ============================ step 1/5 数据 ============================
    norm_mean = [0.485, 0.456, 0.406]
    norm_std = [0.229, 0.224, 0.225]

    # 训练集图像预处理
    train_transform = transforms.Compose([
        transforms.Resize((256)),  # 将短边resize成256，另一条边等比例缩放，而(256, 256) 则是直接resize成256*256
        transforms.CenterCrop(256),  # 中心裁剪出256*256的图片
        transforms.RandomCrop(224),  # 随机裁剪224*224的图片
        transforms.RandomHorizontalFlip(p=0.5),  # 随机水平翻转
        transforms.ToTensor(),
        transforms.Normalize(norm_mean, norm_std),
    ])

    # 验证集图像预处理
    normalizes = transforms.Normalize(norm_mean, norm_std)
    valid_transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.TenCrop(224, vertical_flip=False),  # list形式：10个BCHW形式的数据
        transforms.Lambda(lambda crops: torch.stack([normalizes(transforms.ToTensor()(crop)) for crop in crops])),   # 处理后数据 shape=[B,10,C,H,W]
    ])

    # 构建MyDataset实例
    train_data = CatDogDataset(data_dir=data_dir, mode="train", transform=train_transform)
    valid_data = CatDogDataset(data_dir=data_dir, mode="valid", transform=valid_transform)

    # 构建DataLoder
    train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
    valid_loader = DataLoader(dataset=valid_data, batch_size=4)

    # ============================ step 2/5 模型 ============================
    alexnet_model = get_model(path_state_dict, False)

    num_ftrs = alexnet_model.classifier._modules["6"].in_features  # 修改最后的全连接层
    alexnet_model.classifier._modules["6"] = nn.Linear(num_ftrs, num_classes)

    alexnet_model.to(device)
    # ============================ step 3/5 损失函数 ============================
    criterion = nn.CrossEntropyLoss()
    # ============================ step 4/5 优化器 ============================
    flag = 0
    # flag = 1  #  为不同层设置不同学习率
    if flag:
        fc_params_id = list(map(id, alexnet_model.classifier.parameters()))  # 返回的是parameters的 内存地址
        base_params = filter(lambda p: id(p) not in fc_params_id, alexnet_model.parameters())
        optimizer = optim.SGD([
            {'params': base_params, 'lr': LR * 0.1},  # 0
            {'params': alexnet_model.classifier.parameters(), 'lr': LR}], momentum=0.9)

    else:
        optimizer = optim.SGD(alexnet_model.parameters(), lr=LR, momentum=0.9)  # 选择优化器

    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=lr_decay_step, gamma=0.1)  # 设置学习率下降策略
    # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(patience=5)

# ============================ step 5/5 训练 ============================
    train_curve = list()
    valid_curve = list()

    for epoch in range(start_epoch + 1, MAX_EPOCH):

        loss_mean = 0.
        correct = 0.
        total = 0.

        alexnet_model.train()
        for i, data in enumerate(train_loader):

            # if i > 1:
            #     break

            # forward
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = alexnet_model(inputs)

            # backward
            optimizer.zero_grad()
            loss = criterion(outputs, labels)
            loss.backward()

            # update weights
            optimizer.step()

            # 统计分类情况
            _, predicted = torch.max(outputs.data, dim=1)  # 获取最大值索引，max函数会返回最大值和最大值的索引
            total += labels.size(0)
            correct += (predicted == labels).squeeze().cpu().sum().numpy()  # squeeze 将dim=1的维度去掉

            # 打印训练信息
            loss_mean += loss.item()
            train_curve.append(loss.item())
            if (i+1) % log_interval == 0:
                loss_mean = loss_mean / log_interval
                print("Training:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
                    epoch, MAX_EPOCH, i+1, len(train_loader), loss_mean, correct / total))
                loss_mean = 0.

        scheduler.step()  # 更新学习率

        # validate the model
        if (epoch+1) % val_interval == 0:

            correct_val = 0.
            total_val = 0.
            loss_val = 0.
            alexnet_model.eval()
            with torch.no_grad():
                for j, data in enumerate(valid_loader):
                    inputs, labels = data
                    inputs, labels = inputs.to(device), labels.to(device)

                    bs, ncrops, c, h, w = inputs.size()     # [4, 10, 3, 224, 224]
                    outputs = alexnet_model(inputs.view(-1, c, h, w))  # view处理后获得模型接收的BCHW格式的数据，模型输出为[bs*ncrops= 40,class_num=2]
                    outputs_avg = outputs.view(bs, ncrops, -1).mean(1)  # viem再对数据维度进行变换，然后对ncrops维度求平均获得每张图片的预测

                    loss = criterion(outputs_avg, labels)

                    _, predicted = torch.max(outputs_avg.data, 1)
                    total_val += labels.size(0)
                    correct_val += (predicted == labels).squeeze().cpu().sum().numpy()

                    loss_val += loss.item()

                loss_val_mean = loss_val/len(valid_loader)
                valid_curve.append(loss_val_mean)
                print("Valid:\t Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
                    epoch, MAX_EPOCH, j+1, len(valid_loader), loss_val_mean, correct_val / total_val))
            alexnet_model.train()

    train_x = range(len(train_curve))
    train_y = train_curve

    train_iters = len(train_loader)
    valid_x = np.arange(1, len(valid_curve)+1) * train_iters*val_interval  # 由于valid中记录的是epochloss，需要对记录点进行转换到iterations
    valid_y = valid_curve

    plt.plot(train_x, train_y, label='Train')
    plt.plot(valid_x, valid_y, label='Valid')

    plt.legend(loc='upper right')
    plt.ylabel('loss value')
    plt.xlabel('Iteration')
    plt.show()

DataLoader构建 ：

• 构建时最少需要定义getitem和len两个函数
• 尽可能的减少getitem中的代码量可以加快执行效率 

# -*- coding: utf-8 -*-

import numpy as np
import torch
import os
import random
from PIL import Image
from torch.utils.data import Dataset

random.seed(1)
rmb_label = {"1": 0, "100": 1}


class CatDogDataset(Dataset):
    def __init__(self, data_dir, mode="train", split_n=0.9, rng_seed=620, transform=None):
        """
        分类任务的Dataset
        :param data_dir: str, 数据集所在路径
        :param transform: torch.transform，数据预处理
        """
        self.mode = mode
        self.data_dir = data_dir
        self.rng_seed = rng_seed
        self.split_n = split_n  # 训练集所占比例
        self.data_info = self._get_img_info()  # data_info存储所有图片路径和标签，在DataLoader中通过index读取样本
        self.transform = transform

    def __getitem__(self, index):
        path_img, label = self.data_info[index]
        img = Image.open(path_img).convert('RGB')     # 0~255

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

        return img, label

    def __len__(self):
        if len(self.data_info) == 0:
            raise Exception("\ndata_dir:{} is a empty dir! Please checkout your path to images!".format(self.data_dir))
        return len(self.data_info)

    def _get_img_info(self):
        """
        获取图片路径及标签并存储成list形式
        """

        img_names = os.listdir(self.data_dir)  # 获取文件夹下所有的文件名
        img_names = list(filter(lambda x: x.endswith('.jpg'), img_names))  # 如果文件以jpg结尾添加到list中

        random.seed(self.rng_seed)  # 为了固定图片排序
        random.shuffle(img_names)

        img_labels = [0 if n.startswith('cat') else 1 for n in img_names]  # 设置标签 cat为0

        split_idx = int(len(img_labels) * self.split_n)  # 25000* 0.9 = 22500
        # split_idx = int(100 * self.split_n)
        if self.mode == "train":
            img_set = img_names[:split_idx]     # 数据集90%训练
            # img_set = img_names[:22500]     #  hard code 数据集90%训练
            label_set = img_labels[:split_idx]
        elif self.mode == "valid":
            img_set = img_names[split_idx:]
            label_set = img_labels[split_idx:]
        else:
            raise Exception("self.mode 无法识别，仅支持(train, valid)")

        path_img_set = [os.path.join(self.data_dir, n) for n in img_set]  # 拼接获得绝对路径
        data_info = [(n, l) for n, l in zip(path_img_set, label_set)]

        return data_info

本文为深度之眼paper论文班的学习笔记，仅供自己学习使用，如有问题欢迎讨论！关于课程可以扫描下图二维码