如何训练一个图像分类的模型

2023/12/6

以下给出文字描述以及简略代码，具体实现请根据实际情况进行更改！

1.准备数据集:收集数据集，尽可能包含样本多样，使训练的模型的鲁棒性更好。对数据集进行划分，划分为测试集，训练集，验证集

2.对数据集进行预处理：包括更改尺寸、水平翻转、灰度化、归一化等操作，以防止模型出现拟合的现象。

3.选择模型：可以是vgg、resnet、iception、efficientnet等，根据实际情况对模型进行调整，如人脸真假二分类问题，需要将更改输出的类别数目为2。添加正则化、归一化，防止模型出现过拟合的问题。

4.损失函数和优化器的选择：对于真假二分类的问题损失函数一般是交叉熵损失函数，优化器是Adam优化器。

5.对模型进行训练

6.模型评估：通过准确率、损失率以及混淆矩阵对模型进行评估，再根据结果对模型进行调整。

以下是简单的代码描述，根据实际问题进行修改：
 

1.导入必要的包

import os
os.environ['TORCH_USE_CUDA_DSA'] = '1'
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder
from efficientnet_pytorch import EfficientNet
import torch.nn as nn
import torchvision.models as models
import torch.optim as optim

2.# 自定义数据集类
class MyDataset(Dataset):
    def __init__(self, root, transform=None):
        self.dataset = ImageFolder(root, transform=transform)#transformer对数据进行操纵

 # 获取ImageFolder对象的classes属性，其class属性为train、validation、test
         self.classes = self.dataset.classes 
    def __len__(self):
        return len(self.dataset)
    def __getitem__(self, idx):
        img, label = self.dataset[idx]
        return img, label
# 数据预处理和增强
transform = transforms.Compose([
    transforms.Resize((224, 224)),#更改图像的尺寸大小
    transforms.ToTensor(),#将图像转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])#对图像进行归一化

# 创建自定义数据集实例
dataset = MyDataset(root='这里填写相对项目的相对路径', transform=transform)
# 输出类别信息
print("Classes:", dataset.classes)

# 创建数据加载器
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 加载预训练的 EfficientNet 模型
#model = models.efficientnet_b4(pretrained=True)
model = EfficientNet.from_pretrained('efficientnet-b4')

# 查看模型的结构
print(model)
# 获取全连接层的输入
feature = model._fc.in_features
model._fc = nn.Linear(in_features=feature,out_features=2,bias=True)#更改模型的输出

#查看更改后的模型的参数是否符合自己的分类要求
print(model)

3.# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

4.# 训练模型
num_epochs = 10#根据实际情况进行更改
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)#将模型加载到GPU上

for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(dataloader):
        inputs, labels = inputs.to(device), labels.to(device)  # Move data to GPU

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()#反向传递误差
        optimizer.step()

        running_loss += loss.item()
        if i % 10 == 9:  # Print every 10 mini-batches
            print(f'Epoch [{epoch + 1}/{num_epochs}], '
                  f'Step [{i + 1}/{len(dataloader)}], '
                  f'Loss: {running_loss / 10:.4f}')
            running_loss = 0.0

    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {loss.item()}')

# Save the trained model保存训练好的模型参数
torch.save(model.state_dict(), 'my_model.pth')

----------对训练的优化，可以对上面的训练模型的部分进行添加

# 初始化存储训练过程中的损失和准确率
train_losses = []
train_accuracies = []

for epoch in range(num_epochs):
    running_loss = 0.0
    correct_predictions = 0
    total_samples = 0

    for i, (inputs, labels) in enumerate(dataloader):
        inputs, labels = inputs.to(device), labels.to(device)

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

        unning_loss += loss.item()

        # 计算准确率
        _, predicted = torch.max(outputs.data, 1)
        total_samples += labels.size(0)
        correct_predictions += (predicted == labels).sum().item()

        if i % 10 == 9:
            print(f'Epoch [{epoch + 1}/{num_epochs}], '
                  f'Step [{i + 1}/{len(dataloader)}], '
                  f'Loss: {running_loss / 10:.4f}')
            running_loss = 0.0

    # 计算并保存准确率
    accuracy = correct_predictions / total_samples
    train_accuracies.append(accuracy)

    # 计算平均损失并保存
    average_loss = running_loss / len(dataloader)
    train_losses.append(average_loss)

    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {average_loss}, Accuracy: {accuracy}')

# 绘制训练过程图表
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(range(1, num_epochs + 1), train_losses, label='Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(range(1, num_epochs + 1), train_accuracies, label='Training Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()

plt.show()

小结：暂时记录这么多以后学的多了，再对上面进行修改

贴个代码，这个是图像二分类的问题，可根据自己的实际情况进行修改

import os
os.environ['TORCH_USE_CUDA_DSA'] = '1'
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder
from efficientnet_pytorch import EfficientNet
import torch.nn as nn
import torchvision.models as models
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np


# 自定义数据集类
class MyDataset(Dataset):
    def __init__(self, root, transform=None):
        self.dataset = ImageFolder(root, transform=transform)
        self.classes = self.dataset.classes  # 获取ImageFolder对象的classes属性

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        img, label = self.dataset[idx]
        return img, label

# 数据预处理和增强
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 创建自定义数据集实例
dataset = MyDataset(root='相对于项目的--数据集的路径', transform=transform)

# 输出类别信息
print("Classes:", dataset.classes)

# 创建数据加载器
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 加载预训练的 EfficientNet 模型
#model = models.efficientnet_b4(pretrained=True)
model = EfficientNet.from_pretrained('efficientnet-b4')

# 查看模型的结构
#print(model)
# 获取全连接层
feature = model._fc.in_features
num_classes = 3
model._fc = nn.Linear(in_features=feature,out_features=3,bias=True)
#print(model)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 1
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 初始化存储训练过程中的损失和准确率
train_losses = []
train_accuracies = []
#初始化变量跟踪准确率
total = 0
correct = 0

for epoch in range(num_epochs):
    running_loss = 0.0
    correct_predictions = 0
    total_samples = 0

    for i, data in enumerate(dataloader,0):
        inputs, labels = inputs.to(device), labels.to(device)

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

        running_loss += loss.item()

        # 计算准确率
        _, predicted = torch.max(outputs.data, 1)
        total_samples += labels.size(0)
        correct_predictions += (predicted == labels).sum().item()
# Print statistics every 10 steps每10步记录一次
        if i % 10 == 9:
            print(f'Epoch [{epoch + 1}/{num_epochs}], '
                  f'Step [{i + 1}/{len(dataloader)}], '
                  f'Loss: {running_loss / 10:.4f},'
                  f'Accuracy: {(correct / total) * 100:.2f}%')
                 
            running_loss = 0.0

    # 计算并保存准确率
    accuracy = correct_predictions / total_samples
    train_accuracies.append(accuracy)

    # 计算平均损失并保存
    average_loss = running_loss / len(dataloader)
    train_losses.append(average_loss)

    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {average_loss}, Accuracy: {accuracy}')

# 绘制训练过程图表
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(range(1, num_epochs + 1), train_losses, label='Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(range(1, num_epochs + 1), train_accuracies, label='Training Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()

plt.show()


# Save the trained model
torch.save(model.state_dict(), 'model_efficient.pth')