基于pytorch的图像训练识别

一、准备数据集

分为测试集和训练集，文件如下排放

二、开始识别

数据集准备好后，即可导入到模型开始训练，运行下列代码

import time
from torch.utils.tensorboard import SummaryWriter
from torchvision.datasets import ImageFolder
from torchvision import transforms
from torch.utils.data import DataLoader
import torchvision.models as models
import torch.nn as nn
import torch

print("是否使用GPU训练：{}".format(torch.cuda.is_available()))    #打印是否采用gpu训练
if torch.cuda.is_available:
    print("GPU名称为：{}".format(torch.cuda.get_device_name()))  #打印相应的gpu信息
#数据增强太多也可能造成训练出不好的结果，而且耗时长，宜增强两三倍即可。
normalize=transforms.Normalize(mean=[.5,.5,.5],std=[.5,.5,.5])  #规范化
transform=transforms.Compose([                                  #数据处理
    transforms.Resize((64,64)),
    transforms.ToTensor(),
    normalize
])
dataset_train=ImageFolder('modelphotos/train',transform=transform)     #训练数据集
# print(dataset_tran[0])
dataset_valid=ImageFolder('modelphotos/valid',transform=transform)     #验证或测试数据集
# print(dataset_train.classer)#返回类别
print(dataset_train.class_to_idx)                               #返回类别及其索引
# print(dataset_train.imgs)#返回图片路径
print(dataset_valid.class_to_idx)
train_data_size=len(dataset_train)                              #放回数据集长度
test_data_size=len(dataset_valid)
print("训练数据集的长度为：{}".format(train_data_size))
print("测试数据集的长度为：{}".format(test_data_size))
#torch自带的标准数据集加载函数
dataloader_train=DataLoader(dataset_train,batch_size=4,shuffle=True,num_workers=0,drop_last=True)
dataloader_test=DataLoader(dataset_valid,batch_size=4,shuffle=True,num_workers=0,drop_last=True)

#2.模型加载
model_ft=models.resnet18(pretrained=True)#使用迁移学习，加载预训练权重
# print(model_ft)

in_features=model_ft.fc.in_features
model_ft.fc=nn.Sequential(nn.Linear(in_features,36),
                          nn.Linear(36,6))#将最后的全连接改为（36，6），使输出为六个小数，对应六种植物的置信度
#冻结卷积层函数
# for i,para in enumerate(model_ft.parameters()):
#     if i<18:
#         para.requires_grad=False

# print(model_ft)

# model_ft.half()#可改为半精度，加快训练速度，在这里不适用

model_ft=model_ft.cuda()#将模型迁移到gpu
#3.优化器
loss_fn=nn.CrossEntropyLoss()

loss_fn=loss_fn.cuda()  #将loss迁移到gpu
learn_rate=0.01         #设置学习率
optimizer=torch.optim.SGD(model_ft.parameters(),lr=learn_rate,momentum=0.01)#可调超参数

total_train_step=0
total_test_step=0
epoch=30                #迭代次数
writer=SummaryWriter("logs_train_yaopian")
best_acc=-1
ss_time=time.time()

for i in range(epoch):
    start_time = time.time()
    print("--------第{}轮训练开始---------".format(i+1))
    model_ft.train()
    for data in dataloader_train:
        imgs,targets=data
        # if torch.cuda.is_available():
        # imgs.float()
        # imgs=imgs.float()#为上述改为半精度操作，在这里不适用
        imgs=imgs.cuda()
        targets=targets.cuda()
        # imgs=imgs.half()
        outputs=model_ft(imgs)
        loss=loss_fn(outputs,targets)

        optimizer.zero_grad()   #梯度归零
        loss.backward()         #反向传播计算梯度
        optimizer.step()        #梯度优化

        total_train_step=total_train_step+1
        if total_train_step%100==0:#一轮时间过长可以考虑加一个
            end_time=time.time()
            print("使用GPU训练100次的时间为：{}".format(end_time-start_time))
            print("训练次数：{},loss:{}".format(total_train_step,loss.item()))
            # writer.add_scalar("valid_loss",loss.item(),total_train_step)
    model_ft.eval()
    total_test_loss=0
    total_accuracy=0
    with torch.no_grad():       #验证数据集时禁止反向传播优化权重
        for data in dataloader_test:
            imgs,targets=data
            # if torch.cuda.is_available():
            # imgs.float()
            # imgs=imgs.float()
            imgs = imgs.cuda()
            targets = targets.cuda()
            # imgs=imgs.half()
            outputs=model_ft(imgs)
            loss=loss_fn(outputs,targets)
            total_test_loss=total_test_loss+loss.item()
            accuracy=(outputs.argmax(1)==targets).sum()
            total_accuracy=total_accuracy+accuracy
        print("整体测试集上的loss：{}(越小越好,与上面的loss无关此为测试集的总loss)".format(total_test_loss))
        print("整体测试集上的正确率：{}(越大越好)".format(total_accuracy / len(dataset_valid)))

        writer.add_scalar("valid_loss",(total_accuracy/len(dataset_valid)),(i+1))#选择性使用哪一个
        total_test_step = total_test_step + 1
        if total_accuracy > best_acc:   #保存迭代次数中最好的模型
            print("已修改模型")
            best_acc = total_accuracy
            torch.save(model_ft, "best_model_yaopian.pth")
ee_time=time.time()
zong_time=ee_time-ss_time
print("训练总共用时:{}h:{}m:{}s".format(int(zong_time//3600),int((zong_time%3600)//60),int(zong_time%60))) #打印训练总耗时
writer.close()

上述采用的迁移学习直接使用resnet18的模型进行训练，只对全连接的输出进行修改，是一种十分方便且实用的方法，同样，你也可以自己编写模型，然后使用自己的模型进行训练，但是这种方法显然需要训练更长的时间才能达到拟合。如图所示，只需要修改矩形框内部分，将‘model_ft=models.resnet18(pretrained=True)'改为自己的模型‘model_ft=model’即可。

三、模型应用测试

经过上述的步骤后，我们将会得到一个‘best_model_yaopian.pth’的模型权重文件，最后运行下列代码就可以对图片进行识别了

import os
import torch
import torchvision
from PIL import Image
from torch import nn
i=0 #识别图片计数
root_path="测试_data"         #待测试文件夹
names=os.listdir(root_path)
for name in names:
    print(name)
    i=i+1
    data_class=['cat','dog','duck','mouse']   #按文件索引顺序排列
    image_path=os.path.join(root_path,name)
    image=Image.open(image_path)
    print(image)
    transforms=torchvision.transforms.Compose([torchvision.transforms.Resize((64,64)),
                                              torchvision.transforms.ToTensor()])
    image=transforms(image)
    print(image.shape)

    model_ft=torchvision.models.resnet18()      #需要使用训练时的相同模型
    # print(model_ft)
    in_features=model_ft.fc.in_features
    model_ft.fc=nn.Sequential(nn.Linear(in_features,36),
                              nn.Linear(36,6))     #此处也要与训练模型一致

    model=torch.load("best_model_yaopian.pth",map_location=torch.device("cpu")) #选择训练后得到的模型文件
    # print(model)
    image=torch.reshape(image,(1,3,64,64))      #修改待预测图片尺寸，需要与训练时一致
    model.eval()
    with torch.no_grad():
        output=model(image)
    print(output)               #输出预测结果
    # print(int(output.argmax(1)))
    print("第{}张图片预测为：{}".format(i,data_class[int(output.argmax(1))]))   #对结果进行处理，使直接显示出预测的植物种类

 这里需要注意一下索引顺序，可以通过训练模型时的输出来查看并填入

可以更换标签来改变输出结果

最后，通过上述步骤我们可以得到一个简单的猫狗鸭鼠的智能识别程序，如下图是识别结果说明。