基于迁移学习的手势分类模型训练

news2024/9/23 11:27:46

1、基本原理介绍

这里介绍的单指模型迁移。一般我们训练模型时，往往会自定义一个模型类，这个类中定义了神经网络的结构，训练时将数据集输入，从0开始训练；而迁移学习中（单指模型迁移策略），是在一个已经有过训练基础的模型上，用自己的数据集，进一步训练，使得这个模型能够完成我们需要的任务。

这么做有有这样几个显而易见的好处：

※ 因为模型之前被训练过，所以初始参数不会是0，这样能够加速模型训练

※ 因为预训练模型（什么是预训练模型下文会讲到）在其他数据集上训练过，而其他数据集往往和我们用的数据集存在一定的区别，所以这可以提高模型的泛化能力

※ 通过迁移学习，可以将来自大规模数据的优势转移到小规模或新任务上，提高模型的表现和效果

2、预训练模型

在进行迁移学习时，我们要先找到一个预训练模型。在分类任务领域，比较流行的如resnet系列、mobilenet系列（更轻量化）、vgg（系列）、efficientnet（系列）等等网络，都是比较常用且容易获得的预训练模型，这些模型都能够通过python直接下载。

而且由于上述模型基本都是在ImageNet这一大规模，多分类类别的数据集上进行过训练的，所以对于简单的二分类等少数类别分类，能有较好的效果。

3、训练流程

迁移学习完整的训练流程和一般搭建神经网络的训练模型的流程基本类似：数据预处理->数据集的切分->加载预训练模型（搭建神经网络）->设置超参数/损失函数/优化器等->训练模型

3.1 模型训练

下面的代码是一个利用mobilenet网络训练得到的手势分类模型，该模型能够较准确的分类不同类别手势。

相关解释已在代码中注释说明。

from torchvision.models import mobilenet_v2
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, RandomHorizontalFlip, RandomRotation

# 定义数据预处理和增强器
transform = Compose([
    RandomHorizontalFlip(),  # 随机水平翻转
    RandomRotation(10),      # 随机旋转10度
    Resize((224, 224)),
    CenterCrop(224),
    ToTensor(),
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集并应用预处理和增强器
dataset = ImageFolder(root='data', transform=transform)
# 这里由于数据比较少，将所有数据集全部用来训练，得到的模型直接拿来用了，这其实不算是非常规范的操作，仅供参考


# 定义网络结构
model = mobilenet_v2(pretrained=True)  # 加载预训练模型，也可以试试其他模型，效果差别挺大的
num_ftrs = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_ftrs, 5)  # 假设是5分类问题，具体几分类，改这里的参数就行了

# 将模型移动到设备上
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 定义优化器和损失函数
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-5)
criterion = nn.CrossEntropyLoss()

# 定义训练循环
def train_model(model, criterion, optimizer, num_epochs, train_loader):
    for epoch in range(num_epochs):
        model.train()  # 设置模型为训练模式
        train_loss = 0.0
        correct = 0
        total = 0

        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        epoch_loss = train_loss / total
        epoch_acc = 100. * correct / total

        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%')


# 创建训练集的DataLoader
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

# 开始训练模型
train_model(model, criterion, optimizer, num_epochs=15, train_loader=train_loader)
torch.save(model, 'my_model(1).pth')

3.2 数据集文件结构

当然，你也可以自己定义读取数据集的data_loader类。

3.3 模型推理

这段代码是用训练得到的模型对一张图片进行推理测试的，如果需要对系列图片进行推理，评估模型效果，可自行修改，调用对应函数即可。

import torch
from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
def predict_image(image_path, model_path='my_model(1).pth'):

    image = Image.open(image_path).convert("RGB")
    # 对测试的图片进行预处理，需要和训练时处理的方式一样
    transform = Compose([
        Resize((224, 224)),
        CenterCrop(224),
        ToTensor(),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    image_tensor = transform(image).unsqueeze(0)
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    image_tensor = image_tensor.to(device)


    model = torch.load(model_path,map_location=device)
    model.eval()
    with torch.no_grad():
        output = model(image_tensor)
        _, predicted = torch.max(output.data, 1)  # 获得分类标记
    return predicted.item()
if __name__=="__main__":
    image_path = "test2/6.jpg"
    print(predict_image(image_path))