[Pytorch案例实践006]基于迁移学习-ResNet18的蚂蚁&蜜蜂图像分类实战

一、项目介绍

        此项目的目标是对图像数据集进行分类任务。它使用了`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,