图像分类-CIFAR100项目实战

train.py

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

from torch import nn
from torch.nn import Sequential,Conv2d,MaxPool2d,Flatten,Linear,BatchNorm2d,Softmax
from torch.utils.tensorboard import SummaryWriter
from torch import cuda
import os
import sys
from model.VR import *
from utility.S_C_E import *
from utility.step_lr import *
from utility.data_enhence import *
from utility.lion import Lion  #优化函数
from utility.resume import *


#记录
writer = SummaryWriter("logs")
#添加数据集

transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),  # 图像一半的概率翻转，一半的概率不翻转
    # transforms.RandomVerticalFlip(p=0.5),
    # transforms.RandomRotation(90),
    transforms.Resize([256,256]),
    transforms.ToTensor(),  # 维度转化 由32x32x3  ->3x32x32
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    # transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])

    # R,G,B每层的归一化用到的均值和方差     即参数为变换过程，而非最终结果。
])

transform_test = transforms.Compose([
    transforms.Resize([256,256]),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    # transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])
])


train_data=torchvision.datasets.CIFAR100(root=r'D:\desk\dataset',train=True,transform=transform_train,download=True)
test_data=torchvision.datasets.CIFAR100(root=r'D:\desk\dataset',train=False,transform=transform_test,download=True)


# #打印训练集与测试集的数量
print("训练集的数量为：{}".format(len(train_data)))
print("测试集的数量为:{}".format(len(test_data)))


Batch_size=32
#加载数据集
train_dataloader=DataLoader(train_data,batch_size=Batch_size,shuffle=True,drop_last=True)
test_dataloader=DataLoader(test_data,batch_size=Batch_size,shuffle=True,drop_last=True)


#实例化网络模型
modelA = ViT(
    image_size=256,
    patch_size=32,
    num_classes=10,
    dim=1024,
    depth=6,
    heads=16,
    mlp_dim=2048,
    dropout=0.1,
    emb_dropout=0.1
)
modelB = ResNet34()

#加载预训练模型
T= MyEnsemble(modelA, modelB)

if torch.cuda.is_available():
    T=T.cuda()


#损失函数
loss_fn=nn.CrossEntropyLoss() #分类问题，用交叉熵
if torch.cuda.is_available():
    loss_fn=loss_fn.cuda()



#优化器
learning_rate=1e-4
optim=torch.optim.SGD(T.parameters(),lr=learning_rate,momentum=0.9,weight_decay=5e-4) #params表示对哪一部分进行优化，此处选择网络模型
# optim=Lion(T.parameters(), lr = 1e-4, weight_decay = 1e-2)



#训练时的参数
#训练的次数
total_train_step=0

#测试的次数
total_test_step=0

#总的训练轮数
epochs=300

#最佳模型精度
best_acc = 0.0

#动态学习率
# scheduler = StepLR(optim, learning_rate, epochs)

#模型保存位置
models_save_path=r'D:\desk\project_XHU\models\{}'.format(3)
if not os.path.exists(models_save_path):
    os.makedirs(models_save_path)

for epoch in range(epochs):
    print("第{}轮训练开始".format(epoch+1))
    T.train()#开始训练只对一些特定的模块有用，如dropout、batchnormal有用，有它们一定要用，把模型设置成训练模式。

    for data in train_dataloader:
        imgs,targets=data

        if torch.cuda.is_available():
            imgs=imgs.cuda()
            targets=targets.cuda()

        output=T(imgs)
        train_loss=loss_fn(output,targets.long())#loss_fn的返回的结果为一个损失值（实际与预测的差值）



        #优化器优化模型
        optim.zero_grad()
        train_loss.backward()
        optim.step()

        #记录训练的次数
        total_train_step = total_train_step + 1
        if total_train_step % 100==0:
            print('训练次数：{}，损失为{}'.format(total_train_step,train_loss.item()))


        writer.add_scalar("train_loss", train_loss.item(), global_step=total_train_step)

        # 更新学习率
        # with torch.no_grad():
        #     scheduler(epoch)


    #开始测试，测试步骤，no_grad即没有梯度，无需更新参数，仅仅用来判断模型的训练好坏
    T.eval() #只对一些特定的模块有用，如dropout、batchnormal有用,把模型设置成验证模式。
    total_accuracy=0
    total_loss=0
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets=data

            if torch.cuda.is_available():
                imgs=imgs.cuda()
                targets=targets.cuda()

            output=T(imgs)
            test_loss=loss_fn(output,targets.long())
            total_loss=total_loss+test_loss



            accuracy=(output.argmax(1)==targets).sum()
            total_accuracy=accuracy+total_accuracy

            total_test_step=total_test_step+1
            if total_test_step % 100 ==0:
                print('测试次数为：{}，整体损失为：{}'.format(total_test_step,total_loss.item()))
            writer.add_scalar('test_loss', total_loss, global_step=total_test_step)

        acc=total_accuracy / len(test_data)

        print("total_accuracy / len(test_ds={}/{}".format(total_accuracy,len(test_data)))
        print('测试集精度为：{}'.format(acc))#数据集精确度
    writer.add_scalar('test_acc',acc,global_step=total_test_step)

    if best_acc<acc:
        best_models_epoch=epoch
        print("上次最佳精度为{}，本轮精度为{}".format(best_acc,acc))
        best_acc=acc
        torch.save(T, models_save_path+'/Model({})acc({}).pth'.format(epoch,best_acc))
        print("模型已保存")
    if epoch - best_models_epoch>30:
        print("模型训练已经有{}轮没有提升，因此提前结束训练".format(epoch-best_models_epoch))
        sys.exit(0)

writer.close()

test.py

import torchvision.transforms
from PIL import Image
from torchvision import transforms
import torch
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear
import cv2
from matplotlib import pyplot as plt
import os
import random

Cls=["Cr","In","Pa","PS","RS","Sc"]


res=torchvision.models.resnet101(pretrained=True,progress=True)
res.fc.add_module('linelayer',Linear(in_features=1000,out_features=6))
if torch.cuda.is_available():
    T=res.cuda()

path_dir=r'D:\desk\project_steel\dataset\process_imgs\RS'


a=random.randint(1,99999)
path_save=r"D:\desk\project_steel\test_result\{}".format(a)
print(path_save)

if not os.path.exists(path_save):
    os.mkdir(path_save)



# img=cv2.imread(os.path.join(path_dir,"car (2).jpg"))
# img1=cv2.imread(os.path.join(path_dir,"car (1).jpg"))
# print(img1.shape)



for name in os.listdir(path_dir):

    path=os.path.join(path_dir,name)

    img=plt.imread(path)
    img=Image.open(path)
    img = img.convert('RGB')

    transform=torchvision.transforms.Compose([torchvision.transforms.Resize([200,200]),torchvision.transforms.ToTensor()])
    img=transform(img)


    img=img.cuda()#指定图片是cuda加速了的
    model=torch.load(r'D:\desk\project_steel\models\10\Model(10)acc(1.0).pth',map_location=torch.device('cuda')) #指定模型在哪种条件下训练的，以便加载
    # print(model)
    img=torch.reshape(img,[1,3,200,200])
    # print(img.shape)


    model.eval()
    with torch.no_grad():#此行代码可以减少训练参数，加快训练速度。
        output=model(img)
    # print(output)


    img=cv2.imread(path)
    H,W=img.shape[:2]
    x=int(W/2)
    y=int(H/2)
    img=cv2.putText(img,Cls[output.argmax(1)],(50,y),5,5,(0,0,255),3)

    save_name_pic=os.path.join(path_save,name)
    cv2.imwrite(save_name_pic, img)
    print("save")
    # cv2.imshow("1",img)
    # cv2.waitKey(0)
    print(Cls[output.argmax(1)])







# img = plt.imread(path_data)
# img = Image.open(path_data)
# img = img.convert('RGB')
#
# transform = torchvision.transforms.Compose([torchvision.transforms.Resize([200,200]), torchvision.transforms.ToTensor()])
# img = transform(img)
#
# img = img.cuda()  # 指定图片是cuda加速了的
# model = torch.load(r'D:\desk\project_steel\models\3\Model(0)acc(0.6777777671813965).pth',
#                    map_location=torch.device('cuda'))  # 指定模型在哪种条件下训练的，以便加载
# print(model)
#
#
# # pil_img = Image.open(img)  # pip install pillow
# # img = img.convert('RGB')
#
#
# img = torch.reshape(img, [1, 3, 200,200])
#
# # pil_img = Image.open(img)
# # pil_img = pil_img.convert('RGB')
#
# # print("img.shape:",img.shape)
#
# model.eval()
# with torch.no_grad():  # 此行代码可以减少训练参数，加快训练速度。
#     output = model(img)
# print(output.argmax(1))
#
# img = cv2.imread(path_data)
# H, W = img.shape[:2]
# x = int(W / 2)
# y = int(H / 2)
# img = cv2.putText(img, Cls[output.argmax(1)], (50, y), 5, 5, (0, 0, 255), 3)
#
# save_name_pic = os.path.join(path_save, name[-1])
# print(save_name_pic)
#
# cv2.imwrite(save_name_pic, img)
# print("save")
# cv2.imshow("1",img)
# cv2.waitKey(0)
# print(Cls[output.argmax(1)])