图像分类模板——有数据集即可实现图像分类任务

VaceMan

已于 2023-09-22 14:14:23 修改

阅读量600

点赞数 3

文章标签：人工智能 pytorch 深度学习 python

于 2023-09-22 14:08:54 首次发布

本文链接：https://blog.csdn.net/weixin_46709801/article/details/133168573

版权

2.将数据集划分为训练集和测试集，按照73分

三、图像分类万能模板（基于alexnet的预训练模型）

一.如何制作/获取数据集？数据集的格式？

1.自己制作数据集

按照自己的图像分类任务，依次按类收集图片。（百度平台/自行拍摄……）

2.从网上下载免费数据集

AI Studio平台

3.数据集的格式？

将平台上下载下来的数据集解压后，如下所示：

数据集（目标类1/目标类2/……目标类n）

二、划分数据集

数据集划分为训练集和测试集。

训练集用来训练模型，测试集用来评价模型。

划分前：数据集（目标类1/目标类2/……/目标类n）

划分后：（82分、73分）

（训练集：目标类1 70%/目标类2 70%/……/目标类n 70%）

（测试集：目标类1 30%/目标类2 30%/……/目标类n 30%）

1.选取需要的数据集（以三分类为例）

从11类中选取3类，我想训练一个识别柴犬、拉布拉多和中华田园犬的模型

这样，我把解压后的文件夹中的这3个文件夹拿到我待制作的数据集文件夹中，如下图所示：

柴犬文件夹照片有89张

2.将数据集划分为训练集和测试集，按照73分

创建train文件夹和test文件夹，将柴犬文件夹、拉布拉多文件夹、中华田园犬文件夹放到train文件夹下。如下图所示：

训练集文件下如下图所示：

将柴犬文件夹改为0，拉布拉多文件夹改为1，中华田园犬文件夹改为2

在1中，刚才我们注意到柴犬文件夹有89张，按照73分，训练集：测试集=62.3：26.7

训练集63张，测试集26张

将train文件夹中的柴犬文件夹的图片保留63张，剩余移到test文件夹下的柴犬文件夹下。如下图所示。

上面是0(柴犬)划分后的结果，1(拉布拉多)、2(中华田园犬)同理。

三、图像分类万能模板（基于alexnet的预训练模型）

下面代码只需要修改4处（代码中我已用中文标出），即可使用。

（1）"划分后的训练集文件夹绝对路径" 、 "划分后的测试集文件夹绝对路径"

（2）图像分类任务的分类数

（3）训练轮数

（4）训练后的模型保存路径

from random import shuffle
from matplotlib import pyplot as plt
import torch
import torchvision
import copy

#一、数据集
#1.1图片增强定义
from torchvision import transforms
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])
])
test_transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])
])

#1.2 datasets
train_dir = r"划分后的训练集文件夹绝对路径"
test_dir = r"划分后的测试集文件夹绝对路径"
train_ds = torchvision.datasets.ImageFolder(
    train_dir,
    transform=train_transform
)
test_ds = torchvision.datasets.ImageFolder(
    test_dir,
    transform=test_transform
)
#1.3 dataloader
BATCHSIZE = 32
train_dl = torch.utils.data.DataLoader(
    train_ds,
    batch_size=BATCHSIZE,
    shuffle = True
)
test_dl = torch.utils.data.DataLoader(
    test_ds,
    batch_size=BATCHSIZE
)
#1.4 输出数据集信息
print("训练集数量:",len(train_ds))
print("测试集数量:",len(test_ds))

#二、创建模型
#2.1 加载alexnet预训练模型
from torch import nn
model = torchvision.models.alexnet(pretrained=True)
for p in model.features.parameters():#卷积层梯度更新冻结
    p.requires_grad = False
model.classifier[-1].out_features = 图像分类任务的分类数
print(model.classifier[-1].out_features)
#2.2查看模型信息
print(model)

#三、模型训练与测试
#3.1 train函数
def train(dl,model,loss_fn,optimizer):
    size = len(dl.dataset)  #数据集总个数
    num_batches = len(dl)  #总批次数

    train_loss, correct = 0,0
    model.train()
    for x,y in dl:  #x,y代表每一批次
        x,y = x.to(device),y.to(device)
        pred = model(x)
        loss = loss_fn(pred,y)  #一个批次上的损失
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        with torch.no_grad():
            #上下文管理器，不会计算梯度
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()  #累加每一个batch下的正确个数
            train_loss += loss.item()  #累加每一个batch的loss
    correct /= size
    train_loss /= num_batches
    return correct,train_loss
#3.2 test函数
def test(test_dl,model,loss_fn):
    size = len(test_dl.dataset)  #数据集总个数
    num_batches = len(test_dl)  #总批次数

    test_loss, correct = 0,0
    model.eval()
    with torch.no_grad():
        for x,y in test_dl:  #x,y代表每一批次
            x,y = x.to(device),y.to(device)
            pred = model(x)
            loss = loss_fn(pred,y)  #一个批次上的损失
            test_loss += loss.item()  #累加每一个batch的loss
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()  #累加每一个batch下的正确个数
        correct /= size
        test_loss /= num_batches
        return correct,test_loss
#3.3 一些参数设置
if torch.cuda.is_available():
    model.to('cuda')
loss_fn = nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.classifier.parameters(),lr=0.0001)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)
#训练轮数
epochs = 50
train_loss = []
train_acc = []
test_loss = []
test_acc = []
#训练后的模型保存路径，以,pth结尾
PATH = "./alexnet_pre_first.pth"
best_acc = 0
#3.4 训练和测试循环
#4.开始训练
for epoch in range(epochs):
    epoch_acc,epoch_loss = train(train_dl,model,loss_fn,optim)
    epoch_test_acc,epoch_test_loss = test(test_dl,model,loss_fn)
    if epoch_test_acc > best_acc:
        best_model_wts = copy.deepcopy(model.state_dict())
        best_acc = epoch_test_acc
        torch.save(model, PATH)
#         torch.save(model.state_dict(),PATH)
    train_acc.append(epoch_acc)
    train_loss.append(epoch_loss)
    test_acc.append(epoch_test_acc)
    test_loss.append(epoch_test_loss)
    template = ("epoch:{:2d},train_loss:{:.5f},train_acc:{:.1f},test_loss:{:.5f},test_acc:{:.1f}")
    print(template.format(epoch+1,epoch_loss,epoch_acc*100,epoch_test_loss,epoch_test_acc*100))
print("Done")
#3.5训练结果打印输出
print("train_acc",train_acc)
print("train_loss",train_loss)
print("test_acc",test_acc)
print("test_loss",test_loss)

#这部分是将训练结果绘制图形
plt.plot(range(训练轮数),train_loss,label='train_loss')
plt.plot(range(训练轮数),test_loss,label='test_loss')
plt.legend()
plt.show()

plt.plot(range(训练轮数),train_acc,label='train_acc')
plt.plot(range(训练轮数),test_acc,label='test_acc')
plt.legend()
plt.show()

from PIL import Image
#这部分是使用训练出来的模型去预测其它图片
path = r"要预测的图片路径"
#识别种类
classes = ['分类目标0','分类目标1','分类目标n']
#图片处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
])
#加载预处理模型
model = torch.load(r"训练好的模型文件路径",map_location='cpu')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)
#对图片进行识别
img = Image.open(path)
img = transform(img)
img = torch.unsqueeze(img,dim=0)  #增加batch维度
model.eval()
with torch.no_grad():
    img = img.to(device)
    pred = model(img)
    pre_gailv = torch.softmax(pred,dim=1)
    pre_label = torch.argmax(pre_gailv).cpu().numpy().item()
    print(classes[pre_label])