一、项目介绍

        此项目的目标是对图像数据集进行分类任务。它使用了`resnet18`作为基础模型，并对其进行微调以适应新的数据集。这里采用的是迁移学习的一种常见方式：微调（Fine-tuning）。

        迁移学习是一种机器学习方法，通过在大型数据集上预训练好的模型，然后将这些模型应用于不同的但相关的任务。这种方法可以显著减少新任务所需的训练时间和数据量。

项目细节

1. 定义模型 (`YourModel` 类)

   - 使用了`torchvision.models.resnet18`作为基础网络，并且明确指定了`pretrained=False`来不加载预训练权重。
   - 修改了`resnet18`的最后一层全连接层(`fc`)，使其输出节点数等于新任务中的类别数量。

2. 加载预训练权重

   - 项目开始时尝试从指定路径加载预训练权重。
   - 如果找到了预训练权重，则加载到模型中；如果没有找到，则从头开始训练。

3. 数据准备

   - 使用了`torchvision.transforms`来对图像数据进行预处理，包括随机裁剪、翻转等操作。
   - 使用了`ImageFolder`来加载训练和验证数据集，并且定义了相应的数据加载器(`DataLoader`)。

4. 训练过程 (`train_model` 函数)

   - 实现了训练和验证循环。
   - 使用了交叉熵损失函数(`nn.CrossEntropyLoss`)作为损失计算的标准。
   - 使用了随机梯度下降优化器(`SGD`)来更新模型的参数，并且设置了学习率衰减策略(`lr_scheduler.StepLR`)。

5. 评估和可视化

   - 在每个epoch后记录训练和验证的损失及准确率，并绘制学习曲线。
   - 最终保存训练好的模型，并绘制出训练和验证的损失、准确率以及学习率的变化曲线。

        迁移学习的方法能够有效利用已有模型的知识，并且可以极大地减少训练时间和需要的数据量。这对于小数据集和资源有限的情况非常有用。

二、数据集介绍

        数据集在上个项目介绍过，在此不再赘述。

三、完整代码

训练代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import matplotlib.pyplot as plt
import numpy as np
import os
from tqdm import tqdm

# 定义模型类
import torchvision.models as models

class YourModel(nn.Module):
    def __init__(self, num_classes):
        super(YourModel, self).__init__()
        self.model = models.resnet18(pretrained=False)  # 不使用预训练权重
        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)

    def forward(self, x):
        return self.model(x)

# 训练与验证的函数
def train_model(model, criterion, optimizer, scheduler, dataloaders, device, num_epochs=25):
    train_losses = []
    val_losses = []
    train_accuracies = []
    val_accuracies = []
    lrs = []

    for epoch in range(num_epochs):
        print(f'Epoch {epoch+1}/{num_epochs}')
        print('-' * 10)

        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()
            else:
                model.eval()

            running_loss = 0.0
            running_corrects = 0
            total = 0

            for inputs, labels in tqdm(dataloaders[phase], desc=f'{phase} Epoch {epoch+1}'):
                inputs = inputs.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()

                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
                total += labels.size(0)

            epoch_loss = running_loss / total
            epoch_acc = running_corrects.double() / total

            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

            if phase == 'train':
                train_losses.append(epoch_loss)
                train_accuracies.append(epoch_acc)
                lrs.append(optimizer.param_groups[0]['lr'])
                scheduler.step()
            else:
                val_losses.append(epoch_loss)
                val_accuracies.append(epoch_acc)

        print()

    return train_losses, val_losses, train_accuracies, val_accuracies, lrs

# 保存模型
def save_model(model, path='model.pth'):
    torch.save(model.state_dict(), path)
    print(f'Model saved to {path}')

# 绘制训练曲线
# 绘制训练曲线
def plot_curves(train, val, title, ylabel, filename):
    train = [t.cpu().numpy() if isinstance(t, torch.Tensor) else t for t in train]
    val = [v.cpu().numpy() if isinstance(v, torch.Tensor) else v for v in val]

    epochs = np.arange(1, len(train) + 1)
    plt.plot(epochs, train, 'r', label='Train')
    plt.plot(epochs, val, 'b', label='Validation')
    plt.title(title)
    plt.xlabel('Epoch')
    plt.ylabel(ylabel)
    plt.legend()
    plt.grid(True)
    plt.savefig(filename)
    plt.show()

def main():
    # 超参数设置
    num_epochs = 10
    batch_size = 32
    learning_rate = 0.001
    num_classes = 2  # 根据你的数据集调整类别数

    # 指定预训练权重路径
    pretrained_weights_path = r'I:\code\pytorch\resnet18\resnet18-5c106cde.pth'

    # 数据增强和标准化
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'val': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }

    # 加载数据集
    data_dir = r'I:\code\pytorch\resnet18\datasets'
    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                              data_transforms[x])
                      for x in ['train', 'val']}
    dataloaders = {x: DataLoader(image_datasets[x], batch_size=batch_size,
                                 shuffle=True, num_workers=4)
                   for x in ['train', 'val']}
    class_names = image_datasets['train'].classes

    # 模型实例化
    model = YourModel(num_classes=len(class_names))
    # print(len(class_names))
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    # 加载预训练权重
    if os.path.exists(pretrained_weights_path):
        model.load_state_dict(torch.load(pretrained_weights_path), strict=False)
        print(f"Loaded pretrained weights from {pretrained_weights_path}")
    else:
        print(f"Pretrained weights not found at {pretrained_weights_path}. Starting from scratch.")

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

    # 学习率调度器
    scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

    # 开始训练模型
    train_losses, val_losses, train_accuracies, val_accuracies, lrs = train_model(
        model, criterion, optimizer, scheduler, dataloaders, device, num_epochs)

    # 保存模型
    save_model(model)

    # 绘制并保存曲线
    plot_curves(train_losses, val_losses, 'Loss Curve', 'Loss', 'loss_curve.png')
    plot_curves(train_accuracies, val_accuracies, 'Accuracy Curve', 'Accuracy', 'accuracy_curve.png')

    # 学习率曲线
    plt.plot(np.arange(1, num_epochs+1), lrs, 'g', label='Learning Rate')
    plt.title('Learning Rate Curve')
    plt.xlabel('Epoch')
    plt.ylabel('Learning Rate')
    plt.grid(True)
    plt.savefig('learning_rate_curve.png')
    plt.show()

if __name__ == '__main__':
    main()

测试代码：

import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import matplotlib.pyplot as plt
import os


# 定义模型类
class YourModel(nn.Module):
    def __init__(self, num_classes):
        super(YourModel, self).__init__()
        self.model = models.resnet18(pretrained=False)
        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)

    def forward(self, x):
        return self.model(x)


# 加载模型并进行推理
def load_model(model_path, num_classes):
    model = YourModel(num_classes=num_classes)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.load_state_dict(torch.load(model_path, map_location=device))
    model = model.to(device)
    model.eval()
    return model


def predict(model, image_path, class_names):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 图像预处理
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    image = Image.open(image_path)
    image_tensor = transform(image).unsqueeze(0).to(device)

    # 模型推理
    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = nn.Softmax(dim=1)(outputs)
        confidence, predicted_idx = torch.max(probabilities, 1)
        predicted_idx = predicted_idx.item()
        confidence = confidence.item()
        predicted_class = class_names[predicted_idx]

    return predicted_class, confidence, image


def display_prediction(image, predicted_class, confidence):
    plt.figure(figsize=(8, 8))
    plt.imshow(image)
    plt.title(f"Predicted: {predicted_class} ({confidence * 100:.2f}%)", color='red', fontsize=16)
    plt.axis('off')
    plt.show()


def main():
    # 设置参数
    model_path = r'I:\code\pytorch\resnet18\model.pth'  # 模型路径
    test_image_path = r'I:\code\pytorch\resnet18\datasets\val\ants\263615709_cfb28f6b8e.jpg'  # 测试图片路径
    class_names = ['ants', 'bees']  # 替换为你的类别名称
    num_classes = len(class_names)

    # 加载模型
    model = load_model(model_path, num_classes)

    # 预测结果
    predicted_class, confidence, image = predict(model, test_image_path, class_names)

    # 显示结果
    display_prediction(image, predicted_class, confidence)


if __name__ == '__main__':
    main()

四、项目总结

        本项目成功地利用了迁移学习的技术，通过微调 ResNet-18 模型实现了对特定图像数据集的高效分类。通过数据增强和模型微调，模型在新任务上的性能得到了显著提升。此外，通过绘制训练过程中的曲线，我们可以直观地了解模型的学习状态，从而进一步优化训练过程。