这次我们来看一个非常实用的深度学习实战项目如何用少量图片通过 ResNet 迁移学习快速训练一个可用的图像分类模型。如果你手头只有几十张或几百张图片又想快速得到一个效果不错的分类器那么迁移学习几乎是唯一可行的路径。这篇文章不讲复杂的理论重点解决“能不能跑起来”和“效果怎么样”这两个实际问题。我们将围绕一个具体的实战案例展开使用 PyTorch 框架基于预训练的 ResNet 模型在自定义的小数据集上进行微调Fine-tuning。整个过程会重点关注环境搭建、数据准备、模型修改、训练技巧以及最终的模型验证。无论你是刚入门深度学习还是需要快速为某个特定任务如工业质检、特定物体识别搭建原型这套流程都能直接复用。本文会带你完成从零到一的完整过程从 PyTorch 和 CUDA 环境配置开始到准备你自己的图片数据集接着修改并加载预训练的 ResNet 模型然后进行训练和验证最后评估模型在少量数据上的实际表现。我们重点关注实操中的关键点例如如何应对小样本过拟合、如何设置合适的学习率、以及如何观察训练过程中的损失和准确率变化。1. 核心能力速览在深入代码之前我们先快速了解这个实战项目的核心信息让你判断是否值得投入时间。能力项说明项目类型深度学习实战教程 / 图像分类模型微调技术栈PyTorch, torchvision, ResNet (如 ResNet-18/34/50)核心方法迁移学习 (Transfer Learning) / 微调 (Fine-tuning)数据需求少量图片即可(每类几十到几百张)支持自定义类别硬件门槛支持 CPU 训练(速度慢)推荐使用 GPU(如 NVIDIA GTX 1060 6G 及以上) 以加速训练显存占用取决于模型大小和批次 (Batch Size)。以 ResNet-18 为例Batch Size32 时训练显存占用约 1.5-2.5 GB。主要功能1. 加载 ImageNet 预训练的 ResNet 模型。2. 替换模型最后一层全连接层适配自定义类别数。3. 在自定义小数据集上进行训练和验证。4. 保存训练好的模型并用于推理预测。输出成果训练好的.pth模型文件、训练过程日志、准确率/损失曲线图。适合场景学术研究入门、工业场景快速原型验证如缺陷检测、特定物品识别、个人兴趣项目如花卉、宠物分类。2. 适用场景与使用边界迁移学习特别是基于 ResNet 这类在大规模数据集如 ImageNet上预训练好的模型进行微调是解决小样本学习问题的利器。它最适合谁深度学习初学者希望绕过从零训练的巨大计算成本和数据需求快速体验模型训练全流程并获得成就感。算法工程师/研究者需要为某个垂直领域如医疗影像、遥感图像、特定商品识别快速搭建一个可用的分类基线模型。学生或爱好者手头有自己收集的图片集如自己拍的花、昆虫、手写字符想用 AI 技术做个分类工具。它能解决什么问题数据稀缺当你的标注数据只有几百甚至几十张时从零训练一个深度网络几乎必然过拟合而迁移学习可以极大缓解这个问题。计算资源有限预训练模型已经包含了丰富的通用视觉特征如边缘、纹理、形状微调只需要学习与新任务相关的特定特征训练周期短对算力要求低。快速验证想法在投入大量资源标注数据或训练大模型之前先用迁移学习验证任务可行性。它的边界在哪里领域差异过大如果您的自定义图片与 ImageNet 的自然图像分布差异极大例如全是显微镜下的细胞图片或卫星云图预训练模型提供的底层特征帮助可能有限需要更谨慎地调整训练策略或考虑使用在更相关领域预训练的模型。类别极度不平衡如果某些类别的图片数量远少于其他类别即使使用迁移学习模型也可能偏向多数类。需要结合数据增强、重采样或损失函数加权等技术。并非“万能药”迁移学习能提供一个强大的起点但最终模型性能的上限仍然受限于你的数据质量、数量以及任务本身的复杂度。版权与合规确保你用于训练和测试的图片拥有合法的使用权。对于涉及人脸、个人隐私或商业版权的图片务必在授权范围内使用。本项目代码和预训练模型如 torchvision 提供的通常用于学习和研究目的。3. 环境准备与前置条件工欲善其事必先利其器。下面是我们完成本次实战所需的环境清单。1. 操作系统推荐Ubuntu 18.04/20.04/22.04 LTS 或 Windows 10/11。说明Linux 在深度学习开发中更为常见环境配置问题相对较少。Windows 用户也可通过 Anaconda 顺利运行。2. Python 环境Python 版本3.8, 3.9 或 3.10与 PyTorch 版本兼容。环境管理强烈推荐使用Anaconda或Miniconda创建独立的虚拟环境避免包冲突。3. 深度学习框架PyTorch核心框架。需要根据你的 CUDA 版本选择安装命令。torchvision提供预训练模型ResNet和数据集处理工具。可通过 PyTorch 官网的安装命令生成器获取精确命令。4. 硬件要求CPU现代多核处理器即可。GPU推荐NVIDIA GPU并安装对应版本的CUDA Toolkit和cuDNN。这是加速训练的关键。常见组合PyTorch 1.12 CUDA 11.3或 PyTorch 2.0 CUDA 11.8/12.1。可使用nvidia-smi命令查看 GPU 和驱动信息。内存至少 8GB RAM推荐 16GB 以上。磁盘空间预留 5-10 GB 空间用于安装包、数据集和保存模型。5. 其他依赖库NumPy数值计算。Matplotlib绘制损失和准确率曲线。Pillow (PIL)图像处理。tqdm显示进度条可选但很实用。scikit-learn用于计算分类报告、混淆矩阵等评估指标可选。4. 安装部署与启动方式环境准备好了接下来我们一步步安装必要的软件包并准备好项目目录结构。步骤 1创建并激活 Conda 虚拟环境打开终端Linux/macOS或 Anaconda PromptWindows执行以下命令# 创建一个名为 pytorch_tl 的 Python 3.9 环境 conda create -n pytorch_tl python3.9 -y # 激活环境 conda activate pytorch_tl步骤 2安装 PyTorch 和 torchvision访问 PyTorch 官网 根据你的系统、包管理工具、CUDA 版本如果有GPU选择对应的安装命令。例如对于 Linux 系统、使用 Conda、CUDA 11.8 的用户命令可能如下conda install pytorch torchvision torchaudio pytorch-cuda11.8 -c pytorch -c nvidia对于仅使用 CPU 的用户命令类似conda install pytorch torchvision torchaudio cpuonly -c pytorch安装完成后可以启动 Python 交互环境验证import torch import torchvision print(torch.__version__) print(torch.cuda.is_available()) # 如果返回 True说明 GPU 可用步骤 3安装其他辅助库pip install numpy matplotlib pillow tqdm scikit-learn步骤 4准备项目目录结构一个清晰的项目结构有助于管理代码、数据和模型。建议按如下方式组织your_project/ ├── data/ # 存放数据集 │ ├── train/ # 训练集 │ │ ├── class_1/ # 类别1的图片 │ │ ├── class_2/ # 类别2的图片 │ │ └── ... │ └── val/ # 验证集结构同train ├── src/ # 存放源代码 │ ├── train.py # 主训练脚本 │ ├── predict.py # 推理预测脚本 │ └── utils.py # 工具函数如数据加载、可视化 ├── models/ # 存放训练好的模型文件 (.pth) ├── logs/ # 存放训练日志、曲线图 └── requirements.txt # 项目依赖列表步骤 5准备你的数据集这是最关键的一步。将你的图片按类别放入data/train和data/val的对应子文件夹中。torchvision.datasets.ImageFolder会自动根据文件夹名识别类别标签。训练集/验证集划分建议按 8:2 或 7:3 的比例划分。确保每个类别在训练集和验证集中都有出现。图片格式支持 JPG, PNG 等常见格式。图片大小不需要统一数据加载时会进行 resize 和裁剪。5. 功能测试与效果验证一切就绪让我们开始编写核心的训练代码并观察模型在少量数据上的学习效果。5.1 数据加载与预处理我们使用torchvision提供的ImageFolder和transforms来高效加载和增强数据。# src/utils.py 或 train.py 的一部分 import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader def get_dataloaders(data_dir, batch_size32): 创建训练和验证数据加载器。 Args: data_dir: 数据根目录包含 train 和 val 子目录。 batch_size: 批大小。 Returns: train_loader, val_loader, class_names # 数据预处理和增强管道 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]) # ImageNet均值标准差 ]), 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]) ]), } # 创建数据集 image_datasets { x: datasets.ImageFolder(rootf{data_dir}/{x}, transformdata_transforms[x]) for x in [train, val] } # 创建数据加载器 dataloaders { x: DataLoader(image_datasets[x], batch_sizebatch_size, shuffle(xtrain), num_workers4) for x in [train, val] } # 获取类别名称 class_names image_datasets[train].classes return dataloaders[train], dataloaders[val], class_names关键点说明数据增强对训练集使用RandomResizedCrop和RandomHorizontalFlip这是应对小数据集过拟合最基础且有效的手段。归一化参数Normalize使用的均值[0.485, 0.456, 0.406]和标准差[0.229, 0.224, 0.225]是ImageNet 数据集的统计值。因为我们的预训练模型是在 ImageNet 上训练的输入数据必须使用相同的归一化方式。num_workers设置大于0的数字可以加速数据加载但在 Windows 下有时会出错若出错可设为0。5.2 模型构建加载并修改 ResNet我们将加载预训练的 ResNet-18并替换其最后一层全连接层使其输出维度等于我们的自定义类别数。import torch.nn as nn from torchvision import models def initialize_model(num_classes, feature_extractTrue, use_pretrainedTrue): 初始化一个预训练的ResNet模型并替换其最后一层。 Args: num_classes: 目标数据集的类别数量。 feature_extract: 是否只训练最后一层特征提取模式。True表示冻结前面所有层。 use_pretrained: 是否使用预训练权重。 Returns: model_ft: 修改后的模型。 input_size: 模型要求的输入尺寸对于ResNet是224。 # 选择模型这里以 ResNet18 为例。你也可以尝试 ResNet34, ResNet50 等。 model_ft models.resnet18(pretraineduse_pretrained) # 设置输入尺寸 input_size 224 if feature_extract: # 冻结所有模型参数 for param in model_ft.parameters(): param.requires_grad False # 获取最后一层全连接层fc的输入特征数 num_ftrs model_ft.fc.in_features # 替换为一个新的全连接层输出维度为我们的类别数 model_ft.fc nn.Linear(num_ftrs, num_classes) return model_ft, input_size模式选择特征提取模式 (feature_extractTrue)冻结预训练模型的所有层只训练新添加的最后一层fc。这是最快速、最节省显存的方式适用于数据量非常少或与 ImageNet 任务相似度很高的场景。微调模式 (feature_extractFalse)解冻所有层或部分深层对整个模型进行训练。这需要更多的数据和计算资源但可能获得更好的性能尤其当你的任务与 ImageNet 差异较大时。初学者建议先从特征提取模式开始。5.3 训练循环与验证这是整个流程的核心我们将定义训练和验证的函数并监控损失和准确率。import torch.optim as optim from tqdm import tqdm import copy def train_model(model, dataloaders, criterion, optimizer, num_epochs25, devicecuda): 训练模型。 Args: model: 要训练的模型。 dataloaders: 包含train和val键的字典值为DataLoader。 criterion: 损失函数。 optimizer: 优化器。 num_epochs: 训练轮数。 device: 训练设备cuda 或 cpu。 Returns: model: 训练好的模型。 val_acc_history: 每个epoch的验证集准确率历史记录。 model model.to(device) best_model_wts copy.deepcopy(model.state_dict()) best_acc 0.0 val_acc_history [] 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 # 使用tqdm显示进度条 pbar tqdm(dataloaders[phase], descf{phase.capitalize()} Epoch {epoch}) for inputs, labels in pbar: 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) # 更新进度条信息 pbar.set_postfix({Loss: loss.item()}) epoch_loss running_loss / len(dataloaders[phase].dataset) epoch_acc running_corrects.double() / len(dataloaders[phase].dataset) 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 copy.deepcopy(model.state_dict()) if phase val: val_acc_history.append(epoch_acc.cpu().numpy()) print() print(fBest val Acc: {best_acc:.4f}) # 加载最佳模型权重 model.load_state_dict(best_model_wts) return model, val_acc_history5.4 主训练脚本整合现在我们将所有部分整合到一个主脚本train.py中。# src/train.py import argparse import torch import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler import matplotlib.pyplot as plt import time import os from utils import get_dataloaders, initialize_model def main(): parser argparse.ArgumentParser(descriptionPyTorch ResNet Transfer Learning) parser.add_argument(--data_dir, typestr, default./data, help数据集根目录) parser.add_argument(--model_name, typestr, defaultresnet18, help模型名称 (resnet18, resnet34, resnet50)) parser.add_argument(--num_classes, typeint, requiredTrue, help自定义类别数量) parser.add_argument(--batch_size, typeint, default32, help批大小) parser.add_argument(--num_epochs, typeint, default25, help训练轮数) parser.add_argument(--feature_extract, actionstore_true, help是否使用特征提取模式冻结预训练层) parser.add_argument(--lr, typefloat, default0.001, help学习率) args parser.parse_args() # 设置设备 device torch.device(cuda:0 if torch.cuda.is_available() else cpu) print(fUsing device: {device}) # 1. 初始化数据加载器 print(Initializing dataloaders...) train_loader, val_loader, class_names get_dataloaders(args.data_dir, args.batch_size) print(fClass names: {class_names}) print(fTrain dataset size: {len(train_loader.dataset)}) print(fVal dataset size: {len(val_loader.dataset)}) # 2. 初始化模型 print(Initializing model...) model_ft, input_size initialize_model(args.num_classes, args.feature_extract, use_pretrainedTrue) model_ft model_ft.to(device) # 3. 设置损失函数和优化器 criterion nn.CrossEntropyLoss() # 只优化那些 requires_gradTrue 的参数 params_to_update model_ft.parameters() if args.feature_extract: params_to_update [] for name, param in model_ft.named_parameters(): if param.requires_grad: params_to_update.append(param) optimizer_ft optim.Adam(params_to_update, lrargs.lr) # 每7个epoch将学习率衰减为原来的0.1倍 exp_lr_scheduler lr_scheduler.StepLR(optimizer_ft, step_size7, gamma0.1) # 4. 训练模型 print(Starting training...) since time.time() model_ft, val_acc_history train_model(model_ft, {train: train_loader, val: val_loader}, criterion, optimizer_ft, exp_lr_scheduler, num_epochsargs.num_epochs, devicedevice) time_elapsed time.time() - since print(fTraining complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s) # 5. 保存模型 os.makedirs(./models, exist_okTrue) model_save_path f./models/{args.model_name}_finetuned.pth torch.save(model_ft.state_dict(), model_save_path) print(fModel saved to {model_save_path}) # 6. 绘制验证准确率曲线 plt.plot(val_acc_history) plt.title(Validation Accuracy over Epochs) plt.xlabel(Epoch) plt.ylabel(Accuracy) os.makedirs(./logs, exist_okTrue) plt.savefig(./logs/val_acc_curve.png) plt.show() if __name__ __main__: main()如何启动训练在项目根目录下运行python src/train.py --data_dir ./data --num_classes 5 --batch_size 16 --num_epochs 20 --feature_extract --lr 0.001参数解释--data_dir: 你的数据集路径。--num_classes: 你的数据有多少个类别。--batch_size: 根据你的 GPU 显存调整。如果显存不足如小于 4GB可以设为 8 或 4。--num_epochs: 训练轮数。对于小数据集10-30 轮通常足够。--feature_extract: 使用此标志表示启用特征提取模式冻结预训练层。如果想微调所有层则去掉此标志。--lr: 学习率。特征提取模式下学习率可以稍大如 0.001微调所有层时建议更小如 0.0001。5.5 效果验证与推理训练完成后我们需要加载保存的模型并在新的图片或验证集上进行预测以验证模型效果。# src/predict.py import torch from torchvision import transforms from PIL import Image import matplotlib.pyplot as plt import numpy as np import argparse from utils import initialize_model # 复用之前的函数 def predict_image(model, image_path, class_names, device, transform): 对单张图片进行预测。 model.eval() image Image.open(image_path).convert(RGB) image transform(image).unsqueeze(0) # 增加batch维度 image image.to(device) with torch.no_grad(): outputs model(image) _, predicted torch.max(outputs, 1) prob torch.nn.functional.softmax(outputs, dim1)[0] * 100 return class_names[predicted.item()], prob[predicted.item()].item() def main(): parser argparse.ArgumentParser(descriptionPredict using a trained model) parser.add_argument(--model_path, typestr, requiredTrue, helpPath to the trained .pth model) parser.add_argument(--image_path, typestr, requiredTrue, helpPath to the image to predict) parser.add_argument(--num_classes, typeint, requiredTrue, helpNumber of classes) parser.add_argument(--class_names, typestr, nargs, helpList of class names (optional)) args parser.parse_args() device torch.device(cuda:0 if torch.cuda.is_available() else cpu) # 加载模型结构 model, _ initialize_model(args.num_classes, feature_extractFalse, use_pretrainedFalse) model.load_state_dict(torch.load(args.model_path, map_locationdevice)) model model.to(device) # 定义与验证集相同的数据转换 data_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]) ]) # 预测 predicted_class, confidence predict_image(model, args.image_path, args.class_names, device, data_transform) # 显示结果 img Image.open(args.image_path) plt.imshow(img) plt.title(fPredicted: {predicted_class} ({confidence:.2f}%)) plt.axis(off) plt.show() print(fPredicted class: {predicted_class} with confidence {confidence:.2f}%) if __name__ __main__: main()运行预测脚本python src/predict.py --model_path ./models/resnet18_finetuned.pth --image_path ./test_image.jpg --num_classes 5 --class_names cat dog bird fish rabbit6. 资源占用与性能观察在训练和推理过程中监控资源使用情况至关重要它能帮助你优化配置避免内存溢出。1. GPU 显存占用观察训练阶段显存占用主要取决于模型参数量、批大小 (Batch Size)和输入图像分辨率。ResNet-18参数量约 1100 万。Batch Size32输入 224x224训练时显存占用通常在1.5 GB ~ 2.5 GB之间。ResNet-50参数量约 2500 万。相同条件下显存占用可能达到3.5 GB ~ 5 GB。调整策略如果遇到CUDA out of memory错误首先尝试减小--batch_size如从 32 降到 16 或 8。其次可以尝试降低输入图像分辨率需同步修改模型初始化中的input_size和数据预处理中的Resize和Crop尺寸例如从 224 降到 112但这可能影响精度。推理阶段显存占用远小于训练因为不需要保存中间变量用于反向传播。通常几百 MB 到 1 GB 即可。如何监控在 Linux 终端可以使用watch -n 1 nvidia-smi命令每秒刷新一次 GPU 状态。在代码中可以使用torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()来跟踪显存使用。2. 训练速度CPU vs GPUGPU 训练速度通常是 CPU 的10 倍到 50 倍以上。对于 ResNet-18 在小型数据集上一个 epoch 在 GPU 上可能只需几秒到几十秒而在 CPU 上可能需要几分钟。数据加载瓶颈如果发现 GPU 利用率很低例如使用nvidia-smi看到 GPU-Util 经常为 0%可能是数据加载 (DataLoader) 太慢。可以尝试增加num_workers在 Linux 下或使用更快的存储如 SSD或预先将图片缓存到内存。3. 学习率与收敛观察损失曲线训练初期损失应快速下降。如果损失几乎不变可能是学习率太小如果损失剧烈震荡或变成 NaN可能是学习率太大。验证准确率这是衡量模型泛化能力的核心指标。理想情况是验证准确率随训练轮数稳步上升最终趋于平稳。如果验证准确率在达到一个峰值后开始下降而训练准确率持续上升这是典型的过拟合信号。此时应增加数据增强的强度如随机旋转、颜色抖动。添加正则化如 Dropout在自定义的全连接层后添加nn.Dropout(p0.5)。提前停止训练Early Stopping即保存验证准确率最高的模型。7. 常见问题与排查方法在实践过程中你可能会遇到以下问题。这里提供快速排查思路。问题现象可能原因排查方式解决方案RuntimeError: CUDA out of memory1. 批大小 (batch_size) 太大。2. 模型太大 (如 ResNet-152)。3. 输入图片分辨率太高。4. 有其他程序占用显存。1. 运行nvidia-smi查看显存占用。2. 在代码开始时打印torch.cuda.memory_allocated()。1.首要方案减小--batch_size。2. 换用更小的模型 (如 ResNet-18)。3. 降低输入图片尺寸。4. 关闭不必要的图形界面或程序。训练损失不下降或下降极慢1. 学习率 (lr) 设置过小。2. 优化器选择不当。3. 在特征提取模式下却冻结了所有层包括新加的fc层。1. 检查optimizer的参数列表是否包含需要更新的参数。2. 打印模型参数检查requires_grad属性。1. 适当增大学习率。2. 确保新添加的fc层的requires_gradTrue。3. 尝试使用optim.SGD并搭配动量。验证准确率远低于训练准确率且差距越来越大过拟合。模型记住了训练集的噪声而非通用特征。观察训练和验证的损失/准确率曲线。1.加强数据增强。2. 在fc层后添加Dropout。3. 使用权重衰减(Weight Decay)。4. 收集更多数据。5. 实施早停(Early Stopping)。KeyError: ‘val’或找不到数据集数据目录结构不正确。ImageFolder要求子目录名为train和val。检查--data_dir参数指定的路径下是否有train和val文件夹且其内部是按类别命名的子文件夹。严格按照data/train/class_x/和data/val/class_x/的结构组织图片。预测时类别标签错乱class_names的顺序与训练时DataLoader自动生成的顺序不一致。训练时将class_names image_datasets[‘train’].classes保存到一个文件 (如class_indices.json)。预测时加载该文件。在训练脚本中保存类别到索引的映射with open(‘class_indices.json’, ‘w’) as f:json.dump({v: k for k, v in image_datasets[‘train’].class_to_idx.items()}, f)ImportError: No module named ‘torch’PyTorch 未安装或未安装在当前 Conda 环境中。在终端输入python -c “import torch; print(torch.__version__)”。1. 确认已激活正确的 Conda 环境。2. 在该环境中重新安装 PyTorch。Windows 下DataLoader设置num_workers0报错Windows 系统下 Python 的多进程机制问题。将DataLoader的num_workers参数设为0。在get_dataloaders函数中设置num_workers0。8. 最佳实践与使用建议为了让你的迁移学习项目更稳健、高效遵循以下最佳实践从小开始快速迭代第一次运行时使用最小的配置ResNet-18feature_extractTruebatch_size8num_epochs5。目标是快速验证整个流程能否跑通。流程跑通后再逐步增加num_epochs尝试feature_extractFalse微调或换用更大的模型如ResNet-50。系统化管理实验每次实验不同的超参数组合都保存对应的模型文件、训练日志和验证曲线图。文件名可以包含关键参数如resnet18_fe_bs16_lr0.001_epoch20.pth。考虑使用TensorBoard或Weights Biases等工具来可视化和管理实验。数据是王道质量确保图片清晰、标注准确。模糊、错误的图片会严重干扰模型。增强对于小数据集数据增强是防止过拟合的必备手段。除了水平翻转还可以尝试随机旋转、色彩抖动、随机裁剪等torchvision.transforms提供了丰富选项。平衡尽量让每个类别的图片数量接近。如果无法做到可以使用WeightedRandomSampler在DataLoader中进行重采样。超参数调优学习率 (lr)这是最重要的超参数。特征提取模式下可用0.001微调模式下建议从0.0001开始尝试。使用学习率调度器如StepLR很有帮助。优化器Adam通常是个不错的默认选择。对于微调SGD with momentum如optim.SGD(..., lr0.001, momentum0.9)有时能获得更好的泛化性能。模型选择与扩展轻量级起点ResNet-18在速度和精度上取得了很好的平衡是首选的基准模型。追求精度如果显存和算力允许可以尝试ResNet-50、ResNet-101或EfficientNet、Vision Transformer等更先进的架构。自定义网络头除了简单替换最后一层你还可以在预训练模型后添加一个小型自定义网络如多个全连接层Dropout这有时能更好地适应特定任务。部署与集成训练完成后你可以使用torch.jit.trace或torch.jit.script将模型转换为 TorchScript以便在非 Python 环境中部署。也可以使用ONNX格式将模型导出集成到其他支持 ONNX 的推理引擎中。通过这套流程你可以在短短几小时内利用少量图片和一台普通配置的电脑甚至只有 CPU构建出一个可用的图像分类模型。这不仅是学习深度学习的高效路径也是在实际项目中快速验证 AI 解决方案可行性的强大工具。建议收藏本文在下次遇到小样本图像分类任务时直接按此流程实践。
