大家好我是专注于分享AI与深度学习实战经验的技术博主。在实际项目中我们常常会遇到一个经典难题想要训练一个图像分类模型但手头只有几十张、几百张图片远达不到深度学习模型通常所需的数万甚至数十万的数据量。直接训练不仅容易过拟合效果也往往不尽如人意。这时迁移学习Transfer Learning就成了我们的“救命稻草”。它允许我们利用在大规模数据集如ImageNet上预训练好的强大模型仅用少量数据对其进行微调Fine-tuning就能快速得到一个针对新任务的高性能分类器。本文将围绕“用少量图片进行图像分类”这一核心场景手把手带你完成从数据准备到模型微调的全流程实战。无论你是刚入门深度学习的新手还是希望快速将AI能力应用到具体业务中的开发者都能从本文中获得一套可直接复用的完整方案。我们将使用PyTorch框架以Kaggle上的“猫狗大战”数据集仅取少量样本为例详细拆解每一个步骤。1. 迁移学习与数据准备的核心概念在开始敲代码之前我们必须理解两个核心概念为什么迁移学习能解决小样本问题以及如何为小样本任务准备高质量的数据。1.1 迁移学习站在巨人的肩膀上想象一下一个已经读过千万本书、知识渊博的学者现在要学习一门新的小众语言。他不需要再从字母开始学起而是可以快速地将已有的语言知识语法结构、词根等迁移过来只需学习少量新词汇就能掌握这门新语言。迁移学习的思想与此类似。一个在ImageNet上预训练好的卷积神经网络如ResNet、VGG已经学会了提取通用图像特征的能力比如边缘、纹理、形状甚至部分物体部件。这些底层特征是跨任务通用的。我们的新任务比如区分猫和狗或者区分不同种类的工业零件与原始任务ImageNet的1000类物体分类在底层视觉特征上高度相关。因此迁移学习的典型做法是保留特征提取器冻结预训练模型的大部分层通常是前面的卷积层它们负责提取通用特征。替换并训练分类头替换掉模型最后的全连接层分类头使其输出维度等于我们的新类别数例如2类猫和狗。然后主要训练这个新的分类头以及靠近它的少数几层网络。微调有时为了达到更好效果我们会在训练后期解冻部分或全部卷积层用很小的学习率进行微调让模型的特征提取能力更适应我们的新数据。这样做的好处是巨大的我们只需要很少的新数据就能让模型快速适应新任务大大缩短训练时间并显著提升在小数据集上的性能。1.2 小样本学习下的数据准备质量重于数量当数据量很少时每一张图片都弥足珍贵。低质量或准备不当的数据会直接导致模型无法学习或学习到错误的模式。因此数据准备环节比在大数据场景下更为关键主要目标有两个扩充有效数据量和提升数据质量。数据收集与清洗确保收集的图片都正确对应其标签。对于小样本任何错误标注都是致命的。需要人工仔细检查剔除模糊、不相关或标注错误的图片。数据增强Data Augmentation这是小样本学习的核心技术。通过对原始图像进行一系列随机但合理的变换如旋转、翻转、裁剪、调整亮度对比度等人工创造出“新”的训练样本。这能有效增加数据的多样性防止模型过拟合到训练集有限的变换上提升模型的泛化能力。PyTorch的torchvision.transforms模块提供了丰富的增强方法。数据集划分即使总数据量很少也必须严格划分训练集、验证集和测试集。通常采用8:1:1或7:2:1的比例。验证集用于在训练过程中监控模型表现、调整超参数和进行早停Early Stopping防止过拟合测试集用于最终评估模型的泛化能力在整个训练过程中绝对不能使用。数据标准化Normalization将图像的像素值从 [0, 255] 缩放到一个固定的范围通常是 [-1, 1] 或 [0, 1]并进行减均值、除标准差的操作。这能加速模型收敛提升训练稳定性。关键点是必须使用预训练模型训练时所用的均值和标准差例如ImageNet的mean[0.485, 0.456, 0.406],std[0.229, 0.224, 0.225]而不是我们自己数据的统计量。理解了这些我们就可以开始动手搭建环境了。2. 环境准备与工具说明为了确保代码可复现以下是本次实战所需的环境和工具。建议使用Anaconda管理Python环境。2.1 基础环境操作系统Windows 10/11, macOS 或 Linux (Ubuntu 20.04)。本文指令以Linux/macOS的bash和Windows的PowerShell为例。Python3.8 或 3.9 (推荐)。避免使用过新或过旧的版本。深度学习框架PyTorch 1.12。我们将使用PyTorch及其生态工具torchvision。CUDA/cuDNN如果你有NVIDIA GPU并希望加速训练需要安装对应版本的CUDA和cuDNN。本文代码同时支持CPU和GPU运行。2.2 创建虚拟环境与安装依赖强烈建议创建独立的虚拟环境避免包版本冲突。# 1. 创建并激活虚拟环境 (以 conda 为例) conda create -n pytorch_tl python3.9 -y conda activate pytorch_tl # 2. 安装 PyTorch 和 torchvision # 请根据你的CUDA版本前往 https://pytorch.org/get-started/locally/ 获取最准确的安装命令。 # 例如对于CUDA 11.8 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 如果只用CPU # pip install torch torchvision torchaudio # 3. 安装其他必要的库 pip install numpy pandas matplotlib jupyter notebook opencv-python pillow scikit-learn tqdm # opencv-python 用于可能的图像处理pillow是PIL的友好分支scikit-learn用于评估tqdm用于进度条2.3 项目结构规划在开始前规划一个清晰的项目结构有助于管理代码和数据。ai_study_transfer_learning/ │ ├── data/ # 存放所有数据 │ ├── raw/ # 原始下载的数据 │ ├── processed/ # 处理后的数据划分后的文件夹 │ │ ├── train/ │ │ │ ├── cat/ │ │ │ └── dog/ │ │ ├── val/ │ │ │ ├── cat/ │ │ │ └── dog/ │ │ └── test/ │ │ ├── cat/ │ │ └── dog/ │ └── splits/ # 存放划分好的文件列表可选 │ ├── notebooks/ # Jupyter notebook 用于探索性分析 │ └── 01_data_exploration.ipynb │ ├── src/ # 源代码 │ ├── data_preparation.py # 数据准备脚本 │ ├── dataset.py # 自定义Dataset类 │ ├── train.py # 训练脚本 │ ├── utils.py # 工具函数 │ └── config.py # 配置文件超参数等 │ ├── models/ # 保存训练好的模型 ├── logs/ # 训练日志TensorBoard文件 ├── requirements.txt # 项目依赖 └── README.md接下来我们将进入最核心的环节数据准备。3. 实战少量图片的数据准备全流程我们将以经典的“猫狗分类”任务为例。假设我们从Kaggle的“Dogs vs Cats”数据集中每类只随机抽取150张图片作为训练集20张作为验证集30张作为测试集。这是一个典型的小样本场景。3.1 步骤一获取与组织原始数据首先你需要下载数据。可以从Kaggle官网下载dogs-vs-cats数据集解压后你会得到一个train文件夹里面包含约25000张命名为cat.0.jpg,dog.0.jpg的图片。我们的目标是从中抽取少量样本并按照上述项目结构进行组织。我们编写一个Python脚本src/data_preparation.py来完成这个工作。# src/data_preparation.py import os import shutil import random from pathlib import Path import argparse def split_dataset(raw_data_dir, output_dir, train_samples_per_class150, val_samples_per_class20, test_samples_per_class30, seed42): 从原始数据集中随机抽取样本并划分为train/val/test集。 参数: raw_data_dir: 原始数据目录例如包含 cat.0.jpg, dog.1.jpg 的文件夹。 output_dir: 处理后的输出根目录。 train_samples_per_class: 每个类别在训练集中的样本数。 val_samples_per_class: 每个类别在验证集中的样本数。 test_samples_per_class: 每个类别在测试集中的样本数。 seed: 随机种子确保结果可复现。 random.seed(seed) # 定义类别根据文件名前缀 classes [cat, dog] # 创建输出目录结构 splits [train, val, test] for split in splits: for cls in classes: Path(os.path.join(output_dir, split, cls)).mkdir(parentsTrue, exist_okTrue) # 遍历每个类别进行抽样和分配 for cls in classes: # 获取该类所有图片文件 # 假设文件名格式为 cls.number.jpg pattern f{cls}.*.jpg import glob all_images glob.glob(os.path.join(raw_data_dir, pattern)) # 随机打乱 random.shuffle(all_images) # 计算划分点 total_needed train_samples_per_class val_samples_per_class test_samples_per_class if len(all_images) total_needed: print(f警告: {cls} 类只有 {len(all_images)} 张图片少于需要的 {total_needed} 张。将使用所有可用图片。) all_images all_images[:total_needed] # 或者采取其他策略如复制 train_end train_samples_per_class val_end train_end val_samples_per_class split_ranges { train: (0, train_end), val: (train_end, val_end), test: (val_end, val_end test_samples_per_class) } # 复制文件到对应目录 for split_name, (start, end) in split_ranges.items(): for img_path in all_images[start:end]: filename os.path.basename(img_path) dest_path os.path.join(output_dir, split_name, cls, filename) shutil.copy2(img_path, dest_path) print(f{cls} 类 - {split_name} 集: 复制了 {end-start} 张图片。) print(数据集划分完成) # 打印统计信息 for split in splits: total 0 for cls in classes: count len(os.listdir(os.path.join(output_dir, split, cls))) total count print(f {split}/{cls}: {count}) print(f {split} 集总计: {total}) if __name__ __main__: parser argparse.ArgumentParser(description准备小样本猫狗数据集) parser.add_argument(--raw_dir, typestr, default../data/raw/train, help原始训练图片目录) parser.add_argument(--output_dir, typestr, default../data/processed, help处理后的输出目录) parser.add_argument(--train_num, typeint, default150, help每类训练样本数) parser.add_argument(--val_num, typeint, default20, help每类验证样本数) parser.add_argument(--test_num, typeint, default30, help每类测试样本数) args parser.parse_args() split_dataset(args.raw_dir, args.output_dir, args.train_num, args.val_num, args.test_num)运行这个脚本cd src python data_preparation.py --raw_dir ../../data/raw/train --output_dir ../../data/processed执行后你的data/processed/目录下就会生成结构清晰的train,val,test文件夹。3.2 步骤二创建PyTorch Dataset与DataLoader数据组织好后我们需要用PyTorch的Dataset和DataLoader来加载它们并集成数据增强和标准化。首先我们创建一个自定义的Dataset类它能够从train/class_name/*.jpg这样的文件夹结构中读取数据。# src/dataset.py import torch from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image import os class CatDogDataset(Dataset): 猫狗分类数据集. def __init__(self, root_dir, transformNone): 参数: root_dir: 数据集根目录例如 ‘data/processed/train‘。 transform: 应用于图像的变换数据增强标准化。 self.root_dir root_dir self.transform transform # 获取所有图像路径和对应的标签 self.image_paths [] self.labels [] self.class_to_idx {cat: 0, dog: 1} # 固定类别映射 self.idx_to_class {v: k for k, v in self.class_to_idx.items()} for class_name, label in self.class_to_idx.items(): class_dir os.path.join(root_dir, class_name) if not os.path.isdir(class_dir): continue for img_name in os.listdir(class_dir): if img_name.lower().endswith((.png, .jpg, .jpeg)): self.image_paths.append(os.path.join(class_dir, img_name)) self.labels.append(label) def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img_path self.image_paths[idx] label self.labels[idx] # 使用PIL打开图像确保是RGB三通道 image Image.open(img_path).convert(RGB) if self.transform: image self.transform(image) return image, label def get_class_names(self): return list(self.class_to_idx.keys())接下来在配置文件中定义针对训练集和验证/测试集的不同数据变换。这是小样本学习提升性能的关键。# src/config.py import torchvision.transforms as 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_test_transform transforms.Compose([ transforms.Resize(256), # 将短边缩放到256 transforms.CenterCrop(224), # 中心裁剪到224x224 transforms.ToTensor(), transforms.Normalize(meanIMAGENET_MEAN, stdIMAGENET_STD) ])最后创建DataLoader它负责批量加载数据、打乱顺序训练集等。# 在训练脚本或另一个工具文件中使用 from src.dataset import CatDogDataset from src.config import train_transform, val_test_transform from torch.utils.data import DataLoader # 创建数据集实例 train_dataset CatDogDataset(root_dir../data/processed/train, transformtrain_transform) val_dataset CatDogDataset(root_dir../data/processed/val, transformval_test_transform) test_dataset CatDogDataset(root_dir../data/processed/test, transformval_test_transform) # 创建数据加载器 batch_size 32 # 根据你的GPU内存调整小样本可以小一些 train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workers2, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_sizebatch_size, shuffleFalse, num_workers2, pin_memoryTrue) test_loader DataLoader(test_dataset, batch_sizebatch_size, shuffleFalse, num_workers2)关键点解释RandomResizedCrop和RandomHorizontalFlip是最常用且有效的增强方式。ColorJitter通过轻微改变颜色来模拟光照变化进一步提升模型鲁棒性。验证集和测试集绝对不能使用任何随机性增强否则评估指标将不可靠。Normalize的参数必须使用ImageNet的统计量这是为了与预训练模型的特征分布对齐。num_workers可以加速数据加载pin_memory在GPU训练时能提升数据从CPU到GPU的传输速度。3.3 步骤三可视化与检查数据在投入训练前务必检查数据是否被正确加载和增强。# 可以在 notebook 中运行或写一个简单的检查脚本 import matplotlib.pyplot as plt import numpy as np import torch def imshow(inp, titleNone): 显示一个Tensor图像。 inp inp.numpy().transpose((1, 2, 0)) # 从(C, H, W)转为(H, W, C) # 反标准化 mean np.array(IMAGENET_MEAN) std np.array(IMAGENET_STD) inp std * inp mean inp np.clip(inp, 0, 1) plt.imshow(inp) if title is not None: plt.title(title) plt.axis(off) # 获取一个批次的数据 images, labels next(iter(train_loader)) # 创建一个网格来显示 out torchvision.utils.make_grid(images[:8]) # 显示前8张 imshow(out, title[train_dataset.idx_to_class[l.item()] for l in labels[:8]]) plt.show()运行这段代码你应该能看到一个由8张经过随机裁剪、翻转、颜色变换的猫狗图片组成的网格并且图片已经恢复了可视化的颜色。这确认了数据管道工作正常。4. 构建迁移学习模型数据准备就绪后我们就可以构建模型了。这里以使用torchvision.models中的resnet18预训练模型为例。# src/model.py import torch.nn as nn import torchvision.models as models def get_model(num_classes2, pretrainedTrue, freeze_backboneTrue): 获取一个基于ResNet18的迁移学习模型。 参数: num_classes: 输出类别数。 pretrained: 是否加载在ImageNet上预训练的权重。 freeze_backbone: 是否冻结卷积层特征提取器的参数。 # 加载预训练的ResNet18模型 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None) if freeze_backbone: # 冻结所有卷积层的参数使其在训练中不更新 for param in model.parameters(): param.requires_grad False # 替换最后的全连接层 # ResNet18的fc层输入特征数是512 num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes) # 如果冻结了主干只让新加的fc层可训练 # 如果未冻结则所有层都可训练微调模式 # 注意即使冻结了主干BatchNorm层的running mean/var在训练时依然会更新eval模式除外 return model # 使用示例 model get_model(num_classes2, pretrainedTrue, freeze_backboneTrue) print(model) # 可以打印参数查看哪些被冻结 # for name, param in model.named_parameters(): # print(f{name}: requires_grad{param.requires_grad})模型策略选择策略一特征提取器freeze_backboneTrue。仅训练新添加的全连接层。训练速度快所需数据量极少适合数据量非常少每类100或计算资源有限的情况。策略二微调freeze_backboneFalse。解冻所有层用很小的学习率进行训练。通常能获得比策略一更好的效果但需要更多的数据每类几百张和更谨慎的训练防止过拟合。策略三分层微调先冻结大部分层训练分类头训练几轮后再解冻靠近顶部的几层卷积层进行微调。这是一种折中方案。5. 训练与验证循环有了数据和模型我们就可以编写训练脚本了。这里提供一个完整的训练循环示例包含验证和模型保存。# src/train.py import torch import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler import time import copy from tqdm import tqdm import os from src.model import get_model from src.dataset import CatDogDataset from src.config import train_transform, val_test_transform from torch.utils.data import DataLoader def train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochs25, devicecuda): 训练模型的通用函数。 since time.time() best_model_wts copy.deepcopy(model.state_dict()) best_acc 0.0 # 记录训练过程 history {train_loss: [], train_acc: [], val_loss: [], val_acc: []} 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显示进度条 for inputs, labels in tqdm(dataloaders[phase], descf{phase.capitalize()} Epoch {epoch}): 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 and scheduler is not None: scheduler.step() epoch_loss running_loss / len(dataloaders[phase].dataset) epoch_acc running_corrects.double() / len(dataloaders[phase].dataset) # 记录历史 if phase train: history[train_loss].append(epoch_loss) history[train_acc].append(epoch_acc.cpu().item()) else: history[val_loss].append(epoch_loss) history[val_acc].append(epoch_acc.cpu().item()) 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()) # 保存最佳模型 torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict() if scheduler else None, best_acc: best_acc, history: history, }, ../models/best_model.pth) print(f 保存新的最佳模型准确率: {best_acc:.4f}) print() time_elapsed time.time() - since print(f训练完成用时 {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s) print(f最佳验证准确率: {best_acc:.4f}) # 加载最佳模型权重 model.load_state_dict(best_model_wts) return model, history def main(): # 超参数配置 data_dir ../data/processed batch_size 32 num_epochs 15 # 对于小样本epoch不宜过多防止过拟合 learning_rate 0.001 # 初始学习率 momentum 0.9 step_size 7 # 学习率衰减步长 gamma 0.1 # 学习率衰减因子 # 设备选择 device torch.device(cuda if torch.cuda.is_available() else cpu) print(f使用设备: {device}) # 创建数据加载器字典 image_datasets { train: CatDogDataset(os.path.join(data_dir, train), transformtrain_transform), val: CatDogDataset(os.path.join(data_dir, val), transformval_test_transform), } dataloaders { x: DataLoader(image_datasets[x], batch_sizebatch_size, shuffle(x train), num_workers4, pin_memoryTrue) for x in [train, val] } dataset_sizes {x: len(image_datasets[x]) for x in [train, val]} print(f数据集大小: {dataset_sizes}) # 初始化模型 model get_model(num_classes2, pretrainedTrue, freeze_backboneTrue) # 第一阶段冻结特征提取器 model model.to(device) # 定义损失函数和优化器 criterion nn.CrossEntropyLoss() # 只优化那些 requires_gradTrue 的参数即我们新加的fc层 optimizer optim.SGD(model.parameters(), lrlearning_rate, momentummomentum) # 学习率调度器 scheduler lr_scheduler.StepLR(optimizer, step_sizestep_size, gammagamma) # 训练模型 model, history train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochsnum_epochs, devicedevice) # 可以在这里添加第二阶段微调解冻部分层的代码 print(第一阶段训练完成。) # ... (后续可添加微调代码) if __name__ __main__: main()运行此脚本开始训练。由于数据量小且只训练分类头训练速度会非常快。你应该能看到训练损失下降验证准确率逐步提升。6. 模型评估与测试训练完成后我们需要在从未参与训练的测试集上评估模型的最终泛化性能。# src/evaluate.py import torch from src.model import get_model from src.dataset import CatDogDataset from src.config import val_test_transform from torch.utils.data import DataLoader import numpy as np from sklearn.metrics import classification_report, confusion_matrix import seaborn as sns import matplotlib.pyplot as plt def evaluate_model(model_path, test_data_dir, devicecuda): 在测试集上评估模型。 # 加载模型 checkpoint torch.load(model_path, map_locationdevice) model get_model(num_classes2, pretrainedFalse, freeze_backboneFalse) # 加载时不需要预训练权重 model.load_state_dict(checkpoint[model_state_dict]) model model.to(device) model.eval() # 非常重要设置为评估模式 # 创建测试集加载器 test_dataset CatDogDataset(root_dirtest_data_dir, transformval_test_transform) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse, num_workers2) all_preds [] all_labels [] 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) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) # 计算准确率 accuracy np.sum(np.array(all_preds) np.array(all_labels)) / len(all_labels) print(f测试集准确率: {accuracy:.4f}) # 打印分类报告精确率、召回率、F1分数 print(\n分类报告:) print(classification_report(all_labels, all_preds, target_namestest_dataset.get_class_names())) # 绘制混淆矩阵 cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(6,5)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelstest_dataset.get_class_names(), yticklabelstest_dataset.get_class_names()) plt.ylabel(真实标签) plt.xlabel(预测标签) plt.title(混淆矩阵) plt.tight_layout() plt.savefig(../results/confusion_matrix.png) plt.show() return accuracy if __name__ __main__: evaluate_model(../models/best_model.pth, ../data/processed/test, devicecuda if torch.cuda.is_available() else cpu)运行评估脚本你将得到模型在测试集上的准确率、精确率、召回率和F1分数以及一个可视化的混淆矩阵。对于我们的猫狗二分类小样本任务使用迁移学习通常能达到95%以上的测试准确率这充分展示了迁移学习的威力。7. 常见问题与排查思路在实际操作中你可能会遇到以下问题问题现象可能原因解决思路训练损失不下降准确率始终在50%左右二分类1. 数据加载错误标签与图片不匹配。2. 模型输出层维度错误。3. 学习率设置不当太大或太小。4. 预训练模型权重未正确加载。1. 使用imshow函数可视化几个批次的数据和标签确认对应关系。2. 检查model.fc.out_features是否等于类别数。3. 尝试调整学习率如0.01, 0.001, 0.0001。4. 打印模型参数确认requires_grad状态并检查预训练权重是否加载成功。验证准确率远低于训练准确率过拟合1. 训练数据太少。2. 数据增强不够强或未启用。3. 模型过于复杂或训练轮数太多。4. 没有使用验证集进行早停。1. 尝试收集更多数据或使用更激进的数据增强如RandomAffine,RandomPerspective。2. 增加数据增强的强度或种类。3. 减少训练轮数 (num_epochs)或使用更简单的模型如resnet18而非resnet50。4. 实现早停机制当验证损失连续几个epoch不下降时停止训练。GPU内存不足CUDA out of memory1. 批次大小 (batch_size) 太大。2. 输入图片尺寸太大。3. 模型太大。1. 减小batch_size如从32减到16或8。2. 减小RandomResizedCrop和CenterCrop的尺寸如从224到128。3. 换用更小的预训练模型如resnet18代替resnet50。4. 使用torch.cuda.empty_cache()清理缓存。训练速度很慢1. 未使用GPU。2.num_workers设置过小默认为0。3. 数据增强过于复杂。1. 检查torch.cuda.is_available()并确保model.to(device)和data.to(device)。2. 根据CPU核心数增加num_workers通常设置为CPU核心数。3. 简化数据增强流程或在训练初期使用较弱增强。加载保存的模型时报错1. 模型结构定义与保存时不一致。2. PyTorch版本不兼容。3. 加载路径错误。1. 确保加载模型时使用的get_model函数参数与保存时完全一致特别是num_classes。2. 尽量在相同版本的PyTorch环境下进行保存和加载。3. 使用绝对路径并检查文件是否存在。8. 最佳实践与进阶建议掌握了基础流程后以下建议能帮助你在实际项目中做得更好数据永远是王道即使使用迁移学习高质量、标注准确的数据也是成功的基石。对于小样本可以尝试主动学习Active Learning策略即让模型筛选出最不确定的样本交由人工标注用最少的标注成本获得最大性能提升。分层学习率与差分学习率在微调时不同层应该使用不同的学习率。靠近输入的底层卷积层学习通用特征应使用较小的学习率靠近输出的高层层和分类头应使用较大的学习率。可以使用torch.optim中的param_groups来实现。使用更强大的预训练模型ResNet只是起点。可以尝试EfficientNet,Vision Transformer (ViT),Swin Transformer等更现代的架构它们在ImageNet上的表现更好迁移能力可能更强。交叉验证对于极少量数据如每类少于50张可以考虑使用K折交叉验证来更可靠地评估模型性能并充分利用有限的数据。利用外部数据如果领域内公开数据极少可以考虑使用领域自适应Domain Adaptation技术利用大量相关的公开数据源域来辅助训练再迁移到你的小数据目标域。模型集成训练多个不同的模型如使用不同的预训练网络、不同的数据增强策略、不同的随机种子然后将它们的预测结果进行平均或投票通常能稳定地提升最终性能。自动化与实验管理使用Weights Biases (wandb)、TensorBoard或MLflow等工具记录每次实验的超参数、指标和模型便于分析和复现。通过本文的详细拆解你应该已经掌握了使用迁移学习应对小样本图像分类任务的核心流程从高质量的数据准备清洗、划分、增强、标准化到模型构建与策略选择冻结/微调再到完整的训练、验证与评估循环。这套方法不仅适用于猫狗分类可以轻松迁移到任何新的图像分类任务上例如医学影像分析、工业质检、动植物识别等。关键在于理解每个步骤背后的原理并根据自己的数据和任务特点进行灵活调整。 小样本图像分类实战：基于PyTorch的迁移学习与数据增强应用
大家好我是专注于分享AI与深度学习实战经验的技术博主。在实际项目中我们常常会遇到一个经典难题想要训练一个图像分类模型但手头只有几十张、几百张图片远达不到深度学习模型通常所需的数万甚至数十万的数据量。直接训练不仅容易过拟合效果也往往不尽如人意。这时迁移学习Transfer Learning就成了我们的“救命稻草”。它允许我们利用在大规模数据集如ImageNet上预训练好的强大模型仅用少量数据对其进行微调Fine-tuning就能快速得到一个针对新任务的高性能分类器。本文将围绕“用少量图片进行图像分类”这一核心场景手把手带你完成从数据准备到模型微调的全流程实战。无论你是刚入门深度学习的新手还是希望快速将AI能力应用到具体业务中的开发者都能从本文中获得一套可直接复用的完整方案。我们将使用PyTorch框架以Kaggle上的“猫狗大战”数据集仅取少量样本为例详细拆解每一个步骤。1. 迁移学习与数据准备的核心概念在开始敲代码之前我们必须理解两个核心概念为什么迁移学习能解决小样本问题以及如何为小样本任务准备高质量的数据。1.1 迁移学习站在巨人的肩膀上想象一下一个已经读过千万本书、知识渊博的学者现在要学习一门新的小众语言。他不需要再从字母开始学起而是可以快速地将已有的语言知识语法结构、词根等迁移过来只需学习少量新词汇就能掌握这门新语言。迁移学习的思想与此类似。一个在ImageNet上预训练好的卷积神经网络如ResNet、VGG已经学会了提取通用图像特征的能力比如边缘、纹理、形状甚至部分物体部件。这些底层特征是跨任务通用的。我们的新任务比如区分猫和狗或者区分不同种类的工业零件与原始任务ImageNet的1000类物体分类在底层视觉特征上高度相关。因此迁移学习的典型做法是保留特征提取器冻结预训练模型的大部分层通常是前面的卷积层它们负责提取通用特征。替换并训练分类头替换掉模型最后的全连接层分类头使其输出维度等于我们的新类别数例如2类猫和狗。然后主要训练这个新的分类头以及靠近它的少数几层网络。微调有时为了达到更好效果我们会在训练后期解冻部分或全部卷积层用很小的学习率进行微调让模型的特征提取能力更适应我们的新数据。这样做的好处是巨大的我们只需要很少的新数据就能让模型快速适应新任务大大缩短训练时间并显著提升在小数据集上的性能。1.2 小样本学习下的数据准备质量重于数量当数据量很少时每一张图片都弥足珍贵。低质量或准备不当的数据会直接导致模型无法学习或学习到错误的模式。因此数据准备环节比在大数据场景下更为关键主要目标有两个扩充有效数据量和提升数据质量。数据收集与清洗确保收集的图片都正确对应其标签。对于小样本任何错误标注都是致命的。需要人工仔细检查剔除模糊、不相关或标注错误的图片。数据增强Data Augmentation这是小样本学习的核心技术。通过对原始图像进行一系列随机但合理的变换如旋转、翻转、裁剪、调整亮度对比度等人工创造出“新”的训练样本。这能有效增加数据的多样性防止模型过拟合到训练集有限的变换上提升模型的泛化能力。PyTorch的torchvision.transforms模块提供了丰富的增强方法。数据集划分即使总数据量很少也必须严格划分训练集、验证集和测试集。通常采用8:1:1或7:2:1的比例。验证集用于在训练过程中监控模型表现、调整超参数和进行早停Early Stopping防止过拟合测试集用于最终评估模型的泛化能力在整个训练过程中绝对不能使用。数据标准化Normalization将图像的像素值从 [0, 255] 缩放到一个固定的范围通常是 [-1, 1] 或 [0, 1]并进行减均值、除标准差的操作。这能加速模型收敛提升训练稳定性。关键点是必须使用预训练模型训练时所用的均值和标准差例如ImageNet的mean[0.485, 0.456, 0.406],std[0.229, 0.224, 0.225]而不是我们自己数据的统计量。理解了这些我们就可以开始动手搭建环境了。2. 环境准备与工具说明为了确保代码可复现以下是本次实战所需的环境和工具。建议使用Anaconda管理Python环境。2.1 基础环境操作系统Windows 10/11, macOS 或 Linux (Ubuntu 20.04)。本文指令以Linux/macOS的bash和Windows的PowerShell为例。Python3.8 或 3.9 (推荐)。避免使用过新或过旧的版本。深度学习框架PyTorch 1.12。我们将使用PyTorch及其生态工具torchvision。CUDA/cuDNN如果你有NVIDIA GPU并希望加速训练需要安装对应版本的CUDA和cuDNN。本文代码同时支持CPU和GPU运行。2.2 创建虚拟环境与安装依赖强烈建议创建独立的虚拟环境避免包版本冲突。# 1. 创建并激活虚拟环境 (以 conda 为例) conda create -n pytorch_tl python3.9 -y conda activate pytorch_tl # 2. 安装 PyTorch 和 torchvision # 请根据你的CUDA版本前往 https://pytorch.org/get-started/locally/ 获取最准确的安装命令。 # 例如对于CUDA 11.8 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 如果只用CPU # pip install torch torchvision torchaudio # 3. 安装其他必要的库 pip install numpy pandas matplotlib jupyter notebook opencv-python pillow scikit-learn tqdm # opencv-python 用于可能的图像处理pillow是PIL的友好分支scikit-learn用于评估tqdm用于进度条2.3 项目结构规划在开始前规划一个清晰的项目结构有助于管理代码和数据。ai_study_transfer_learning/ │ ├── data/ # 存放所有数据 │ ├── raw/ # 原始下载的数据 │ ├── processed/ # 处理后的数据划分后的文件夹 │ │ ├── train/ │ │ │ ├── cat/ │ │ │ └── dog/ │ │ ├── val/ │ │ │ ├── cat/ │ │ │ └── dog/ │ │ └── test/ │ │ ├── cat/ │ │ └── dog/ │ └── splits/ # 存放划分好的文件列表可选 │ ├── notebooks/ # Jupyter notebook 用于探索性分析 │ └── 01_data_exploration.ipynb │ ├── src/ # 源代码 │ ├── data_preparation.py # 数据准备脚本 │ ├── dataset.py # 自定义Dataset类 │ ├── train.py # 训练脚本 │ ├── utils.py # 工具函数 │ └── config.py # 配置文件超参数等 │ ├── models/ # 保存训练好的模型 ├── logs/ # 训练日志TensorBoard文件 ├── requirements.txt # 项目依赖 └── README.md接下来我们将进入最核心的环节数据准备。3. 实战少量图片的数据准备全流程我们将以经典的“猫狗分类”任务为例。假设我们从Kaggle的“Dogs vs Cats”数据集中每类只随机抽取150张图片作为训练集20张作为验证集30张作为测试集。这是一个典型的小样本场景。3.1 步骤一获取与组织原始数据首先你需要下载数据。可以从Kaggle官网下载dogs-vs-cats数据集解压后你会得到一个train文件夹里面包含约25000张命名为cat.0.jpg,dog.0.jpg的图片。我们的目标是从中抽取少量样本并按照上述项目结构进行组织。我们编写一个Python脚本src/data_preparation.py来完成这个工作。# src/data_preparation.py import os import shutil import random from pathlib import Path import argparse def split_dataset(raw_data_dir, output_dir, train_samples_per_class150, val_samples_per_class20, test_samples_per_class30, seed42): 从原始数据集中随机抽取样本并划分为train/val/test集。 参数: raw_data_dir: 原始数据目录例如包含 cat.0.jpg, dog.1.jpg 的文件夹。 output_dir: 处理后的输出根目录。 train_samples_per_class: 每个类别在训练集中的样本数。 val_samples_per_class: 每个类别在验证集中的样本数。 test_samples_per_class: 每个类别在测试集中的样本数。 seed: 随机种子确保结果可复现。 random.seed(seed) # 定义类别根据文件名前缀 classes [cat, dog] # 创建输出目录结构 splits [train, val, test] for split in splits: for cls in classes: Path(os.path.join(output_dir, split, cls)).mkdir(parentsTrue, exist_okTrue) # 遍历每个类别进行抽样和分配 for cls in classes: # 获取该类所有图片文件 # 假设文件名格式为 cls.number.jpg pattern f{cls}.*.jpg import glob all_images glob.glob(os.path.join(raw_data_dir, pattern)) # 随机打乱 random.shuffle(all_images) # 计算划分点 total_needed train_samples_per_class val_samples_per_class test_samples_per_class if len(all_images) total_needed: print(f警告: {cls} 类只有 {len(all_images)} 张图片少于需要的 {total_needed} 张。将使用所有可用图片。) all_images all_images[:total_needed] # 或者采取其他策略如复制 train_end train_samples_per_class val_end train_end val_samples_per_class split_ranges { train: (0, train_end), val: (train_end, val_end), test: (val_end, val_end test_samples_per_class) } # 复制文件到对应目录 for split_name, (start, end) in split_ranges.items(): for img_path in all_images[start:end]: filename os.path.basename(img_path) dest_path os.path.join(output_dir, split_name, cls, filename) shutil.copy2(img_path, dest_path) print(f{cls} 类 - {split_name} 集: 复制了 {end-start} 张图片。) print(数据集划分完成) # 打印统计信息 for split in splits: total 0 for cls in classes: count len(os.listdir(os.path.join(output_dir, split, cls))) total count print(f {split}/{cls}: {count}) print(f {split} 集总计: {total}) if __name__ __main__: parser argparse.ArgumentParser(description准备小样本猫狗数据集) parser.add_argument(--raw_dir, typestr, default../data/raw/train, help原始训练图片目录) parser.add_argument(--output_dir, typestr, default../data/processed, help处理后的输出目录) parser.add_argument(--train_num, typeint, default150, help每类训练样本数) parser.add_argument(--val_num, typeint, default20, help每类验证样本数) parser.add_argument(--test_num, typeint, default30, help每类测试样本数) args parser.parse_args() split_dataset(args.raw_dir, args.output_dir, args.train_num, args.val_num, args.test_num)运行这个脚本cd src python data_preparation.py --raw_dir ../../data/raw/train --output_dir ../../data/processed执行后你的data/processed/目录下就会生成结构清晰的train,val,test文件夹。3.2 步骤二创建PyTorch Dataset与DataLoader数据组织好后我们需要用PyTorch的Dataset和DataLoader来加载它们并集成数据增强和标准化。首先我们创建一个自定义的Dataset类它能够从train/class_name/*.jpg这样的文件夹结构中读取数据。# src/dataset.py import torch from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image import os class CatDogDataset(Dataset): 猫狗分类数据集. def __init__(self, root_dir, transformNone): 参数: root_dir: 数据集根目录例如 ‘data/processed/train‘。 transform: 应用于图像的变换数据增强标准化。 self.root_dir root_dir self.transform transform # 获取所有图像路径和对应的标签 self.image_paths [] self.labels [] self.class_to_idx {cat: 0, dog: 1} # 固定类别映射 self.idx_to_class {v: k for k, v in self.class_to_idx.items()} for class_name, label in self.class_to_idx.items(): class_dir os.path.join(root_dir, class_name) if not os.path.isdir(class_dir): continue for img_name in os.listdir(class_dir): if img_name.lower().endswith((.png, .jpg, .jpeg)): self.image_paths.append(os.path.join(class_dir, img_name)) self.labels.append(label) def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img_path self.image_paths[idx] label self.labels[idx] # 使用PIL打开图像确保是RGB三通道 image Image.open(img_path).convert(RGB) if self.transform: image self.transform(image) return image, label def get_class_names(self): return list(self.class_to_idx.keys())接下来在配置文件中定义针对训练集和验证/测试集的不同数据变换。这是小样本学习提升性能的关键。# src/config.py import torchvision.transforms as 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_test_transform transforms.Compose([ transforms.Resize(256), # 将短边缩放到256 transforms.CenterCrop(224), # 中心裁剪到224x224 transforms.ToTensor(), transforms.Normalize(meanIMAGENET_MEAN, stdIMAGENET_STD) ])最后创建DataLoader它负责批量加载数据、打乱顺序训练集等。# 在训练脚本或另一个工具文件中使用 from src.dataset import CatDogDataset from src.config import train_transform, val_test_transform from torch.utils.data import DataLoader # 创建数据集实例 train_dataset CatDogDataset(root_dir../data/processed/train, transformtrain_transform) val_dataset CatDogDataset(root_dir../data/processed/val, transformval_test_transform) test_dataset CatDogDataset(root_dir../data/processed/test, transformval_test_transform) # 创建数据加载器 batch_size 32 # 根据你的GPU内存调整小样本可以小一些 train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workers2, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_sizebatch_size, shuffleFalse, num_workers2, pin_memoryTrue) test_loader DataLoader(test_dataset, batch_sizebatch_size, shuffleFalse, num_workers2)关键点解释RandomResizedCrop和RandomHorizontalFlip是最常用且有效的增强方式。ColorJitter通过轻微改变颜色来模拟光照变化进一步提升模型鲁棒性。验证集和测试集绝对不能使用任何随机性增强否则评估指标将不可靠。Normalize的参数必须使用ImageNet的统计量这是为了与预训练模型的特征分布对齐。num_workers可以加速数据加载pin_memory在GPU训练时能提升数据从CPU到GPU的传输速度。3.3 步骤三可视化与检查数据在投入训练前务必检查数据是否被正确加载和增强。# 可以在 notebook 中运行或写一个简单的检查脚本 import matplotlib.pyplot as plt import numpy as np import torch def imshow(inp, titleNone): 显示一个Tensor图像。 inp inp.numpy().transpose((1, 2, 0)) # 从(C, H, W)转为(H, W, C) # 反标准化 mean np.array(IMAGENET_MEAN) std np.array(IMAGENET_STD) inp std * inp mean inp np.clip(inp, 0, 1) plt.imshow(inp) if title is not None: plt.title(title) plt.axis(off) # 获取一个批次的数据 images, labels next(iter(train_loader)) # 创建一个网格来显示 out torchvision.utils.make_grid(images[:8]) # 显示前8张 imshow(out, title[train_dataset.idx_to_class[l.item()] for l in labels[:8]]) plt.show()运行这段代码你应该能看到一个由8张经过随机裁剪、翻转、颜色变换的猫狗图片组成的网格并且图片已经恢复了可视化的颜色。这确认了数据管道工作正常。4. 构建迁移学习模型数据准备就绪后我们就可以构建模型了。这里以使用torchvision.models中的resnet18预训练模型为例。# src/model.py import torch.nn as nn import torchvision.models as models def get_model(num_classes2, pretrainedTrue, freeze_backboneTrue): 获取一个基于ResNet18的迁移学习模型。 参数: num_classes: 输出类别数。 pretrained: 是否加载在ImageNet上预训练的权重。 freeze_backbone: 是否冻结卷积层特征提取器的参数。 # 加载预训练的ResNet18模型 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None) if freeze_backbone: # 冻结所有卷积层的参数使其在训练中不更新 for param in model.parameters(): param.requires_grad False # 替换最后的全连接层 # ResNet18的fc层输入特征数是512 num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes) # 如果冻结了主干只让新加的fc层可训练 # 如果未冻结则所有层都可训练微调模式 # 注意即使冻结了主干BatchNorm层的running mean/var在训练时依然会更新eval模式除外 return model # 使用示例 model get_model(num_classes2, pretrainedTrue, freeze_backboneTrue) print(model) # 可以打印参数查看哪些被冻结 # for name, param in model.named_parameters(): # print(f{name}: requires_grad{param.requires_grad})模型策略选择策略一特征提取器freeze_backboneTrue。仅训练新添加的全连接层。训练速度快所需数据量极少适合数据量非常少每类100或计算资源有限的情况。策略二微调freeze_backboneFalse。解冻所有层用很小的学习率进行训练。通常能获得比策略一更好的效果但需要更多的数据每类几百张和更谨慎的训练防止过拟合。策略三分层微调先冻结大部分层训练分类头训练几轮后再解冻靠近顶部的几层卷积层进行微调。这是一种折中方案。5. 训练与验证循环有了数据和模型我们就可以编写训练脚本了。这里提供一个完整的训练循环示例包含验证和模型保存。# src/train.py import torch import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler import time import copy from tqdm import tqdm import os from src.model import get_model from src.dataset import CatDogDataset from src.config import train_transform, val_test_transform from torch.utils.data import DataLoader def train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochs25, devicecuda): 训练模型的通用函数。 since time.time() best_model_wts copy.deepcopy(model.state_dict()) best_acc 0.0 # 记录训练过程 history {train_loss: [], train_acc: [], val_loss: [], val_acc: []} 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显示进度条 for inputs, labels in tqdm(dataloaders[phase], descf{phase.capitalize()} Epoch {epoch}): 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 and scheduler is not None: scheduler.step() epoch_loss running_loss / len(dataloaders[phase].dataset) epoch_acc running_corrects.double() / len(dataloaders[phase].dataset) # 记录历史 if phase train: history[train_loss].append(epoch_loss) history[train_acc].append(epoch_acc.cpu().item()) else: history[val_loss].append(epoch_loss) history[val_acc].append(epoch_acc.cpu().item()) 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()) # 保存最佳模型 torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict() if scheduler else None, best_acc: best_acc, history: history, }, ../models/best_model.pth) print(f 保存新的最佳模型准确率: {best_acc:.4f}) print() time_elapsed time.time() - since print(f训练完成用时 {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s) print(f最佳验证准确率: {best_acc:.4f}) # 加载最佳模型权重 model.load_state_dict(best_model_wts) return model, history def main(): # 超参数配置 data_dir ../data/processed batch_size 32 num_epochs 15 # 对于小样本epoch不宜过多防止过拟合 learning_rate 0.001 # 初始学习率 momentum 0.9 step_size 7 # 学习率衰减步长 gamma 0.1 # 学习率衰减因子 # 设备选择 device torch.device(cuda if torch.cuda.is_available() else cpu) print(f使用设备: {device}) # 创建数据加载器字典 image_datasets { train: CatDogDataset(os.path.join(data_dir, train), transformtrain_transform), val: CatDogDataset(os.path.join(data_dir, val), transformval_test_transform), } dataloaders { x: DataLoader(image_datasets[x], batch_sizebatch_size, shuffle(x train), num_workers4, pin_memoryTrue) for x in [train, val] } dataset_sizes {x: len(image_datasets[x]) for x in [train, val]} print(f数据集大小: {dataset_sizes}) # 初始化模型 model get_model(num_classes2, pretrainedTrue, freeze_backboneTrue) # 第一阶段冻结特征提取器 model model.to(device) # 定义损失函数和优化器 criterion nn.CrossEntropyLoss() # 只优化那些 requires_gradTrue 的参数即我们新加的fc层 optimizer optim.SGD(model.parameters(), lrlearning_rate, momentummomentum) # 学习率调度器 scheduler lr_scheduler.StepLR(optimizer, step_sizestep_size, gammagamma) # 训练模型 model, history train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochsnum_epochs, devicedevice) # 可以在这里添加第二阶段微调解冻部分层的代码 print(第一阶段训练完成。) # ... (后续可添加微调代码) if __name__ __main__: main()运行此脚本开始训练。由于数据量小且只训练分类头训练速度会非常快。你应该能看到训练损失下降验证准确率逐步提升。6. 模型评估与测试训练完成后我们需要在从未参与训练的测试集上评估模型的最终泛化性能。# src/evaluate.py import torch from src.model import get_model from src.dataset import CatDogDataset from src.config import val_test_transform from torch.utils.data import DataLoader import numpy as np from sklearn.metrics import classification_report, confusion_matrix import seaborn as sns import matplotlib.pyplot as plt def evaluate_model(model_path, test_data_dir, devicecuda): 在测试集上评估模型。 # 加载模型 checkpoint torch.load(model_path, map_locationdevice) model get_model(num_classes2, pretrainedFalse, freeze_backboneFalse) # 加载时不需要预训练权重 model.load_state_dict(checkpoint[model_state_dict]) model model.to(device) model.eval() # 非常重要设置为评估模式 # 创建测试集加载器 test_dataset CatDogDataset(root_dirtest_data_dir, transformval_test_transform) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse, num_workers2) all_preds [] all_labels [] 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) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) # 计算准确率 accuracy np.sum(np.array(all_preds) np.array(all_labels)) / len(all_labels) print(f测试集准确率: {accuracy:.4f}) # 打印分类报告精确率、召回率、F1分数 print(\n分类报告:) print(classification_report(all_labels, all_preds, target_namestest_dataset.get_class_names())) # 绘制混淆矩阵 cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(6,5)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelstest_dataset.get_class_names(), yticklabelstest_dataset.get_class_names()) plt.ylabel(真实标签) plt.xlabel(预测标签) plt.title(混淆矩阵) plt.tight_layout() plt.savefig(../results/confusion_matrix.png) plt.show() return accuracy if __name__ __main__: evaluate_model(../models/best_model.pth, ../data/processed/test, devicecuda if torch.cuda.is_available() else cpu)运行评估脚本你将得到模型在测试集上的准确率、精确率、召回率和F1分数以及一个可视化的混淆矩阵。对于我们的猫狗二分类小样本任务使用迁移学习通常能达到95%以上的测试准确率这充分展示了迁移学习的威力。7. 常见问题与排查思路在实际操作中你可能会遇到以下问题问题现象可能原因解决思路训练损失不下降准确率始终在50%左右二分类1. 数据加载错误标签与图片不匹配。2. 模型输出层维度错误。3. 学习率设置不当太大或太小。4. 预训练模型权重未正确加载。1. 使用imshow函数可视化几个批次的数据和标签确认对应关系。2. 检查model.fc.out_features是否等于类别数。3. 尝试调整学习率如0.01, 0.001, 0.0001。4. 打印模型参数确认requires_grad状态并检查预训练权重是否加载成功。验证准确率远低于训练准确率过拟合1. 训练数据太少。2. 数据增强不够强或未启用。3. 模型过于复杂或训练轮数太多。4. 没有使用验证集进行早停。1. 尝试收集更多数据或使用更激进的数据增强如RandomAffine,RandomPerspective。2. 增加数据增强的强度或种类。3. 减少训练轮数 (num_epochs)或使用更简单的模型如resnet18而非resnet50。4. 实现早停机制当验证损失连续几个epoch不下降时停止训练。GPU内存不足CUDA out of memory1. 批次大小 (batch_size) 太大。2. 输入图片尺寸太大。3. 模型太大。1. 减小batch_size如从32减到16或8。2. 减小RandomResizedCrop和CenterCrop的尺寸如从224到128。3. 换用更小的预训练模型如resnet18代替resnet50。4. 使用torch.cuda.empty_cache()清理缓存。训练速度很慢1. 未使用GPU。2.num_workers设置过小默认为0。3. 数据增强过于复杂。1. 检查torch.cuda.is_available()并确保model.to(device)和data.to(device)。2. 根据CPU核心数增加num_workers通常设置为CPU核心数。3. 简化数据增强流程或在训练初期使用较弱增强。加载保存的模型时报错1. 模型结构定义与保存时不一致。2. PyTorch版本不兼容。3. 加载路径错误。1. 确保加载模型时使用的get_model函数参数与保存时完全一致特别是num_classes。2. 尽量在相同版本的PyTorch环境下进行保存和加载。3. 使用绝对路径并检查文件是否存在。8. 最佳实践与进阶建议掌握了基础流程后以下建议能帮助你在实际项目中做得更好数据永远是王道即使使用迁移学习高质量、标注准确的数据也是成功的基石。对于小样本可以尝试主动学习Active Learning策略即让模型筛选出最不确定的样本交由人工标注用最少的标注成本获得最大性能提升。分层学习率与差分学习率在微调时不同层应该使用不同的学习率。靠近输入的底层卷积层学习通用特征应使用较小的学习率靠近输出的高层层和分类头应使用较大的学习率。可以使用torch.optim中的param_groups来实现。使用更强大的预训练模型ResNet只是起点。可以尝试EfficientNet,Vision Transformer (ViT),Swin Transformer等更现代的架构它们在ImageNet上的表现更好迁移能力可能更强。交叉验证对于极少量数据如每类少于50张可以考虑使用K折交叉验证来更可靠地评估模型性能并充分利用有限的数据。利用外部数据如果领域内公开数据极少可以考虑使用领域自适应Domain Adaptation技术利用大量相关的公开数据源域来辅助训练再迁移到你的小数据目标域。模型集成训练多个不同的模型如使用不同的预训练网络、不同的数据增强策略、不同的随机种子然后将它们的预测结果进行平均或投票通常能稳定地提升最终性能。自动化与实验管理使用Weights Biases (wandb)、TensorBoard或MLflow等工具记录每次实验的超参数、指标和模型便于分析和复现。通过本文的详细拆解你应该已经掌握了使用迁移学习应对小样本图像分类任务的核心流程从高质量的数据准备清洗、划分、增强、标准化到模型构建与策略选择冻结/微调再到完整的训练、验证与评估循环。这套方法不仅适用于猫狗分类可以轻松迁移到任何新的图像分类任务上例如医学影像分析、工业质检、动植物识别等。关键在于理解每个步骤背后的原理并根据自己的数据和任务特点进行灵活调整。