图像分类网络-train.py

        

        首先在进行训练之前最重要的是对数据进行处理，包括：1.数据预处理（主要是图像数据张量化及正则化处理）；2.数据导入（主要是ImageFolder及DataLoader）。

        其次是定义损失函数、优化器及学习率。

        最后是加载数据进行模型训练，记录数据、画图。

        代码部分参考b站视频

目录

1.数据预处理——torchvision.transforms

2.数据导入—torchvision.datasets.ImageFolder、torch.utils.Data.DataLoader

3.定义损失函数、优化器及学习率

4.训练函数train()

5.验证函数val()

6.正式开始训练并画出训练、验证集损失和准确率

7.附完整代码

---

1.数据预处理——torchvision.transforms

        数据预处理主要包括图像大小调整Resize()、图像竖直翻转RandomVertivalFlip()、将numpy格式的图像数据转换为torch的FloatTensor格式ToTensor()、平均值和标准差来正则化一个tensor图像使其映射到[-1,1]Normalize()。

        其中transforms.ToTensor()映射到[0,1]。       

        其中 transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])中的[0.5,0.5,0.5]分别表示平均值和标准差，对输出张量的三个通道进行先减后除操作，[a，b]'->[（a-0.5）/0.5，（b-0.5）/0.5]从而使其映射到[-1,1]中。

normalize=transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])

train_transform=transforms.Compose([
    transforms.Resize((224,224)),
    transforms.RandomVerticalFlip(),
    transforms.ToTensor(),
    normalize
])

val_transform=transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    normalize
])

2.数据导入—torchvision.datasets.ImageFolder、torch.utils.Data.DataLoader

        其中ImageFolder是一个通用的数据加载器，将以上预处理的方式来组织数据集的训练、验证或者测试图片。训练集图像根路径为./train,不要具体到./train/dog或者./train/cat路径，否则会报RuntimeError: Found 0 files in subfolders of: ./demoImg/dog的错误。ImageFolder将train路径下的两个文件夹cat和dog作为两个类别，train_dataset.classes打印输出['cat','dog']，Demo_DataSet.imgs打印输出为[('./demoImg\\dog\\dog.0.jpg', 1), ('./demoImg\\dog\\dog.1.jpg', 1), ...]

        ——train

                ——cat

                ——dog

        其中DataLoader在训练模型时使用到此函数，用来把训练数据分成多个小组，此函数每次抛出一组数据即batch_size个数据。

ROOT_TRAIN=r"C:\Users\ice\PycharmProjects\demo\data\train"
ROOT_TEST=r"C:\Users\ice\PycharmProjects\demo\data\val"

train_dataset=ImageFolder(ROOT_TRAIN,transform=train_transform)
val_dataset=ImageFolder(ROOT_TEST,transforms=val_transform)

train_dataloader=DataLoader(train_dataset,batch_size=8,shuffle=True)
val_dataloader=DataLoader(val_dataset,batch_size=8,shuffle=True)

3.定义损失函数、优化器及学习率

        torch.nn.CrossEntropyLoss()的输出就是log_softmax()的值作为nll_loss()的输入得到的，输入标签值和网络输出值，输出两者间的交叉熵损失。

        torch.optim.SGD()中的参数：params (iterable)  待优化参数的iterable（即parameters方法）或者是定义了参数组的dict，在优化器中常采用model.parameters()方法，在保存网络模型时候常使用model.state_dict()方法，会同时保存每一层操作和数值。

        torch.optim.lr_scheduler.StepLR()是用于调整学习率，每step_size步之后衰减gamma。

        优化器和学习率都是通过.step（）进行启动的。

#定义一个损失函数
loss_fn=nn.CrossEntropyLoss()
#定义一个优化器
optimizer=torch.optim.SGD(model.parameters(),lr=0.01,momentum=0.9)
#学习率每隔10轮变为原来的0.5
lr_scheduler=lr_scheduler.StepLR(optimizer,step_size=10,gamma=0.5)

4.训练函数train()

        dataloader返回的是image图像数据（batchsize,channel,H,W）及标签label,如下代码所示。       
        img.shape  -----------------  (batchsize, channel, H, W)

        label.shape ----------------- batch

        label           ----------------- 一个batch图片对应的label 

        其次torch.max(output,1)返回的是output张量中每行的最大值，及其对应的索引号。

        使用.item（）会取高精度进行返回，一般用于计算loss和准确率。

def train(dataloader,model,loss_fn,optimizer):
    loss,current,n=0.0,0.0,0.0
    for batch,(img,label) in enumerate(dataloader):
        image,y=img.to(device),label.to(device)
        output=model(image)
        curr_loss=loss_fn(output,label)
        #返回每行中最大值，及其对应的索引号
        _,pred=torch.max(output,axis=1)
        #output.shape[0]返回行数，即取每行的最大值（属于哪个类别的可能性最大），共多少个结果
        cur_acc=torch.sum(label==pred)/output.shape[0]

        #反向传播
        optimizer.zero_grad()
        curr_loss.backward()
        optimizer.step()
        loss+=curr_loss.item()
        current+=cur_acc.item()
        n=n+1

    train_loss=loss/n
    train_acc=current/n
    print("train_loss"+str(train_loss))
    print("train_acc"+str(train_acc))
    return train_loss,train_acc

5.验证函数val()

        验证函数和训练函数最大的区别在于，验证函数不需要反向传播，故需要将模型转换为验证模式model.eval(). 使用with torch.no_grad()的output没有grad_fn=<AddmmBackward>属性，不需要进行反向回传。

def val(dataloader, model, loss_fn):
    #将模型转换为验证模式
    model.eval()
    loss, current, n = 0.0, 0.0, 0.0
    with torch.no_grad():
        for batch, (img, label) in enumerate(dataloader):
            image, y = img.to(device), label.to(device)
            output = model(image)
            curr_loss = loss_fn(output, label)
            _, pred = torch.max(output, axis=1)
            cur_acc = torch.sum(label == pred) / output.shape[0]

            loss += curr_loss.item()
            current += cur_acc.item()
            n = n + 1

    val_loss = loss / n
    val_acc = current / n
    print("val_loss" + str(val_loss))
    print("val_acc" + str(val_acc))
    return val_loss, val_acc

6.正式开始训练并画出训练、验证集损失和准确率

        torch.save(model.state_dict(),dir) 保存的是模型的weights及偏置，同样也可以保存优化器及epoch。

        关于保存model.state_dict()和model的区别:可以参考该篇博客的内容，https://blog.csdn.net/strive_for_future/article/details/83240081

if __name__ == '__main__':
    loss_train=[]
    acc_train=[]

    loss_val=[]
    acc_val=[]

    epoch=20
    min_acc=0

    for t in range(epoch):
        lr_scheduler.step()
        print(f"epoch{t+1}\n----------")
        train_loss,train_acc=train(train_dataloader,model,loss_fn,optimizer)
        val_loss,val_acc=val(val_dataloader,model,loss_fn)

        loss_train.append(train_loss)
        acc_train.append(train_acc)
    
        loss_val.append(val_loss)
        acc_val.append(val_acc)

        #保存最好的模型权重
        if val_acc>min_acc:
            folder="sava_model"
            if not os.path.exists(folder):
                os.makedirs(folder)
            min_acc=val_acc
            print(f"save best model,第{t+1}轮")
            torch.save(model.state_dict(),'sava_model/best_model.pth')
        #保存最后一轮得到权重文件
        if t==epoch-1:
            torch.save(model.state_dict(), 'sava_model/last_model.pth')

7.附完整代码

from  net import myNet
import torch
from torch import nn
import numpy as np

from torch.optim import lr_scheduler
import os

from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

import matplotlib.pyplot as plt

#解决中文显示问题
plt.rcParams['font.sans-serif']=['SimHei']
plt.rcParams['axes.unicode_minus']=False

ROOT_TRAIN=r"C:\data\train"
ROOT_TEST=r"C:\data\val"


#定义一个画图函数
def matplot_loss(train_loss,val_loss):
    plt.plot(train_loss,label="train_loss")
    plt.plot(val_loss,label="val_loss")
    #图例显示为best
    plt.legend(loc="best")
    plt.ylabel("loss")
    plt.xlabel("epoch")
    plt.title("训练集和验证集loss值对比图")
    plt.show()

def matplot_acc(train_acc,val_acc):
    plt.plot(train_acc,label="train_acc")
    plt.plot(val_acc,label="val_acc")
    #图例显示为best
    plt.legend(loc="best")
    plt.ylabel("acc")
    plt.xlabel("epoch")
    plt.title("训练集和验证集acc值对比图")
    plt.show()


##1.数据预处理
#将图像的像素值归一化[-1,1]之间
normalize=transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])

train_transform=transforms.Compose([
    transforms.Resize((224,224)),
    transforms.RandomVerticalFlip(),
    transforms.ToTensor(),
    normalize
])

val_transform=transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    normalize
])

##2.导入数据
train_dataset=ImageFolder(ROOT_TRAIN,transform=train_transform)
val_dataset=ImageFolder(ROOT_TEST,transform=val_transform)


train_dataloader=DataLoader(train_dataset,batch_size=8,shuffle=True)
val_dataloader=DataLoader(val_dataset,batch_size=8,shuffle=True)

device='cuda' if torch.cuda.is_available() else 'cpu'

print(device)

model=myNet().to(device)
model_parameters=model.parameters()

#定义一个损失函数
loss_fn=nn.CrossEntropyLoss()
#定义一个优化器
optimizer=torch.optim.SGD(model.parameters(),lr=0.01,momentum=0.9)
#学习率每隔10轮变为原来的0.5
lr_scheduler=lr_scheduler.StepLR(optimizer,step_size=10,gamma=0.5)
#定义一个训练函数
def train(dataloader,model,loss_fn,optimizer):
    loss,current,n=0.0,0.0,0.0
    for batch,(img,label) in enumerate(dataloader):
        image,y=img.to(device),label.to(device)
        output=model(image)
        curr_loss=loss_fn(output,label)
        #返回每行中最大值，及其对应的索引号
        _,pred=torch.max(output,axis=1)
        #output.shape[0]返回行数，即取每行的最大值（属于哪个类别的可能性最大），共多少个结果
        cur_acc=torch.sum(label==pred)/output.shape[0]

        #反向传播
        optimizer.zero_grad()
        curr_loss.backward()
        optimizer.step()
        loss+=curr_loss.item()
        current+=cur_acc.item()
        n=n+1

    train_loss=loss/n
    train_acc=current/n
    print("train_loss"+str(train_loss))
    print("train_acc"+str(train_acc))
    return train_loss,train_acc

#定义一个验证函数
def val(dataloader, model, loss_fn):
    #将模型转换为验证模式
    model.eval()
    loss, current, n = 0.0, 0.0, 0.0
    with torch.no_grad():
        for batch, (img, label) in enumerate(dataloader):
            image, y = img.to(device), label.to(device)
            output = model(image)
            curr_loss = loss_fn(output, label)
            _, pred = torch.max(output, axis=1)
            cur_acc = torch.sum(label == pred) / output.shape[0]

            loss += curr_loss.item()
            current += cur_acc.item()
            n = n + 1

    val_loss = loss / n
    val_acc = current / n
    print("val_loss" + str(val_loss))
    print("val_acc" + str(val_acc))
    return val_loss, val_acc

#开始训练
if __name__ == '__main__':
    loss_train=[]
    acc_train=[]

    loss_val=[]
    acc_val=[]


    epoch=20
    min_acc=0

    for t in range(epoch):
        lr_scheduler.step()
        print(f"epoch{t+1}\n----------")
        train_loss,train_acc=train(train_dataloader,model,loss_fn,optimizer)
        val_loss,val_acc=val(val_dataloader,model,loss_fn)

        loss_train.append(train_loss)
        acc_train.append(train_acc)

        loss_val.append(val_loss)
        acc_val.append(val_acc)

        #保存最好的模型权重
        if val_acc>min_acc:
            folder="sava_model"
            if not os.path.exists(folder):
                os.makedirs(folder)
            min_acc=val_acc
            print(f"save best model,第{t+1}轮")
            torch.save(model.state_dict(),'sava_model/best_model.pth')
        #保存最后一轮得到权重文件
        if t==epoch-1:
            torch.save(model.state_dict(), 'sava_model/last_model.pth')