备数据 → 做预处理 → 加载模型 → 冻结骨干 → 配损失优化器 → 前向传播算损失 → 反向更新 → 验证 → 保存 → 推理 → 解冻精调
- 准备数据集(划分训练集 / 验证集)
- 数据预处理 + 数据增强
- 构建 / 加载模型(自建模型 或 预训练模型)
- 微调专属:冻结骨干网络
- 配置训练参数(设备、损失函数、优化器)
- 训练循环:前向传播 → 算损失 → 反向传播 → 更新参数
- 每轮验证模型准确率
- 保存最优模型权重
- 推理预测(用训练好的模型做预测)
- 微调进阶:解冻部分层,小学习率再训练
一、逐步骤详细讲解 + 核心代码
第 1 步:准备数据集
把数据分成两部分:
train训练集:教模型学习val验证集:看模型学没学会
文件夹标准格式(CV 通用):
plaintext
data/
train/
类别1/ 图片
类别2/ 图片
val/
类别1/ 图片
类别2/ 图片
第 2 步:数据预处理 + 数据加载
做三件事:
- 缩放图片、翻转、裁剪(数据增强)
- 归一化
- 用
Dataset + DataLoader批量喂给模型
python
运行
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
# 1. 数据增强
train_transform = transforms.Compose([
transforms.Resize((224,224)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])
# 2. 加载数据集
train_dataset = datasets.ImageFolder("data/train", transform=train_transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
第 3 步:加载模型(两种情况)
情况 A:从零自己搭模型
python
运行
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 2)
def forward(self, x):
return self.fc(x)
model = Net()
情况 B:微调 加载预训练模型(最常用)
python
运行
from torchvision import models
# 加载预训练ResNet
model = models.resnet18(pretrained=True)
第 4 步:微调关键 ------ 冻结骨干网络
冻结 :让预训练的主干网络参数不更新、不被破坏,只训练自己的分类头。
python
运行
# 冻结所有层
for param in model.parameters():
param.requires_grad = False
# 替换最后一层分类头,适配自己的类别数
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2) # 2分类
第 5 步:配置训练必备组件
python
运行
import torch
import torch.optim as optim
# 设备(有GPU用GPU,没GPU用CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 损失函数(分类必用)
loss_fn = nn.CrossEntropyLoss()
# 优化器(更新参数)
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)
第 6 步:训练循环(核心)
固定四件套:
- 前向传播
model(x) - 计算损失
loss_fn - 反向传播
loss.backward() - 参数更新
optimizer.step()
python
运行
model.train() # 切换到训练模式
for x, label in train_loader:
x, label = x.to(device), label.to(device)
# 1.前向传播
out = model(x)
# 2.算损失
loss = loss_fn(out, label)
# 3.梯度清零 + 反向传播 + 更新
optimizer.zero_grad()
loss.backward()
optimizer.step()
第 7 步:验证模型
每训练一轮,用验证集测准确率,看模型效果:
python
运行
model.eval() # 推理模式,关闭dropout/BN
with torch.no_grad(): # 不计算梯度,省显存
# 跑验证集,计算正确率
pass
第 8 步:保存模型
保存训练好的权重,下次直接用,不用重新训练:
python
运行
torch.save(model.state_dict(), "best_model.pth")
第 9 步:模型推理(拿来做预测)
加载模型 → 输入数据 → 出结果
python
运行
# 加载权重
model.load_state_dict(torch.load("best_model.pth"))
model.eval()
with torch.no_grad():
pred = model(输入数据)
第 10 步:微调进阶 ------ 解冻再训练
当分类头训练好后:
- 解冻骨干后面几层
- 用更小学习率继续训练
- 模型效果进一步提升
python
运行
# 解冻最后一层残差块
for param in model.layer4.parameters():
param.requires_grad = True
# 更小学习率
optimizer = optim.Adam(model.parameters(), lr=1e-4)
二、从零训练 vs 微调 区别
1. 从零训练
- 不用冻结网络
- 全部参数都训练
- 学习率正常 0.001
- 需要大量数据
2. 模型微调(Finetune)
- 先冻结骨干,只训分类头
- 再解冻部分层,小学习率精调
- 少量数据就能训出好效果
- 工业项目、比赛、毕设全用这个