深度学习模型训练流程和代码框架
0 概述
在训练深度学习模型(如 ResNet、Transformer、Diffusion 等)时,通常需要遵循标准化的训练流程,包括 数据准备、模型定义、训练、验证、测试、评估与保存等步骤。下面是训练流程概述:
1 数据准备
数据的质量决定了模型的上限,所以数据处理是至关重要的一步。
1.1 数据收集
- 下载、抓取或合成数据,例如:
- 公开数据集(ImageNet、COCO、MNIST 等)
- 自定义数据(遥感图像、医学图像、草图数据等)
1.2 数据增强和预处理
- 数据增强
- 分为:离线增强方法、在线增强方法
- 增强操作:旋转、翻转、模糊、马赛克增强、混合增强等
- 数据预处理
- 归一化(Normalization):\((x - mean) / std\)
- 图像缩放:把图像缩放到网络要求的输入大小
- 文本处理:文本转换为数字、零填充等
- 格式转换:将numpy类型转换到tensor类型
1.3 数据加载
-
定义数据集
-
定义数据加载器
- 训练集加载器:
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) - 验证集加载器:
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False) - 测试集加载器:
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
- 训练集加载器:
2 定义模型
- 选择合适的架构(CNN、Transformer、GAN、Diffusion、Graph Neural Network)
- 确定模型的层数、参数量
- 加载预训练权重(可选)python
3 定义损失函数 & 优化器
3.1 选择损失函数
- 分类任务:交叉熵损失
loss_fn = nn.CrossEntropyLoss()
- 回归任务:均方误差
loss_fn = nn.MSELoss()
- 自定义损失:如Hausdorff距离、CLIPLoss等
class MyCustomLoss(nn.Module):
3.2 选择优化器
- SGD(适用于 ResNet、CNN)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
- AdamW(适用于 Transformer)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
3.3 学习率调度
- StepLR(每隔 X 轮衰减)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
- Cosine Annealing(适用于 Transformer、Diffusion)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
4. 训练模型
在每个epoch中:
- 训练阶段
- 前向传播(forward)
- 计算损失(loss)
- 反向传播(backward)
- 更新参数(optimizer.step())
- 验证阶段
- 计算在验证集上的损失
- 评估指标(如 Accuracy、IoU、F1-Score)
5 评估 & 测试
- 计算指标(accuracy、mIoU、F1-score、PSNR)
- 绘制 loss 曲线
- 生成可视化结果(GradCAM、Attention Map)
6 模型保存 & 记录
6.1 保存最佳模型
torch.save(model.state_dict(), "best_model.pth")
【精度监测|模型保存】深度学习卷积神经网络的训练过程中如何监测训练集和验证集的精度并保存最优模型呢??- CSDN
6.2 记录日志
使用wandb或TensorBoard记录训练日志
import wandb
wandb.init(project="my_experiment")
for epoch inrange(epochs):
wandb.log({"epoch": epoch, "loss": loss_value})
7 部署 & 推理
8 超参数调优
- Grid Search / Random Search
- 使用
Optuna自动调参
import optuna
def objective(trial):
lr = trial.suggest_loguniform("lr", 1e-5, 1e-1)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
...
return val_loss
study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=50)
9 图示小结

10 完整的模型训练套路:代码模板
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader,Dataset
# 定义模型
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.fc1=nn.Linear(100,50)
self.fc2=nn.Linear(50,10)
def forward(self,x):
x=torch.relu(self.fc1(x))
x=self.fc2(x)
return x
# 定义训练集
class TrainDataset(Dataset):
def __init__(self,data_path,target_path):
super(TrainDataset,self).__init__()
pass # 初始化
def __getitem__(self, index):
pass# 根据index返回数据
# return data,target
def __len__(self):
pass # 返回数据的长度
# return len(xxx)
# 定义验证集
class ValDataset(Dataset):
def __init__(self, data_path, target_path):
super(ValDataset, self).__init__()
pass # 初始化
def __getitem__(self, index):
pass # 根据index返回数据
# return data,target
def __len__(self):
pass # 返回数据的长度
# return len(xxx)
# 定义测试集
class TestDataset(Dataset):
def __init__(self, data_path, target_path):
super(TestDataset, self).__init__()
pass # 初始化
def __getitem__(self, index):
pass # 根据index返回数据
# return data,target
def __len__(self):
pass # 返回数据的长度
# return len(xxx)
# 训练
def train(model,optimizer,criterion,train_loader,device):
total_loss=0
# 将模型设置为训练模式
model.train()
# 遍历数据集
for data,target in train_loader:
# 将数据移到设备上
data,target=data.to(device),target.to(device)
# 梯度清零
optimizer.zero_grad()
# 前向传播
output=model(data)
# 计算损失
loss=criterion(output,target)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
# 统计损失
total_loss+=loss.item()
# 计算平均损失
avg_loss=total_loss/len(train_loader)
return avg_loss
# 验证
def validate(model,criterion,val_loader,device):
total_loss=0
# 将模型设置为验证模式
model.eval()
with torch.no_grad():
#遍历验证集
for data,target in val_loader:
# 将数据移到设备上
data,target=data.to(device),target.to(device)
# 前向传播
output=model(data)
# 计算损失
loss=criterion(output,target)
# 统计损失
total_loss+=loss.item()
# 计算平均损失
avg_loss=total_loss/len(val_loader)
return avg_loss
def test(model,criterion,test_loader,device):
total_loss=0
total_correct=0
# 将模型设置为测试模式
model.eval()
with torch.no_grad():
# 遍历测试集
for data,target in test_loader:
# 将数据迁移到设备上
data,target=data.to(device),target.to(device)
# 前向传播
output=model(data)
# 计算损失
loss=criterion(output,target)
# 统计损失
total_loss+=loss.item()
# 计算准确率
_,predicted =torch.max(output.data,1)
total_correct+=(predicted==target).sum().item()
# 计算平均损失
avg_loss=total_loss/len(test_loader)
# 计算准确率
accuracy=total_correct/len(test_loader.dataset)
return avg_loss,accuracy
# 主函数
def main():
# 定义超参数
epoch=10
lr=0.01
batch_size=32
# 定义设备
device =torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 定义训练集、验证集、测试集 的数据
train_data=TrainDataset(data_path='',target_path='')
val_data = ValDataset(data_path='', target_path='')
test_data = TestDataset(data_path='', target_path='')
# 定义训练集、验证集、测试集 的数据加载器
train_loader=DataLoader(train_data,batch_size=batch_size,shuffle=True)
val_loader=DataLoader(val_data,batch_size=batch_size,shuffle=True)
test_loader=DataLoader(test_data,batch_size=batch_size,shuffle=True)
# 定义模型
model=Model().to(device)
# 加载预训练权重
model.load_state_dict(torch.load('xxx.pth')
# 定义损失函数和优化器
criterion=nn.CrossEntropyLoss()
optimizer=optim.SGD(model.parameters(),lr=lr)
# 训练和验证
for epoch in range(1,epoch+1):
train_loss=train(model,optimizer,criterion,train_loader,device)
val_loss=validate(model,criterion,val_loader,device)
# 显示训练集和验证集的损失
if (epoch%display_iter)==0:
print(f"Epoch: {epoch}, Train loss:{train_loss:.4f},Val loss: {val_loss: .4f}")
# 保存权重
if (epoch%snapshot_iter)==0:
torch.save(model.state_dict(),"xxx.pth")
# 测试
test_loss,test_accuracy=test(model,criterion,test_loader,device)
print(f"Test loss:{test_loss:.4f},Test accuracy:{test_accuracy:.4f}")
if __name__=='__main__':
main()
参考链接
深度学习模型训练完整流程(系统整理) - 向阳红羊羊的文章 - 知乎
深度学习pytorch训练代码模板(个人习惯) - wfnian的文章 - 知乎
一个完整的Pytorch深度学习项目代码,项目结构是怎样的? - 知乎
示例网络:
# 下采样模块
class DownSample(nn.Module):
def __init__(self, in_channels, out_channels):
super(DownSample, self).__init__()
# self.layers = nn.Sequential(
# nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1),
# nn.BatchNorm2d(out_channels),
# nn.LeakyReLU(0.2)
# )
self.conv_1 = nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)
self.bn_1 = nn.BatchNorm2d(out_channels)
self.l_relu = nn.LeakyReLU(0.2)
def forward(self, x):
# out = self.layers(x)
out = self.l_relu(self.bn_1(self.conv_1(x)))
return out
# 上采样模块
class UpSample(nn.Module):
def __init__(self, in_channels, out_channels):
super(UpSample, self).__init__()
self.layers = nn.Sequential(
nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU()
)
def forward(self, x):
return self.layers(x)
# 完整网络
class CustomNet(nn.Module):
def __init__(self):
super(CustomNet, self).__init__()
# 定义网络结构
self.down1 = DownSample(3, 64)
self.down2 = DownSample(64, 128)
self.up1 = UpSample(128, 64)
self.up2 = UpSample(64, 3)
# 参数初始化
self.apply(self._init_weights)
def _init_weights(self, m):
"""Kaiming初始化方法"""
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
if m.bias is not None:
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
def forward(self, x):
# 下采样路径
d1 = self.down1(x)
d2 = self.down2(d1)
# 上采样路径
u1 = self.up1(d2)
u2 = self.up2(u1)
return u2
pass
pass
请问,我用for循环遍历model.state_dict()打印模型各层的信息,其中有如下一部分输出,是nn.BatchNorm2d层的信息,这个信息是什么意思?
down1.bn_1.running_mean
down1.bn_1.running_var
down1.bn_1.num_batches_tracked
在 PyTorch 的 nn.BatchNorm2d 层中,state_dict() 输出的以下参数是批归一化(Batch Normalization)层的核心统计量,其作用如下:
-
down1.bn_1.running_mean
作用:存储训练过程中所有 batch 数据的 滑动平均均值。
在训练阶段,每个 batch 的均值会被动态计算,并通过指数滑动平均(EMA)更新此值。
在推理阶段,直接使用此值替代当前 batch 的均值,保证模型输出的稳定性。 -
down1.bn_1.running_var
作用:存储训练过程中所有 batch 数据的 滑动平均方差。
与 running_mean 类似,通过滑动平均更新,用于推理阶段的归一化。 -
down1.bn_1.num_batches_tracked
作用:记录模型训练过程中 已处理的 batch 总数。
在加载预训练模型时,此值用于恢复训练进度,避免动量参数(momentum)计算错误。
若模型从未训练过(如初始化状态),此值为 0;每训练一个 batch 后自动加1 -
关键注意事项
训练与推理模式差异:
训练模式:使用当前 batch 的均值和方差归一化数据,并更新 running_mean 和 running_var。
推理模式:固定使用 running_mean 和 running_var 进行归一化。
模型保存与加载:
若需恢复训练,必须同时保存 running_mean、running_var 和 num_batches_tracked,否则会导致统计量错误6。
参数形状:
running_mean 和 running_var 的形状与输入通道数相同(如 down1.bn_1 对应输入通道数为 64,则形状为 (64,))。
请问running_mean 和 running_var这两个参数会计算梯度吗?是否包含在计算图中?梯度反向传播会更新这两个参数吗?
-
梯度计算与计算图
不计算梯度:running_mean 和 running_var 是 非可学习参数,其更新不依赖梯度反向传播,而是通过前向传播时的滑动平均(EMA)计算。
不包含在计算图中:这两个参数的值通过 momentum 和当前 batch 的统计量直接更新,与计算图无关。

浙公网安备 33010602011771号