这次我们来看一个非常实用的深度学习实战项目如何用少量图片通过 ResNet 迁移学习来完成图像分类任务。对于很多开发者来说收集海量标注数据是训练深度学习模型的最大障碍。这个项目直接解决了这个痛点它演示了如何利用 PyTorch 和预训练的 ResNet 模型在仅有少量数据的情况下快速构建一个有效的图像分类器。项目的核心价值在于其“轻量”和“高效”。你不需要准备 TB 级的数据集也不需要从零开始训练一个庞大的模型。通过迁移学习我们可以站在“巨人”即在大规模数据集上预训练好的 ResNet的肩膀上只针对自己的少量特定图片进行微调就能获得相当不错的分类效果。这对于学术研究、原型验证、或者特定垂直领域如工业质检、医学影像辅助分析的应用开发具有极高的实践意义。本文将带你完整走通这个流程。我们会重点关注几个实操要点环境如何搭建、数据如何准备、模型如何修改与加载、训练过程如何监控、以及最终如何验证模型效果。整个过程会基于 PyTorch 框架代码清晰步骤明确。无论你是刚入门深度学习的新手还是希望快速验证某个分类想法的开发者这篇文章都能提供直接的帮助。1. 核心能力速览在深入代码之前我们先快速了解这个项目的关键信息让你判断是否值得投入时间。能力项说明项目类型深度学习实战教程 / 图像分类代码实现技术栈PyTorch, torchvision, ResNet 预训练模型核心方法迁移学习 (Transfer Learning)数据要求少量标注图像即可无需海量数据硬件门槛支持 CPU 训练GPU 可大幅加速。GPU 显存需求取决于图像尺寸和批次大小通常 2GB 以上显存即可进行微调。主要功能1. 加载预训练 ResNet 模型2. 替换并重新训练模型最后一层分类头3. 在自定义小数据集上进行模型微调4. 评估模型在测试集上的分类准确率输出成果一个针对特定类别优化过的 PyTorch 模型文件 (.pth)适合场景学术实验、项目原型验证、特定领域的轻量级图像分类应用开发2. 适用场景与使用边界适合谁深度学习初学者想通过一个完整的项目理解迁移学习的整个流程。算法工程师/研究者需要快速为一个新概念或新品类构建图像分类原型验证想法的可行性。特定领域开发者例如开发识别特定植物病害、工业零件缺陷、零售商品分类等应用但缺乏大规模标注数据。能解决什么问题数据稀缺问题在只有几百张甚至几十张标注图片的情况下训练一个可用的分类模型。训练效率问题避免从零开始训练深度网络节省大量计算资源和时间。模型泛化起点高利用在 ImageNet 等大数据集上学到的通用特征如边缘、纹理、形状让模型在特定任务上更快收敛、效果更好。不适合什么场景超细粒度分类如果需要区分同一犬种下不同血统仅靠微调最后一层可能不够需要更深入的网络调整和数据增强。与预训练任务差异极大的领域如果您的图像风格与自然图像ImageNet截然不同如显微图像、卫星云图底层特征迁移效果可能打折扣可能需要部分或全部重新训练特征提取层。追求极致精度如果拥有充足的数据和算力从零训练或更复杂的架构调优可能获得更好上限。合规与伦理提醒数据版权确保用于训练和测试的图像数据拥有合法使用权避免侵犯他人版权或肖像权。应用边界本技术为通用图像分类方法请勿将其用于任何违反法律法规或公序良俗的图像识别场景。3. 环境准备与前置条件开始之前请确保你的开发环境满足以下要求。这是一个标准的 PyTorch 深度学习环境。3.1 基础软件环境操作系统: Windows 10/11, Linux (如 Ubuntu 18.04), 或 macOS。Python: 版本 3.7 至 3.10 为佳。推荐使用 Anaconda 或 Miniconda 管理环境。包管理工具:pip或conda。3.2 核心依赖包我们将使用 PyTorch 及其视觉工具库 torchvision。以下是通过pip安装的命令。请根据你是否使用 GPU 以及 CUDA 版本选择合适的 PyTorch 安装命令。# 方案一安装支持 CUDA 的 PyTorch (推荐有 NVIDIA GPU 的用户) # 请访问 https://pytorch.org/get-started/locally/ 获取最新安装命令例如对于 CUDA 11.8 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 方案二安装仅支持 CPU 的 PyTorch (无 GPU 或仅测试) pip install torch torchvision torchaudio # 通用工具库 pip install numpy pandas matplotlib tqdm scikit-learn pillow3.3 硬件要求检查GPU (可选但推荐): NVIDIA GPU并安装对应版本的 CUDA 和 cuDNN。使用nvidia-smi命令可查看 GPU 状态。显存: 微调 ResNet-18/34 时若使用常见批次大小如 32和图像尺寸224x2244GB 显存通常足够。ResNet-50 等更大模型需要更多显存。内存: 建议 8GB 以上系统内存。磁盘空间: 预留至少 2-5GB 空间用于存放数据集、预训练模型缓存和训练好的模型。4. 数据准备构建小型图像数据集迁移学习成功的关键之一是正确的数据组织。我们采用 PyTorchImageFolder期望的目录结构。4.1 目录结构假设你的项目根目录为ai_study_transfer按如下方式组织数据ai_study_transfer/ ├── data/ │ ├── train/ │ │ ├── class_a/ │ │ │ ├── image1.jpg │ │ │ └── image2.jpg │ │ └── class_b/ │ │ ├── image1.jpg │ │ └── image2.jpg │ └── val/ (或 test/) │ ├── class_a/ │ │ └── image3.jpg │ └── class_b/ │ └── image3.jpg ├── src/ │ └── train.py └── models/ └── (训练好的模型将保存于此)train/: 训练集。每个子文件夹代表一个类别文件夹名即为类别标签文件夹内放置该类别的图片。val/或test/: 验证集/测试集。结构与训练集相同用于评估模型性能。4.2 数据量要求每个类别至少需要几十到上百张图片。如果数据极少如每类少于20张需要更强的数据增强Data Augmentation来防止过拟合。4.3 数据预处理与增强我们将在代码中使用torchvision.transforms来定义数据转换流程。通常包括训练集转换随机裁剪、水平翻转、颜色抖动、归一化等以增加数据多样性。验证/测试集转换中心裁剪、缩放、归一化不进行随机性增强。5. 模型构建修改与加载 ResNet这是迁移学习的核心步骤。我们以 ResNet-18 为例。5.1 加载预训练模型PyTorch 的torchvision.models提供了预训练的 ResNet 模型在 ImageNet 上训练过。import torch import torchvision.models as models import torch.nn as nn # 检查是否有可用的 GPU device torch.device(cuda:0 if torch.cuda.is_available() else cpu) print(fUsing device: {device}) # 加载预训练的 ResNet-18 模型 # pretrainedTrue 会自动下载模型权重首次运行需要网络 model models.resnet18(pretrainedTrue) # 将模型移动到指定设备GPU/CPU model model.to(device) # 冻结所有网络参数在微调初期我们通常先冻结特征提取层 for param in model.parameters(): param.requires_grad False5.2 修改最后一层分类头预训练 ResNet 的最后一层全连接层 (model.fc) 输出是 1000 维对应 ImageNet 的1000个类。我们需要将其替换为适合自己类别数量的新层。# 假设我们的数据有 num_classes 个类别 num_classes 5 # 例如有5种不同的花朵 # 获取原全连接层的输入特征数 num_ftrs model.fc.in_features # 替换整个全连接层 # 新的全连接层默认 requires_gradTrue model.fc nn.Linear(num_ftrs, num_classes) # 将新的全连接层移动到设备上 model.fc model.fc.to(device) # 现在只有 model.fc 的参数是需要训练requires_gradTrue的 # 我们可以选择在训练后期解冻部分卷积层进行微调6. 训练流程实现接下来我们编写完整的训练脚本。脚本将包括数据加载、模型训练、验证和保存。6.1 数据加载与预处理import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 数据增强和归一化配置 # ImageNet 的均值和标准差常用于迁移学习 data_transforms { train: transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪并缩放至224x224 transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.ToTensor(), # 转为Tensor并归一化至[0,1] transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 标准化 ]), val: transforms.Compose([ transforms.Resize(256), # 将短边缩放到256 transforms.CenterCrop(224), # 中心裁剪224x224 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } # 数据目录路径 data_dir ./data image_datasets { x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in [train, val] } # 创建数据加载器 (DataLoader) dataloaders { x: DataLoader(image_datasets[x], batch_size32, shuffle(xtrain), num_workers4) for x in [train, val] } dataset_sizes {x: len(image_datasets[x]) for x in [train, val]} class_names image_datasets[train].classes print(fClass names: {class_names}) print(fDataset sizes: {dataset_sizes})6.2 定义损失函数与优化器我们只训练新替换的全连接层因此优化器只传入需要更新的参数。import torch.optim as optim # 定义损失函数交叉熵损失适用于多分类 criterion nn.CrossEntropyLoss() # 观察哪些参数需要训练 params_to_update [] for name, param in model.named_parameters(): if param.requires_grad: params_to_update.append(param) print(f\t{name}) # 定义优化器这里使用带动量的SGD是训练CNN的常见选择 # 只传入需要更新的参数可以节省内存和计算量 optimizer optim.SGD(params_to_update, lr0.001, momentum0.9) # 定义学习率调度器每7个epoch将学习率乘以0.1 scheduler optim.lr_scheduler.StepLR(optimizer, step_size7, gamma0.1)6.3 训练与验证循环这是核心的训练逻辑每个 epoch 会遍历训练集进行参数更新并在验证集上评估性能。def train_model(model, criterion, optimizer, scheduler, num_epochs25): best_acc 0.0 best_model_wts model.state_dict().copy() for epoch in range(num_epochs): print(fEpoch {epoch}/{num_epochs - 1}) print(- * 10) # 每个epoch都有训练和验证阶段 for phase in [train, val]: if phase train: model.train() # 设置模型为训练模式 else: model.eval() # 设置模型为评估模式 running_loss 0.0 running_corrects 0 # 遍历数据 for inputs, labels in dataloaders[phase]: 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) if phase train: scheduler.step() # 更新学习率 epoch_loss running_loss / dataset_sizes[phase] epoch_acc running_corrects.double() / dataset_sizes[phase] print(f{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}) # 深度拷贝并保存最佳模型权重 if phase val and epoch_acc best_acc: best_acc epoch_acc best_model_wts model.state_dict().copy() # 可以在这里保存模型 torch.save(model.state_dict(), f./models/best_model_epoch{epoch}.pth) print(f Best model saved at epoch {epoch} with acc {best_acc:.4f}) print() print(fBest val Acc: {best_acc:.4f}) # 加载最佳模型权重 model.load_state_dict(best_model_wts) return model # 开始训练 model train_model(model, criterion, optimizer, scheduler, num_epochs20)7. 模型测试与推理训练完成后我们需要用测试集或保留的验证集评估最终模型并编写一个简单的推理函数。7.1 加载最佳模型进行测试# 假设我们有一个独立的测试集目录 data/test/ test_dataset datasets.ImageFolder(os.path.join(data_dir, test), data_transforms[val]) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse, num_workers4) model.eval() # 非常重要确保模型在评估模式 test_corrects 0 total 0 with torch.no_grad(): # 关闭梯度计算节省内存和计算 for inputs, labels in test_loader: inputs inputs.to(device) labels labels.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) test_corrects torch.sum(preds labels.data) total labels.size(0) test_acc test_corrects.double() / total print(fTest Accuracy: {test_acc:.4f})7.2 单张图片推理函数这个函数可以方便地用于实际应用。from PIL import Image def predict_single_image(image_path, model, class_names, devicecuda:0): 对单张图片进行预测 Args: image_path: 图片文件路径 model: 加载好权重的PyTorch模型 class_names: 类别名称列表 device: 运行设备 Returns: pred_class: 预测的类别名称 confidence: 预测置信度概率 # 加载并预处理图像 image Image.open(image_path).convert(RGB) preprocess data_transforms[val] # 使用验证集的转换 image_tensor preprocess(image).unsqueeze(0) # 增加一个批次维度 [1, C, H, W] image_tensor image_tensor.to(device) model.eval() with torch.no_grad(): outputs model(image_tensor) # 应用Softmax获取概率 probabilities torch.nn.functional.softmax(outputs[0], dim0) # 获取最高置信度的索引和值 confidence, predicted_idx torch.max(probabilities, 0) predicted_idx predicted_idx.item() confidence confidence.item() pred_class class_names[predicted_idx] return pred_class, confidence # 使用示例 # model ... # 加载你训练好的模型 # class_names [daisy, dandelion, roses, sunflowers, tulips] # 你的类别 # pred, conf predict_single_image(./test_image.jpg, model, class_names, device) # print(fPredicted: {pred} with confidence {conf:.2%})8. 资源占用与性能观察在本地运行此项目时监控资源使用情况有助于优化和排查问题。8.1 显存与内存观察训练阶段显存占用主要取决于batch_size和模型大小。对于 ResNet-18 和 batch_size32在 224x224 输入下GPU 显存占用通常在1.5GB ~ 3GB之间。如果显存不足首要降低batch_size如改为 16 或 8。推理阶段显存占用远小于训练因为不需要保存中间变量用于反向传播。CPU 内存数据加载器 (DataLoader) 的num_workers参数会创建子进程加载数据增加 CPU 内存占用。如果内存不足可将其设为 0。8.2 性能优化建议使用 GPU这是最大的加速因素。确保torch.cuda.is_available()返回True。调整批量大小在显存允许范围内较大的batch_size通常能使训练更稳定、更快。可通过torch.cuda.memory_allocated()监控显存。使用pin_memory在 GPU 训练时设置DataLoader的pin_memoryTrue可以加速主机到设备的数据传输。梯度累积如果显存只能支持很小的batch_size可以采用梯度累积技术多次前向传播后再更新一次参数模拟大batch_size的效果。9. 常见问题与排查方法在实践过程中你可能会遇到以下问题。这里提供排查思路。问题现象可能原因排查方式解决方案运行时错误CUDA out of memory批次大小 (batch_size) 过大或图像尺寸过大。运行nvidia-smi观察显存占用。在代码开始时打印torch.cuda.memory_allocated()。1. 减小batch_size。2. 减小输入图像尺寸如从 224 降到 128。3. 使用更小的模型如 ResNet-18 而非 ResNet-50。训练损失 (Loss) 不下降学习率 (lr) 设置不当模型层被冻结参数未更新数据或标签有问题。1. 检查params_to_update列表是否非空。2. 打印几个批次的损失看是否有变化。3. 可视化几张训练图片和对应标签。1. 调整学习率尝试 0.01, 0.001, 0.0001。2. 确保需要训练的层如model.fc的requires_gradTrue。3. 检查数据路径和ImageFolder读取是否正确。验证准确率远低于训练准确率模型过拟合。数据量太少或数据增强不够。对比训练集和验证集的 Loss/Acc 曲线。1. 增强数据增强力度如增加随机旋转、颜色抖动。2. 添加正则化如 Dropout。3. 收集更多数据。4. 更早地停止训练Early Stopping。ImageFolder找不到图片或类别数为0数据目录结构错误图片格式不被支持。打印image_datasets[‘train’].classes和dataset_sizes。检查文件夹命名和图片后缀。严格按照train/class_name/*.jpg的结构组织。确保文件夹非空图片为常见格式jpg, png。预测结果全部为同一类别类别不平衡严重模型未充分训练最后一层初始化有问题。检查训练集中每个类别的图片数量。观察训练初期损失是否快速下降。1. 对训练数据进行重采样或使用类别权重。2. 增加训练轮数 (num_epochs)。3. 尝试解冻部分卷积层进行微调。10. 进阶技巧与最佳实践掌握了基础流程后这些技巧可以帮助你获得更好的效果或应对更复杂的场景。10.1 解冻部分卷积层进行微调在训练后期可以解冻靠近顶部的卷积层让它们也进行微调以适应新任务。# 例如解冻 model.layer4 和 model.layer3 的参数 for name, param in model.named_parameters(): # 冻结所有参数 param.requires_grad False # 解冻最后一层fc和最后两个卷积块layer4, layer3 unfreeze_layers [fc, layer4, layer3] for name, param in model.named_parameters(): for unfreeze_name in unfreeze_layers: if unfreeze_name in name: param.requires_grad True break # 重新定义优化器只传入需要梯度的参数 optimizer optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr0.0001, momentum0.9)10.2 使用不同的预训练模型除了 ResNet-18torchvision 提供了多种模型可以根据任务复杂度和硬件条件选择。# ResNet 系列复杂度递增 model models.resnet34(pretrainedTrue) model models.resnet50(pretrainedTrue) # 更深特征更丰富需要更多显存 model models.resnet101(pretrainedTrue) # 轻量级模型 model models.mobilenet_v3_small(pretrainedTrue) model models.efficientnet_b0(pretrainedTrue) # 修改分类头的方法类似但属性名可能不同需查看模型结构 # 例如 MobileNet V3: model.classifier[-1] nn.Linear(...) # 例如 EfficientNet: model.classifier[1] nn.Linear(...)10.3 自动化实验记录使用 TensorBoard 或 Weights Biases 记录损失、准确率曲线方便分析和比较不同实验。# 使用 TensorBoard 的简单示例 from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/experiment_1) for epoch in range(num_epochs): # ... 训练代码 ... writer.add_scalar(Loss/train, epoch_loss, epoch) writer.add_scalar(Accuracy/train, epoch_acc, epoch) # ... 验证代码 ... writer.add_scalar(Loss/val, val_loss, epoch) writer.add_scalar(Accuracy/val, val_acc, epoch) writer.close()通过这个项目你不仅学会了用 ResNet 做迁移学习的代码实现更重要的是掌握了一套在小数据场景下启动图像分类任务的标准化流程。从环境搭建、数据准备、模型修改、训练验证到问题排查每一步都是实际开发中会反复遇到的环节。建议你用自己的图片数据集哪怕是手机随手拍的十几种不同物品跑一遍整个流程这种亲手调试、观察模型从零开始“学会”分类的过程远比阅读理论更能加深理解。
这次我们来看一个非常实用的深度学习实战项目如何用少量图片通过 ResNet 迁移学习来完成图像分类任务。对于很多开发者来说收集海量标注数据是训练深度学习模型的最大障碍。这个项目直接解决了这个痛点它演示了如何利用 PyTorch 和预训练的 ResNet 模型在仅有少量数据的情况下快速构建一个有效的图像分类器。项目的核心价值在于其“轻量”和“高效”。你不需要准备 TB 级的数据集也不需要从零开始训练一个庞大的模型。通过迁移学习我们可以站在“巨人”即在大规模数据集上预训练好的 ResNet的肩膀上只针对自己的少量特定图片进行微调就能获得相当不错的分类效果。这对于学术研究、原型验证、或者特定垂直领域如工业质检、医学影像辅助分析的应用开发具有极高的实践意义。本文将带你完整走通这个流程。我们会重点关注几个实操要点环境如何搭建、数据如何准备、模型如何修改与加载、训练过程如何监控、以及最终如何验证模型效果。整个过程会基于 PyTorch 框架代码清晰步骤明确。无论你是刚入门深度学习的新手还是希望快速验证某个分类想法的开发者这篇文章都能提供直接的帮助。1. 核心能力速览在深入代码之前我们先快速了解这个项目的关键信息让你判断是否值得投入时间。能力项说明项目类型深度学习实战教程 / 图像分类代码实现技术栈PyTorch, torchvision, ResNet 预训练模型核心方法迁移学习 (Transfer Learning)数据要求少量标注图像即可无需海量数据硬件门槛支持 CPU 训练GPU 可大幅加速。GPU 显存需求取决于图像尺寸和批次大小通常 2GB 以上显存即可进行微调。主要功能1. 加载预训练 ResNet 模型2. 替换并重新训练模型最后一层分类头3. 在自定义小数据集上进行模型微调4. 评估模型在测试集上的分类准确率输出成果一个针对特定类别优化过的 PyTorch 模型文件 (.pth)适合场景学术实验、项目原型验证、特定领域的轻量级图像分类应用开发2. 适用场景与使用边界适合谁深度学习初学者想通过一个完整的项目理解迁移学习的整个流程。算法工程师/研究者需要快速为一个新概念或新品类构建图像分类原型验证想法的可行性。特定领域开发者例如开发识别特定植物病害、工业零件缺陷、零售商品分类等应用但缺乏大规模标注数据。能解决什么问题数据稀缺问题在只有几百张甚至几十张标注图片的情况下训练一个可用的分类模型。训练效率问题避免从零开始训练深度网络节省大量计算资源和时间。模型泛化起点高利用在 ImageNet 等大数据集上学到的通用特征如边缘、纹理、形状让模型在特定任务上更快收敛、效果更好。不适合什么场景超细粒度分类如果需要区分同一犬种下不同血统仅靠微调最后一层可能不够需要更深入的网络调整和数据增强。与预训练任务差异极大的领域如果您的图像风格与自然图像ImageNet截然不同如显微图像、卫星云图底层特征迁移效果可能打折扣可能需要部分或全部重新训练特征提取层。追求极致精度如果拥有充足的数据和算力从零训练或更复杂的架构调优可能获得更好上限。合规与伦理提醒数据版权确保用于训练和测试的图像数据拥有合法使用权避免侵犯他人版权或肖像权。应用边界本技术为通用图像分类方法请勿将其用于任何违反法律法规或公序良俗的图像识别场景。3. 环境准备与前置条件开始之前请确保你的开发环境满足以下要求。这是一个标准的 PyTorch 深度学习环境。3.1 基础软件环境操作系统: Windows 10/11, Linux (如 Ubuntu 18.04), 或 macOS。Python: 版本 3.7 至 3.10 为佳。推荐使用 Anaconda 或 Miniconda 管理环境。包管理工具:pip或conda。3.2 核心依赖包我们将使用 PyTorch 及其视觉工具库 torchvision。以下是通过pip安装的命令。请根据你是否使用 GPU 以及 CUDA 版本选择合适的 PyTorch 安装命令。# 方案一安装支持 CUDA 的 PyTorch (推荐有 NVIDIA GPU 的用户) # 请访问 https://pytorch.org/get-started/locally/ 获取最新安装命令例如对于 CUDA 11.8 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 方案二安装仅支持 CPU 的 PyTorch (无 GPU 或仅测试) pip install torch torchvision torchaudio # 通用工具库 pip install numpy pandas matplotlib tqdm scikit-learn pillow3.3 硬件要求检查GPU (可选但推荐): NVIDIA GPU并安装对应版本的 CUDA 和 cuDNN。使用nvidia-smi命令可查看 GPU 状态。显存: 微调 ResNet-18/34 时若使用常见批次大小如 32和图像尺寸224x2244GB 显存通常足够。ResNet-50 等更大模型需要更多显存。内存: 建议 8GB 以上系统内存。磁盘空间: 预留至少 2-5GB 空间用于存放数据集、预训练模型缓存和训练好的模型。4. 数据准备构建小型图像数据集迁移学习成功的关键之一是正确的数据组织。我们采用 PyTorchImageFolder期望的目录结构。4.1 目录结构假设你的项目根目录为ai_study_transfer按如下方式组织数据ai_study_transfer/ ├── data/ │ ├── train/ │ │ ├── class_a/ │ │ │ ├── image1.jpg │ │ │ └── image2.jpg │ │ └── class_b/ │ │ ├── image1.jpg │ │ └── image2.jpg │ └── val/ (或 test/) │ ├── class_a/ │ │ └── image3.jpg │ └── class_b/ │ └── image3.jpg ├── src/ │ └── train.py └── models/ └── (训练好的模型将保存于此)train/: 训练集。每个子文件夹代表一个类别文件夹名即为类别标签文件夹内放置该类别的图片。val/或test/: 验证集/测试集。结构与训练集相同用于评估模型性能。4.2 数据量要求每个类别至少需要几十到上百张图片。如果数据极少如每类少于20张需要更强的数据增强Data Augmentation来防止过拟合。4.3 数据预处理与增强我们将在代码中使用torchvision.transforms来定义数据转换流程。通常包括训练集转换随机裁剪、水平翻转、颜色抖动、归一化等以增加数据多样性。验证/测试集转换中心裁剪、缩放、归一化不进行随机性增强。5. 模型构建修改与加载 ResNet这是迁移学习的核心步骤。我们以 ResNet-18 为例。5.1 加载预训练模型PyTorch 的torchvision.models提供了预训练的 ResNet 模型在 ImageNet 上训练过。import torch import torchvision.models as models import torch.nn as nn # 检查是否有可用的 GPU device torch.device(cuda:0 if torch.cuda.is_available() else cpu) print(fUsing device: {device}) # 加载预训练的 ResNet-18 模型 # pretrainedTrue 会自动下载模型权重首次运行需要网络 model models.resnet18(pretrainedTrue) # 将模型移动到指定设备GPU/CPU model model.to(device) # 冻结所有网络参数在微调初期我们通常先冻结特征提取层 for param in model.parameters(): param.requires_grad False5.2 修改最后一层分类头预训练 ResNet 的最后一层全连接层 (model.fc) 输出是 1000 维对应 ImageNet 的1000个类。我们需要将其替换为适合自己类别数量的新层。# 假设我们的数据有 num_classes 个类别 num_classes 5 # 例如有5种不同的花朵 # 获取原全连接层的输入特征数 num_ftrs model.fc.in_features # 替换整个全连接层 # 新的全连接层默认 requires_gradTrue model.fc nn.Linear(num_ftrs, num_classes) # 将新的全连接层移动到设备上 model.fc model.fc.to(device) # 现在只有 model.fc 的参数是需要训练requires_gradTrue的 # 我们可以选择在训练后期解冻部分卷积层进行微调6. 训练流程实现接下来我们编写完整的训练脚本。脚本将包括数据加载、模型训练、验证和保存。6.1 数据加载与预处理import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 数据增强和归一化配置 # ImageNet 的均值和标准差常用于迁移学习 data_transforms { train: transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪并缩放至224x224 transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.ToTensor(), # 转为Tensor并归一化至[0,1] transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 标准化 ]), val: transforms.Compose([ transforms.Resize(256), # 将短边缩放到256 transforms.CenterCrop(224), # 中心裁剪224x224 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } # 数据目录路径 data_dir ./data image_datasets { x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in [train, val] } # 创建数据加载器 (DataLoader) dataloaders { x: DataLoader(image_datasets[x], batch_size32, shuffle(xtrain), num_workers4) for x in [train, val] } dataset_sizes {x: len(image_datasets[x]) for x in [train, val]} class_names image_datasets[train].classes print(fClass names: {class_names}) print(fDataset sizes: {dataset_sizes})6.2 定义损失函数与优化器我们只训练新替换的全连接层因此优化器只传入需要更新的参数。import torch.optim as optim # 定义损失函数交叉熵损失适用于多分类 criterion nn.CrossEntropyLoss() # 观察哪些参数需要训练 params_to_update [] for name, param in model.named_parameters(): if param.requires_grad: params_to_update.append(param) print(f\t{name}) # 定义优化器这里使用带动量的SGD是训练CNN的常见选择 # 只传入需要更新的参数可以节省内存和计算量 optimizer optim.SGD(params_to_update, lr0.001, momentum0.9) # 定义学习率调度器每7个epoch将学习率乘以0.1 scheduler optim.lr_scheduler.StepLR(optimizer, step_size7, gamma0.1)6.3 训练与验证循环这是核心的训练逻辑每个 epoch 会遍历训练集进行参数更新并在验证集上评估性能。def train_model(model, criterion, optimizer, scheduler, num_epochs25): best_acc 0.0 best_model_wts model.state_dict().copy() for epoch in range(num_epochs): print(fEpoch {epoch}/{num_epochs - 1}) print(- * 10) # 每个epoch都有训练和验证阶段 for phase in [train, val]: if phase train: model.train() # 设置模型为训练模式 else: model.eval() # 设置模型为评估模式 running_loss 0.0 running_corrects 0 # 遍历数据 for inputs, labels in dataloaders[phase]: 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) if phase train: scheduler.step() # 更新学习率 epoch_loss running_loss / dataset_sizes[phase] epoch_acc running_corrects.double() / dataset_sizes[phase] print(f{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}) # 深度拷贝并保存最佳模型权重 if phase val and epoch_acc best_acc: best_acc epoch_acc best_model_wts model.state_dict().copy() # 可以在这里保存模型 torch.save(model.state_dict(), f./models/best_model_epoch{epoch}.pth) print(f Best model saved at epoch {epoch} with acc {best_acc:.4f}) print() print(fBest val Acc: {best_acc:.4f}) # 加载最佳模型权重 model.load_state_dict(best_model_wts) return model # 开始训练 model train_model(model, criterion, optimizer, scheduler, num_epochs20)7. 模型测试与推理训练完成后我们需要用测试集或保留的验证集评估最终模型并编写一个简单的推理函数。7.1 加载最佳模型进行测试# 假设我们有一个独立的测试集目录 data/test/ test_dataset datasets.ImageFolder(os.path.join(data_dir, test), data_transforms[val]) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse, num_workers4) model.eval() # 非常重要确保模型在评估模式 test_corrects 0 total 0 with torch.no_grad(): # 关闭梯度计算节省内存和计算 for inputs, labels in test_loader: inputs inputs.to(device) labels labels.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) test_corrects torch.sum(preds labels.data) total labels.size(0) test_acc test_corrects.double() / total print(fTest Accuracy: {test_acc:.4f})7.2 单张图片推理函数这个函数可以方便地用于实际应用。from PIL import Image def predict_single_image(image_path, model, class_names, devicecuda:0): 对单张图片进行预测 Args: image_path: 图片文件路径 model: 加载好权重的PyTorch模型 class_names: 类别名称列表 device: 运行设备 Returns: pred_class: 预测的类别名称 confidence: 预测置信度概率 # 加载并预处理图像 image Image.open(image_path).convert(RGB) preprocess data_transforms[val] # 使用验证集的转换 image_tensor preprocess(image).unsqueeze(0) # 增加一个批次维度 [1, C, H, W] image_tensor image_tensor.to(device) model.eval() with torch.no_grad(): outputs model(image_tensor) # 应用Softmax获取概率 probabilities torch.nn.functional.softmax(outputs[0], dim0) # 获取最高置信度的索引和值 confidence, predicted_idx torch.max(probabilities, 0) predicted_idx predicted_idx.item() confidence confidence.item() pred_class class_names[predicted_idx] return pred_class, confidence # 使用示例 # model ... # 加载你训练好的模型 # class_names [daisy, dandelion, roses, sunflowers, tulips] # 你的类别 # pred, conf predict_single_image(./test_image.jpg, model, class_names, device) # print(fPredicted: {pred} with confidence {conf:.2%})8. 资源占用与性能观察在本地运行此项目时监控资源使用情况有助于优化和排查问题。8.1 显存与内存观察训练阶段显存占用主要取决于batch_size和模型大小。对于 ResNet-18 和 batch_size32在 224x224 输入下GPU 显存占用通常在1.5GB ~ 3GB之间。如果显存不足首要降低batch_size如改为 16 或 8。推理阶段显存占用远小于训练因为不需要保存中间变量用于反向传播。CPU 内存数据加载器 (DataLoader) 的num_workers参数会创建子进程加载数据增加 CPU 内存占用。如果内存不足可将其设为 0。8.2 性能优化建议使用 GPU这是最大的加速因素。确保torch.cuda.is_available()返回True。调整批量大小在显存允许范围内较大的batch_size通常能使训练更稳定、更快。可通过torch.cuda.memory_allocated()监控显存。使用pin_memory在 GPU 训练时设置DataLoader的pin_memoryTrue可以加速主机到设备的数据传输。梯度累积如果显存只能支持很小的batch_size可以采用梯度累积技术多次前向传播后再更新一次参数模拟大batch_size的效果。9. 常见问题与排查方法在实践过程中你可能会遇到以下问题。这里提供排查思路。问题现象可能原因排查方式解决方案运行时错误CUDA out of memory批次大小 (batch_size) 过大或图像尺寸过大。运行nvidia-smi观察显存占用。在代码开始时打印torch.cuda.memory_allocated()。1. 减小batch_size。2. 减小输入图像尺寸如从 224 降到 128。3. 使用更小的模型如 ResNet-18 而非 ResNet-50。训练损失 (Loss) 不下降学习率 (lr) 设置不当模型层被冻结参数未更新数据或标签有问题。1. 检查params_to_update列表是否非空。2. 打印几个批次的损失看是否有变化。3. 可视化几张训练图片和对应标签。1. 调整学习率尝试 0.01, 0.001, 0.0001。2. 确保需要训练的层如model.fc的requires_gradTrue。3. 检查数据路径和ImageFolder读取是否正确。验证准确率远低于训练准确率模型过拟合。数据量太少或数据增强不够。对比训练集和验证集的 Loss/Acc 曲线。1. 增强数据增强力度如增加随机旋转、颜色抖动。2. 添加正则化如 Dropout。3. 收集更多数据。4. 更早地停止训练Early Stopping。ImageFolder找不到图片或类别数为0数据目录结构错误图片格式不被支持。打印image_datasets[‘train’].classes和dataset_sizes。检查文件夹命名和图片后缀。严格按照train/class_name/*.jpg的结构组织。确保文件夹非空图片为常见格式jpg, png。预测结果全部为同一类别类别不平衡严重模型未充分训练最后一层初始化有问题。检查训练集中每个类别的图片数量。观察训练初期损失是否快速下降。1. 对训练数据进行重采样或使用类别权重。2. 增加训练轮数 (num_epochs)。3. 尝试解冻部分卷积层进行微调。10. 进阶技巧与最佳实践掌握了基础流程后这些技巧可以帮助你获得更好的效果或应对更复杂的场景。10.1 解冻部分卷积层进行微调在训练后期可以解冻靠近顶部的卷积层让它们也进行微调以适应新任务。# 例如解冻 model.layer4 和 model.layer3 的参数 for name, param in model.named_parameters(): # 冻结所有参数 param.requires_grad False # 解冻最后一层fc和最后两个卷积块layer4, layer3 unfreeze_layers [fc, layer4, layer3] for name, param in model.named_parameters(): for unfreeze_name in unfreeze_layers: if unfreeze_name in name: param.requires_grad True break # 重新定义优化器只传入需要梯度的参数 optimizer optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr0.0001, momentum0.9)10.2 使用不同的预训练模型除了 ResNet-18torchvision 提供了多种模型可以根据任务复杂度和硬件条件选择。# ResNet 系列复杂度递增 model models.resnet34(pretrainedTrue) model models.resnet50(pretrainedTrue) # 更深特征更丰富需要更多显存 model models.resnet101(pretrainedTrue) # 轻量级模型 model models.mobilenet_v3_small(pretrainedTrue) model models.efficientnet_b0(pretrainedTrue) # 修改分类头的方法类似但属性名可能不同需查看模型结构 # 例如 MobileNet V3: model.classifier[-1] nn.Linear(...) # 例如 EfficientNet: model.classifier[1] nn.Linear(...)10.3 自动化实验记录使用 TensorBoard 或 Weights Biases 记录损失、准确率曲线方便分析和比较不同实验。# 使用 TensorBoard 的简单示例 from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/experiment_1) for epoch in range(num_epochs): # ... 训练代码 ... writer.add_scalar(Loss/train, epoch_loss, epoch) writer.add_scalar(Accuracy/train, epoch_acc, epoch) # ... 验证代码 ... writer.add_scalar(Loss/val, val_loss, epoch) writer.add_scalar(Accuracy/val, val_acc, epoch) writer.close()通过这个项目你不仅学会了用 ResNet 做迁移学习的代码实现更重要的是掌握了一套在小数据场景下启动图像分类任务的标准化流程。从环境搭建、数据准备、模型修改、训练验证到问题排查每一步都是实际开发中会反复遇到的环节。建议你用自己的图片数据集哪怕是手机随手拍的十几种不同物品跑一遍整个流程这种亲手调试、观察模型从零开始“学会”分类的过程远比阅读理论更能加深理解。
