ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

PyTorch迁移学习实战:从ImageNet预训练模型到自定义图像分类

PyTorch迁移学习实战:从ImageNet预训练模型到自定义图像分类 在深度学习项目实践中我们常常面临一个现实困境自己的数据集规模有限从头训练一个复杂的神经网络不仅耗时漫长对计算资源要求极高而且模型性能往往难以达到理想状态。尤其是在图像分类、目标检测等任务中直接训练一个像ResNet、VGG这样的深度模型对于个人开发者或中小团队来说几乎是一项不可能完成的任务。迁移学习正是解决这一难题的“银弹”。它允许我们站在巨人的肩膀上利用在大规模数据集如ImageNet上预训练好的模型快速适配到自己的特定任务上极大地缩短了开发周期并显著提升了模型在小数据集上的表现。本文将围绕“PyTorch实战15-迁移学习”这一主题为你提供一套从理论到实践的完整闭环解决方案。无论你是刚接触PyTorch的新手希望快速上手一个图像分类项目还是有一定基础的开发者想系统掌握迁移学习的工程化技巧与调优策略本文都能提供直接的帮助。我们将从核心概念讲起手把手带你完成环境搭建、模型加载、数据准备、微调训练、模型评估的全流程并深入探讨不同微调策略的适用场景最后附上实战中高频出现的“坑点”与解决方案。学完本文你将能够独立使用PyTorch的torchvision.models模块高效地将预训练模型应用到自己的业务场景中。1. 迁移学习核心概念与为什么选择它在深入代码之前我们必须先理解迁移学习“是什么”以及“为什么有效”。这能帮助你在后续实践中做出更明智的技术决策而不是盲目套用代码。1.1 什么是迁移学习迁移学习是一种机器学习方法其核心思想是将从一个任务源任务中学到的知识通常是模型参数或特征表示应用到另一个相关但不同的任务目标任务上。在深度学习的语境下这通常意味着源任务在一个大规模、通用数据集如ImageNet包含1000个类别、1400万张图片上训练一个深度神经网络。学到的知识模型在训练过程中学习到的、能够有效识别图像中低级特征如边缘、纹理和高级语义特征如物体部件、整体形状的权重参数。目标任务我们自己的、数据量可能较小的特定任务如识别猫狗品种、医学影像分析、卫星图像分类等。通过迁移学习我们不是从随机初始化的权重开始训练而是从一个已经具备强大特征提取能力的“预训练模型”的权重开始。这相当于让模型从一个很高的起点开始学习我们的新任务。1.2 为什么迁移学习在图像领域如此有效这源于卷积神经网络CNN的层级结构特性。一个典型的CNN如VGG、ResNet可以看作一个特征提取器浅层网络靠近输入学习的是通用、低级的视觉特征如边缘、角点、颜色、纹理。这些特征对于绝大多数视觉任务都是有用的具有很高的可迁移性。深层网络靠近输出学习的是与特定任务高度相关的高级、抽象语义特征如“车轮”、“猫脸”、“建筑轮廓”等。这些特征的可迁移性相对较低。因此当我们把在ImageNet上预训练的模型迁移到自己的猫狗分类任务时模型浅层学到的“边缘检测器”、“纹理识别器”完全可以直接复用。我们只需要用自己数据集的数据去重新训练微调模型的最后几层全连接层让它们学习如何将通用的高级特征组合成“猫”或“狗”的特定概念。这大大降低了对数据量的需求并加速了收敛。1.3 PyTorch中的迁移学习方案PyTorch通过torchvision.models子模块提供了丰富的预训练模型极大简化了迁移学习的流程。针对不同的目标任务和数据情况我们主要有两种策略特征提取器Feature Extractor做法冻结预训练模型的所有层将其权重设为不可训练仅将其作为一个固定的特征提取器。然后移除其原有的分类头通常是最后的全连接层并添加一个全新的、针对自己任务类别数的分类器新的全连接层。训练只训练我们新添加的分类器层。适用场景目标任务数据集非常小且与源任务ImageNet差异较大。此时微调所有层容易导致过拟合。微调Fine-tuning做法不冻结预训练模型而是让其所有权重均可训练。同样我们需要修改最后的分类头以适应新的类别数。训练训练整个网络但通常使用一个较小的学习率以避免破坏预训练模型中已经学到的良好特征。适用场景目标任务数据集相对充足例如每个类别有几百到几千张图片且与源任务有一定相关性。这是最常用、效果通常也更好的策略。本文将重点讲解微调Fine-tuning策略因为它在实际项目中应用最广。同时我们也会对比介绍特征提取器模式让你能根据自身情况灵活选择。2. 环境准备与项目结构工欲善其事必先利其器。一个清晰的环境和项目结构是成功的第一步。2.1 软硬件环境说明操作系统Windows 10/11 macOS 或 Linux (如Ubuntu 20.04)均可。本文命令以Linux/macOS的bash为例Windows用户可在PowerShell或Anaconda Prompt中运行对应命令。Python推荐使用Python 3.8 或 3.9。版本过高或过低可能导致某些库的兼容性问题。深度学习框架PyTorch。我们将使用其torch和torchvision库。GPU可选但强烈推荐虽然CPU可以运行但训练CNN模型非常缓慢。建议使用NVIDIA GPU并安装对应版本的CUDA和cuDNN以加速训练。你可以通过nvidia-smi命令查看GPU信息。包管理工具推荐使用conda或pip。conda能更好地处理环境隔离和依赖冲突。2.2 使用Conda创建虚拟环境并安装PyTorch虚拟环境可以避免项目间的包版本冲突。这是Python项目开发的最佳实践。# 1. 创建一个名为 pytorch-transfer 的虚拟环境并指定Python版本 conda create -n pytorch-transfer python3.9 # 2. 激活该环境 conda activate pytorch-transfer # 3. 安装PyTorch。请务必访问PyTorch官网(https://pytorch.org/get-started/locally/)获取最适合你系统的安装命令。 # 以下是一个示例针对CUDA 11.8的Linux系统 conda install pytorch torchvision torchaudio pytorch-cuda11.8 -c pytorch -c nvidia # 如果你没有GPU或者想先使用CPU版本 # conda install pytorch torchvision torchaudio cpuonly -c pytorch # 4. 安装其他必要的库 pip install matplotlib opencv-python pillow pandas scikit-learn jupyter # matplotlib: 绘图 # opencv-python: 图像处理 # pillow: Python图像处理库 # pandas: 数据处理 # scikit-learn: 评估指标 # jupyter: 可选用于交互式开发重要提示PyTorch版本与CUDA版本必须严格匹配。安装前请确认你的NVIDIA驱动支持的CUDA版本通过nvidia-smi查看然后去官网选择对应命令。版本不匹配是导致“CUDA不可用”等错误的常见原因。2.3 验证安装创建一个Python脚本或直接在激活的环境下启动Python解释器运行以下代码import torch import torchvision print(fPyTorch版本: {torch.__version__}) print(fTorchvision版本: {torchvision.__version__}) print(fCUDA是否可用: {torch.cuda.is_available()}) if torch.cuda.is_available(): print(fGPU设备: {torch.cuda.get_device_name(0)})如果输出显示CUDA可用并且版本号正常则环境配置成功。2.4 项目目录结构在开始编码前建议建立如下清晰的目录结构这有助于管理数据、代码和模型。pytorch_transfer_learning/ │ ├── data/ # 存放数据集 │ ├── train/ # 训练集 │ │ ├── class1/ # 类别1的图片 │ │ ├── class2/ # 类别2的图片 │ │ └── ... │ ├── val/ # 验证集结构同train │ └── test/ # 测试集可选结构同train │ ├── src/ # 源代码 │ ├── dataset.py # 自定义数据集类如果需要 │ ├── model.py # 模型定义与加载 │ ├── train.py # 训练脚本 │ ├── utils.py # 工具函数如可视化、指标计算 │ └── config.py # 配置文件超参数、路径等 │ ├── outputs/ # 输出目录 │ ├── models/ # 保存训练好的模型权重 │ ├── logs/ # 训练日志如TensorBoard │ └── predictions/ # 模型预测结果 │ ├── notebooks/ # Jupyter Notebook文件用于探索性分析 │ ├── requirements.txt # 项目依赖包列表 └── README.md # 项目说明对于本教程为了简化我们将在一个Jupyter Notebook或一个Python脚本中完成所有步骤。但了解这个结构对构建真实项目至关重要。3. 数据准备以猫狗分类为例没有数据再好的模型也无用武之地。我们将使用一个经典的猫狗分类小型数据集作为示例。你可以从Kaggle下载“Dogs vs Cats”数据集或者使用任何其他分类数据集如花卉分类、手写数字等。关键是数据集的目录结构必须符合ImageFolder的期望。3.1 数据集结构与划分torchvision.datasets.ImageFolder要求数据按以下结构组织data_root/ train/ class_a/ image1.jpg image2.jpg ... class_b/ image1.jpg image2.jpg ... val/ class_a/ ... class_b/ ...每个子文件夹的名称就是该类别的标签。ImageFolder会自动根据文件夹名创建标签映射如class_a- 0,class_b- 1。假设我们的data目录下已经有cat和dog两个文件夹分别存放猫和狗的图片。我们需要手动或编写脚本将其划分为训练集和验证集通常比例是8:2或7:3。3.2 使用torchvision.transforms进行数据增强与标准化数据预处理是提升模型泛化能力的关键。对于使用ImageNet预训练的模型我们必须使用ImageNet的均值和标准差进行标准化因为模型是在这种数据分布下训练的。import torch from torchvision import transforms # 定义训练和验证/测试的数据预处理管道 # ImageNet的均值和标准差 IMAGENET_MEAN [0.485, 0.456, 0.406] IMAGENET_STD [0.229, 0.224, 0.225] # 训练集变换增强 标准化 train_transform transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪并缩放到224x224 transforms.RandomHorizontalFlip(p0.5), # 随机水平翻转 transforms.RandomRotation(degrees15), # 随机旋转 transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.1), # 颜色抖动 transforms.ToTensor(), # 转换为Tensor并归一化到[0,1] transforms.Normalize(meanIMAGENET_MEAN, stdIMAGENET_STD) # 标准化 ]) # 验证集/测试集变换仅标准化不进行增强 val_transform transforms.Compose([ transforms.Resize(256), # 将短边缩放到256 transforms.CenterCrop(224), # 中心裁剪到224x224 transforms.ToTensor(), transforms.Normalize(meanIMAGENET_MEAN, stdIMAGENET_STD) ]) # 解释 # 1. RandomResizedCrop: 数据增强模拟物体在不同位置、不同尺度的出现。 # 2. RandomHorizontalFlip: 对于猫狗这类水平对称的物体翻转是有效的增强。 # 3. Normalize: 至关重要使用预训练模型对应的统计量能加速收敛并稳定训练。3.3 创建DataLoaderDataLoader负责批量加载数据并支持多进程数据加载是训练循环中的核心组件。from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader # 假设数据路径 data_dir ./data train_dir f{data_dir}/train val_dir f{data_dir}/val # 创建数据集对象 train_dataset ImageFolder(roottrain_dir, transformtrain_transform) val_dataset ImageFolder(rootval_dir, transformval_transform) # 查看数据集信息 print(f训练集类别: {train_dataset.classes}) # 输出: [cat, dog] print(f训练集类别到索引的映射: {train_dataset.class_to_idx}) # 输出: {cat: 0, dog: 1} print(f训练集样本数: {len(train_dataset)}) print(f验证集样本数: {len(val_dataset)}) # 创建数据加载器 batch_size 32 num_workers 4 # 根据你的CPU核心数设置用于并行加载数据 train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workersnum_workers, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_sizebatch_size, shuffleFalse, num_workersnum_workers, pin_memoryTrue) # 解释参数 # shuffleTrue: 训练时每个epoch打乱数据顺序防止模型学习到数据顺序。 # num_workers: 子进程数用于加速数据加载。Windows下有时设为0可避免问题。 # pin_memoryTrue: 将数据固定在页锁定内存加速GPU数据传输当使用GPU时。现在数据管道已经搭建完成。我们可以通过迭代DataLoader来查看一批数据。# 获取一个批次的数据 images, labels next(iter(train_loader)) print(f图像批次形状: {images.shape}) # [batch_size, channels, height, width] - [32, 3, 224, 224] print(f标签批次形状: {labels.shape}) # [32]4. 加载与修改预训练模型这是迁移学习的核心步骤。我们将以ResNet18为例展示如何加载预训练模型并修改其最后的全连接层以适应我们的二分类任务。4.1 加载预训练ResNet18PyTorch的torchvision.models提供了pretrainedTrue参数可以自动下载并加载在ImageNet上预训练的权重。import torchvision.models as models import torch.nn as nn # 加载预训练的ResNet18模型 # pretrainedTrue 会自动下载权重。第一次运行需要联网。 model models.resnet18(pretrainedTrue) # 查看模型结构 print(model)观察输出你会看到ResNet18的最后一部分是一个全连接层 (model.fc):... (avgpool): AdaptiveAvgPool2d(output_size(1, 1)) (fc): Linear(in_features512, out_features1000, biasTrue)in_features512是ResNet18最后一个卷积层输出的特征维度out_features1000对应ImageNet的1000个类别。4.2 修改分类头全连接层我们的任务是猫狗二分类所以需要将out_features改为2。同时我们可能还想冻结模型的一部分层只训练最后几层。# 方案一特征提取器模式冻结所有卷积层只训练新分类头 def create_feature_extractor(model, num_classes2): # 1. 冻结所有模型参数 for param in model.parameters(): param.requires_grad False # 2. 获取最后一个全连接层的输入特征数 num_ftrs model.fc.in_features # 3. 用一个新的、适合我们任务的全连接层替换原来的fc层 # 这个新层的参数 requires_grad 默认为 True model.fc nn.Linear(num_ftrs, num_classes) return model # 方案二微调模式解冻所有层或部分层并修改分类头 def create_fine_tune_model(model, num_classes2, feature_extractFalse): Args: model: 预训练模型 num_classes: 目标类别数 feature_extract: 如果为True则冻结所有层特征提取器模式 如果为False则解冻所有层微调模式。 if feature_extract: # 特征提取器模式冻结所有参数 for param in model.parameters(): param.requires_grad False # 否则微调模式所有参数默认 requires_grad True (预训练权重加载时已是True) # 修改分类头 num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes) return model # 使用微调模式创建我们的模型 model_ft create_fine_tune_model(model, num_classes2, feature_extractFalse) print(model_ft)现在model_ft.fc是一个新的Linear(512, 2)层它的参数是可训练的。模型前面的卷积层参数也是可训练的因为feature_extractFalse。4.3 将模型移动到GPU如果可用device torch.device(cuda:0 if torch.cuda.is_available() else cpu) print(f使用设备: {device}) model_ft model_ft.to(device)5. 训练与验证循环有了模型和数据接下来就是编写训练循环。我们将定义损失函数、优化器并实现完整的训练和验证流程。5.1 定义损失函数与优化器对于二分类任务我们使用交叉熵损失nn.CrossEntropyLoss。它内部已经包含了Softmax操作所以模型的最后一层不需要加Softmax激活函数但在推理时需要。优化器我们选择Adam它通常比SGD收敛更快且对学习率不那么敏感。对于微调我们通常为新添加的分类头设置较大的学习率为预训练的卷积层设置较小的学习率以防止破坏已有的好特征。import torch.optim as optim criterion nn.CrossEntropyLoss() # 损失函数 # 观察模型中哪些参数需要训练 params_to_update [] for name, param in model_ft.named_parameters(): if param.requires_grad: params_to_update.append(param) print(f\t{name}) # 打印所有可训练的参数名 # 为不同层设置不同的学习率 # 我们将模型参数分为两组特征提取层卷积层和分类头fc层 feature_params [] classifier_params [] for name, param in model_ft.named_parameters(): if not param.requires_grad: continue # 跳过冻结的参数 if fc in name: # 分类头的参数 classifier_params.append(param) else: # 特征提取层的参数 feature_params.append(param) # 创建优化器为两组参数设置不同的学习率 optimizer optim.Adam([ {params: feature_params, lr: 1e-4}, # 预训练层小学习率 {params: classifier_params, lr: 1e-3} # 新分类头大学习率 ], weight_decay1e-4) # weight_decay是L2正则化防止过拟合5.2 编写训练与验证函数def train_one_epoch(model, dataloader, criterion, optimizer, device): 训练一个epoch model.train() # 设置为训练模式启用Dropout, BatchNorm等 running_loss 0.0 running_corrects 0 for inputs, labels in dataloader: inputs inputs.to(device) labels labels.to(device) # 清零梯度 optimizer.zero_grad() # 前向传播 outputs model(inputs) loss criterion(outputs, labels) # 获取预测结果最大概率的类别 _, preds torch.max(outputs, 1) # 反向传播与优化 loss.backward() optimizer.step() # 统计 running_loss loss.item() * inputs.size(0) running_corrects torch.sum(preds labels.data) epoch_loss running_loss / len(dataloader.dataset) epoch_acc running_corrects.double() / len(dataloader.dataset) return epoch_loss, epoch_acc def validate(model, dataloader, criterion, device): 验证/测试模型 model.eval() # 设置为评估模式禁用Dropout, 固定BatchNorm的统计量 running_loss 0.0 running_corrects 0 with torch.no_grad(): # 不计算梯度节省内存和计算 for inputs, labels in dataloader: inputs inputs.to(device) labels labels.to(device) outputs model(inputs) loss criterion(outputs, labels) _, preds torch.max(outputs, 1) running_loss loss.item() * inputs.size(0) running_corrects torch.sum(preds labels.data) epoch_loss running_loss / len(dataloader.dataset) epoch_acc running_corrects.double() / len(dataloader.dataset) return epoch_loss, epoch_acc5.3 主训练循环现在我们将所有部分组合起来进行多轮epoch训练。import time import copy num_epochs 10 best_model_wts copy.deepcopy(model_ft.state_dict()) # 深拷贝保存最佳权重 best_acc 0.0 history {train_loss: [], train_acc: [], val_loss: [], val_acc: []} for epoch in range(num_epochs): print(fEpoch {epoch1}/{num_epochs}) print(- * 20) # 训练阶段 train_loss, train_acc train_one_epoch(model_ft, train_loader, criterion, optimizer, device) print(fTrain Loss: {train_loss:.4f} Acc: {train_acc:.4f}) # 验证阶段 val_loss, val_acc validate(model_ft, val_loader, criterion, device) print(fVal Loss: {val_loss:.4f} Acc: {val_acc:.4f}) print() # 记录历史 history[train_loss].append(train_loss) history[train_acc].append(train_acc) history[val_loss].append(val_loss) history[val_acc].append(val_acc) # 如果验证准确率提升则保存模型 if val_acc best_acc: best_acc val_acc best_model_wts copy.deepcopy(model_ft.state_dict()) # 可以在这里保存模型到文件 torch.save(model_ft.state_dict(), f./outputs/models/best_model_epoch{epoch1}.pth) print(f 模型已保存当前最佳准确率: {best_acc:.4f}) # 训练结束后加载最佳模型权重 model_ft.load_state_dict(best_model_wts) print(f训练完成最佳验证准确率: {best_acc:.4f})5.4 可视化训练过程使用Matplotlib绘制损失和准确率曲线有助于分析模型是否过拟合或欠拟合。import matplotlib.pyplot as plt plt.figure(figsize(12, 4)) # 绘制损失曲线 plt.subplot(1, 2, 1) plt.plot(history[train_loss], labelTrain Loss) plt.plot(history[val_loss], labelVal Loss) plt.title(Training and Validation Loss) plt.xlabel(Epoch) plt.ylabel(Loss) plt.legend() # 绘制准确率曲线 plt.subplot(1, 2, 2) plt.plot(history[train_acc], labelTrain Acc) plt.plot(history[val_acc], labelVal Acc) plt.title(Training and Validation Accuracy) plt.xlabel(Epoch) plt.ylabel(Accuracy) plt.legend() plt.tight_layout() plt.show()理想情况训练损失和验证损失都稳步下降训练准确率和验证准确率都稳步上升且两者最终差距不大。过拟合迹象训练损失持续下降但验证损失在某个点后开始上升训练准确率远高于验证准确率。欠拟合迹象训练损失和验证损失都很高且下降缓慢准确率提升有限。6. 模型推理与测试训练完成后我们需要用测试集或单张图片来评估模型的最终性能。6.1 在测试集上评估假设我们有一个独立的测试集test_loader评估方式与验证阶段相同。# 加载最佳模型如果之前保存了 # model_ft.load_state_dict(torch.load(./outputs/models/best_model.pth)) model_ft.eval() test_loss, test_acc validate(model_ft, test_loader, criterion, device) print(f测试集损失: {test_loss:.4f}, 测试集准确率: {test_acc:.4f})6.2 单张图片预测在实际应用中我们更常需要对单张图片进行预测。以下是完整的推理流程from PIL import Image import numpy as np def predict_single_image(image_path, model, transform, device, class_names): 对单张图片进行预测 Args: image_path: 图片路径 model: 训练好的模型 transform: 验证时使用的数据预处理管道 device: CPU或GPU class_names: 类别名称列表如 [cat, dog] Returns: pred_class: 预测的类别名 prob: 预测的概率 # 1. 加载并预处理图片 image Image.open(image_path).convert(RGB) # 确保是RGB三通道 image_tensor transform(image).unsqueeze(0) # 增加一个批次维度 - [1, C, H, W] image_tensor image_tensor.to(device) # 2. 模型推理 model.eval() with torch.no_grad(): outputs model(image_tensor) # outputs 形状: [1, num_classes] # 3. 获取预测结果 probabilities torch.nn.functional.softmax(outputs, dim1) # 转换为概率 confidence, predicted_idx torch.max(probabilities, 1) # 4. 返回结果 pred_class class_names[predicted_idx.item()] prob confidence.item() return pred_class, prob # 使用示例 class_names train_dataset.classes # [cat, dog] img_path ./data/test/dog_example.jpg # 替换为你的测试图片路径 pred_class, prob predict_single_image(img_path, model_ft, val_transform, device, class_names) print(f预测类别: {pred_class}, 置信度: {prob:.4f}) # 可视化预测结果可选 img Image.open(img_path) plt.imshow(img) plt.title(fPrediction: {pred_class} ({prob:.2%})) plt.axis(off) plt.show()7. 常见问题与排查思路避坑指南在实际操作中你可能会遇到各种问题。下面列出了一些高频“坑点”及其解决方案。问题现象可能原因排查思路与解决方案CUDA out of memory1. 批次大小batch_size太大。2. 模型太大。3. 显卡显存不足。1.减小batch_size如从32降到16或8。2. 使用更小的预训练模型如ResNet18代替ResNet50。3. 使用torch.cuda.empty_cache()清理缓存。4. 在代码中使用with torch.no_grad():减少验证时的显存占用。5. 使用梯度累积gradient accumulation模拟大批次。验证准确率远低于训练准确率过拟合1. 训练数据太少。2. 模型复杂度太高。3. 数据增强不够。4. 训练时间太长过拟合训练集。1.增加数据增强的强度和多样性。2. 在模型中添加Dropout层如果原模型没有。3. 使用更早的停止Early Stopping根据验证集损失不再下降时停止训练。4. 增强L2正则化weight_decay。5. 尝试特征提取器模式冻结卷积层如果数据量真的非常少。训练损失不下降或准确率不提升1. 学习率设置不当太大或太小。2. 数据预处理错误如未标准化。3. 模型权重未正确加载或修改。4. 标签错误或数据损坏。1.调整学习率尝试一个范围如1e-2, 1e-3, 1e-4。使用学习率调度器如StepLR。2.检查transforms.Normalize的均值和标准差是否正确必须使用ImageNet的。3.打印模型参数确认新分类头的参数requires_gradTrue且被优化器收录。4.可视化一批训练数据确认图片和标签对应正确。RuntimeError: Expected all tensors to be on the same device模型和数据不在同一个设备上一个在CPU一个在GPU。确保在训练循环开始前将模型和每一批数据都移动到同一个设备model.to(device)inputs, labels inputs.to(device), labels.to(device)预训练模型下载慢或失败网络连接问题。1. 使用国内镜像源如清华源。可以手动下载权重文件然后通过model.load_state_dict(torch.load(‘path/to/weights.pth’))加载。2. 检查PyTorch版本与模型版本的兼容性。ImageFolder报错Found 0 files in subfolders数据集目录结构不符合要求。严格按照root/class_name/*.jpg的结构组织数据。检查文件夹名称和图片后缀名。微调后模型性能反而变差学习率对于预训练层可能太大破坏了好的特征。采用差分学习率为预训练层设置更小的学习率如1e-4/1e-5为新分类头设置较大的学习率如1e-3。这正是我们在优化器部分所做的。8. 进阶技巧与最佳实践掌握了基础流程后以下技巧能帮助你进一步提升模型效果和工程效率。8.1 学习率调度Learning Rate Scheduling固定学习率可能不是最优的。使用学习率调度器可以在训练过程中动态调整学习率例如在验证指标停滞时降低学习率。from torch.optim import lr_scheduler # 创建优化器后添加调度器 # StepLR: 每 step_size 个epoch将学习率乘以 gamma # scheduler lr_scheduler.StepLR(optimizer, step_size7, gamma0.1) # ReduceLROnPlateau: 当监控的指标如val_loss停止改善时降低学习率。 scheduler lr_scheduler.ReduceLROnPlateau(optimizer, modemin, factor0.1, patience3, verboseTrue) # modemin 表示监控的指标越低越好如loss。 # patience3 表示容忍指标3个epoch没有改善然后降低学习率。 # 在验证阶段之后每个epoch调用 # scheduler.step(val_loss) # 对于ReduceLROnPlateau # scheduler.step() # 对于StepLR8.2 使用TensorBoard可视化训练过程TensorBoard可以更直观地监控损失、准确率、权重分布等。from torch.utils.tensorboard import SummaryWriter import os # 创建SummaryWriter log_dir ./outputs/logs os.makedirs(log_dir, exist_okTrue) writer SummaryWriter(log_dirlog_dir) # 在训练循环中记录标量 for epoch in range(num_epochs): # ... 训练和验证 ... train_loss, train_acc ... val_loss, val_acc ... # 记录到TensorBoard writer.add_scalar(Loss/train, train_loss, epoch) writer.add_scalar(Accuracy/train, train_acc, epoch) writer.add_scalar(Loss/val, val_loss, epoch) writer.add_scalar(Accuracy/val, val_acc, epoch) # 还可以记录学习率 for i, param_group in enumerate(optimizer.param_groups): writer.add_scalar(fLearningRate/group_{i}, param_group[lr], epoch) # 训练结束后关闭writer writer.close()在终端运行tensorboard --logdir./outputs/logs即可在浏览器查看。8.3 尝试不同的预训练模型torchvision.models提供了多种模型适用于不同的计算资源和精度要求。轻量级/快速mobilenet_v3_small,shufflenet_v2_x1_0平衡型resnet18,resnet34,efficientnet_b0高精度resnet50,resnet101,efficientnet_b3,vit_b_16更换模型通常只需要修改一行代码model models.resnet50(pretrainedTrue) # 改为ResNet50 # 注意不同模型的分类头属性名可能不同如ResNet是.fcVGG是.classifier[6]需要查看模型结构后相应修改。8.4 模型保存与加载的完整流程保存和加载整个模型包括结构或仅保存状态字典推荐。# 保存 # 方法1保存整个模型不推荐因为与代码绑定过紧 torch.save(model_ft, full_model.pth) # 方法2仅保存状态字典推荐 torch.save({ epoch: epoch, model_state_dict: model_ft.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict() if scheduler else None, best_acc: best_acc, class_to_idx: train_dataset.class_to_idx, # 保存类别映射 }, ./outputs/models/checkpoint.pth) # 加载 checkpoint torch.load(./outputs/models/checkpoint.pth, map_locationdevice) # 重新初始化模型结构必须和保存时一致 model_ft models.resnet18(pretrainedFalse) # 不加载预训练权重 num_ftrs model_ft.fc.in_features model_ft.fc nn.Linear(num_ftrs, 2) # 修改分类头类别数需一致 model_ft.load_state_dict(checkpoint[model_state_dict]) model_ft.to(device) # 加载其他状态如果需要继续训练 epoch checkpoint[epoch] best_acc checkpoint[best_acc] optimizer.load_state_dict(checkpoint[optimizer_state_dict]) if scheduler and checkpoint[scheduler_state_dict]: scheduler.load_state_dict(checkpoint[scheduler_state_dict]) class_to_idx checkpoint[class_to_idx]8.5 在生产环境部署的考虑模型导出为TorchScript如果需要脱离Python环境部署如C服务可以使用torch.jit.trace或torch.jit.script将模型转换为TorchScript。example_input torch.rand(1, 3, 224, 224).to(device) traced_script_module torch.jit.trace(model_ft, example_input) traced_script_module.save(model_scripted.pt)使用ONNX格式如果需要与其他框架如TensorFlow, OpenVINO交互可以导出为ONNX格式。优化推理速度使用半精度torch.float16推理、模型剪枝、量化等技术。构建API服务使用Flask、FastAPI等框架将模型封装为REST API。迁移学习是深度学习工程化应用中最实用、最高效的技术之一。通过本文的实战演练你应该已经掌握了使用PyTorch进行迁移学习的完整流程从理解核心思想、准备数据、加载并修改预训练模型到编写训练循环、评估模型以及处理常见问题。关键在于根据自己数据集的大小和与预训练任务的相似性灵活选择特征提取或微调策略并合理设置差分学习率。下一步你可以更换数据集尝试将自己的业务数据如工控缺陷检测、遥感图像分类套用此流程。尝试不同模型对比ResNet、EfficientNet、Vision Transformer等不同架构在你自己任务上的效果。探索更高级的技巧如使用自定义学习率调度器、集成学习、测试时增强TTA等。深入模型可解释性使用Grad-CAM等工具可视化模型关注的图像区域理解其决策依据。记住实践出真知。多动手修改代码、调整参数、分析结果是掌握迁移学习乃至深度学习的最佳途径。如果在实践中遇到新的问题欢迎在社区交流讨论。
返回列表