深度学习——图像分类模型最简单的train.py

import os 
import sys
import json 
import torch
import torch.nn as nn
from torchvision import transforms, datasets 
import torch.optim as optim 
from tqdm import tqdm  
#from classic_models.alexnet import AlexNet
from classic_models.googlenet_v1 import  GoogLeNet

def main():
    # 判断可用设备
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    # 注意改成自己的数据集路径
    data_path = "G:\\flower"
    assert os.path.exists(data_path), "{} path does not exist.".format(data_path) 

    # 数据预处理与增强
    """ 
    ToTensor()能够把灰度范围从0-255变换到0-1之间的张量.
    transform.Normalize()则把0-1变换到(-1,1). 具体地说, 对每个通道而言, Normalize执行以下操作: image=(image-mean)/std
    其中mean和std分别通过(0.5,0.5,0.5)和(0.5,0.5,0.5)进行指定。原来的0-1最小值0则变成(0-0.5)/0.5=-1; 而最大值1则变成(1-0.5)/0.5=1. 
    也就是一个均值为0, 方差为1的正态分布. 这样的数据输入格式可以使神经网络更快收敛。
    """
    data_transform = {
        "train": transforms.Compose([transforms.Resize(224),   # 将图片的短边缩放到224，图片的长边和短边的比值不变，即不能保证每张图片都是224*224大小，那么下一步的裁剪就有必要了
                                     transforms.CenterCrop(224),  # 由中心向两边进行裁剪，裁剪的尺寸为224*224
                                     transforms.ToTensor(),       # 可以将PIL和numpy格式的数据从[0,255]范围转换到[0,1] 。另外原始数据的shape是（H x W x C），这步后shape会变为（C x H x W）
                                     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]), # Normalize(mean, std, inplace=False)，三通道中Normalize里面一般是Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))，将上一步的数据范围由[0,1]转换为[-1,1]

        "val": transforms.Compose([transforms.Resize((224, 224)),  # val不需要任何数据增强
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}


    # 使用ImageFlolder加载数据集中的图像，并使用指定的预处理操作来处理图像， ImageFlolder会同时返回图像和对应的标签。 (image path, class_index) tuples
    train_dataset = datasets.ImageFolder(root=os.path.join(data_path, "train"), transform=data_transform["train"]) # root：图片存储的根目录，即各类别文件夹所在目录的上一级目录。
    validate_dataset = datasets.ImageFolder(root=os.path.join(data_path, "val"), transform=data_transform["val"])  # transform：对图片进行预处理的操作（函数）。在data_transform中已经定义好
    train_num = len(train_dataset) # 计算train_dataset里面的图片个数
    val_num = len(validate_dataset) # 计算validate_dataset里面的图片个数

    # 使用class_to_idx给类别一个index，作为训练时的标签： {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
    flower_list = train_dataset.class_to_idx # class_to_idx是train_dataset里面的一个函数，返回一个字典，即flower_list是一个字典

    # 创建一个字典，存储index和类别的对应关系，在模型推理阶段会用到。
    cla_dict = dict((val, key) for key, val in flower_list.items())  # items()方法将字典里对应的一对键和值以元组的形式（键, 值），存储为所生成序列里的单个元素

    # 将字典写成一个json文件
    json_str = json.dumps(cla_dict, indent=4)  # json.dumps()是把python对象转换成json对象的一个过程，生成的是字符串。
    with open(os.path.join(data_path, 'class_indices.json') , 'w') as json_file:
        json_file.write(json_str)

    batch_size = 32 # batch_size大小，是超参，可调，如果模型跑不起来，尝试调小batch_size
 
    # 使用 DataLoader 将 ImageFloder 加载的数据集处理成批量（batch）加载模式
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True )
    validate_loader = torch.utils.data.DataLoader(validate_dataset, batch_size=4, shuffle=False ) # 注意，验证集不需要shuffle
    print("using {} images for training, {} images for validation.".format(train_num, val_num))
    
    # 实例化模型，并送进设备
    net = GoogLeNet(num_classes = 5) # 使用GoogleNet来定义网络模型，分类数为5
    #net = AlexNet(num_classes=5 )
    net.to(device)

    # 指定损失函数用于计算损失；指定优化器用于更新模型参数；指定训练迭代的轮数，训练权重的存储地址
    loss_function = nn.CrossEntropyLoss() # 交叉熵函数
    optimizer = optim.Adam(net.parameters(), lr=0.0002)
    epochs = 1
    save_path = os.path.abspath(os.path.join(os.getcwd(), './results/weights/alexnet'))  # os.getcwd()返回当前的文件目录，也就是后面的文件目录
    if not os.path.exists(save_path):    
        os.makedirs(save_path)       # 创建名为save_path的目录

    best_acc = 0.0 # 初始化验证集上最好的准确率，以便后面用该指标筛选模型最优参数。  
    for epoch in range(epochs):
        ############################################################## train ######################################################
        net.train() 
        acc_num = torch.zeros(1).to(device)    # 初始化，用于计算训练过程中预测正确的数量
        sample_num = 0                         # 初始化，用于记录当前迭代中，已经计算了多少个样本
        # tqdm是一个进度条显示器，可以在终端打印出现在的训练进度
        # train_loader：是需要迭代的对象，通常为列表或者生成器，其中包含训练数据。进度条会遍历该对象，并相应地更新进度。
        # file=sys.stdout：这个参数指定了进度条应该写入其输出的位置。在这种情况下，它被设置为sys.stdout，表示标准输出流（通常是控制台）。因此，进度条将显示在控制台中。
        # ncols=100：这个参数设置进度条的宽度，以字符为单位。在这里，它被设置为100个字符
        train_bar = tqdm(train_loader, file=sys.stdout, ncols=100)
        for data in train_bar :
            images, labels = data 
            sample_num += images.shape[0] #[32, 3, 224, 224]
            optimizer.zero_grad()         # 梯度初始化为零，把loss关于weight的导数变成0，避免梯度的叠加效应
            outputs = net(images.to(device)) # output_shape: [batch_size, num_classes]
            pred_class = torch.max(outputs, dim=1)[1] # torch.max 返回值是一个tuple，第一个元素是max值，第二个元素是max值的索引。
            acc_num += torch.eq(pred_class, labels.to(device)).sum()  # 是一个比较操作，它会将预测的类别（pred_class）和标签（labels）进行逐元素的比较，返回一个布尔类型的张量，表示对应位置上两个值是否相等。
                                                                      # sum() 是对布尔类型的张量进行求和操作，将所有为 True 的元素相加，得到一个标量值，表示预测正确的样本数量。
            loss = loss_function(outputs, labels.to(device)) # 求损失
            loss.backward() # 自动求导
            optimizer.step() # 梯度下降

            # print statistics 
            train_acc = acc_num.item() / sample_num 
            # .desc是进度条tqdm中的成员变量，作用是描述信息
            train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,  epochs, loss)

        # validate
        net.eval() #不启用 BatchNormalization 和 Dropout。此时pytorch会自动把BN和DropOut固定住，不会取平均，而是用训练好的值。不然的话，一旦test的batch_size过小，很容易就会因BN层导致模型performance损失较大；
        acc_num = 0.0  # accumulate accurate number per epoch
        with torch.no_grad(): 
            for val_data in validate_loader:
                val_images, val_labels = val_data
                outputs = net(val_images.to(device)) 
                predict_y = torch.max(outputs, dim=1)[1] 
                acc_num += torch.eq(predict_y, val_labels.to(device)).sum().item() 

        val_accurate = acc_num / val_num
        print('[epoch %d] train_loss: %.3f  train_acc: %.3f  val_accuracy: %.3f' %  (epoch + 1, loss, train_acc, val_accurate))   
        # 判断当前验证集的准确率是否是最大的，如果是，则更新之前保存的权重
        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(net.state_dict(), os.path.join(save_path, "AlexNet.pth") ) # state_dict()返回一个包含了模型所有参数（权重和偏置）的字典。这个字典中的键是参数的名称，而对应的值则是该参数的张量。

        # 每次迭代后清空这些指标，重新计算 
        train_acc = 0.0
        val_accurate = 0.0

    print('Finished Training')

 
# if __name__ == '__main__':
#     main()
main()