【YOLOv3 Train】YOLOv3训练全过程
文章目录
- 1 网络训练流程概述
- 2 训练前先注意
- 3 整体代码理解
- 4 一轮训练与验证代码
- 5 感谢链接
1 网络训练流程概述
想得到一个model_best模型，宽泛的考虑：拿到数据，训练一下，谁效果好，谁就是model_best。
细致考虑的话，需要注意以下要点：
是否使用cuda？
数据集准备和划分比例关系
训练多少个epoch？
是否要先冻结特征提取部分参数训练，然后再解冻训练？ 通常不是
学习率大小和batch_size大小关系
学习率衰减方式，可使用哪些trick？
网络模型初始化，权重初始化方式，加载预训练模型？
输入图像尺寸
…
2 训练前先注意
从整体看，训练自己的目标检测模型需要注意以下几点：
-
训练前检查数据集格式是否为VOC格式，要有输入图片和标签
(1) 输入图片为.jpg图片，无需固定大小，传入训练前会自动进行resize。
(2) 灰度图会自动转成RGB图片进行训练，无需自己修改。
(3) 输入图片如果后缀非jpg，需要自己批量转成jpg后再开始训练。
(4) 标签为.xml格式，文件中有需要检测的目标信息，标签文件和输入图片文件相对应。 -
训练好的权值文件保存在logs文件夹中
(1) 在训练过程中，并不是只保存最低损失的。
(2) 每个epoch都会保存一次，如果只是训练了几个step是不会保存的 -
损失值的大小用于判断是否收敛
(1) 验证集损失不断下降，当验证集损失基本不改变时，模型基本上就收敛了。
(2) 损失值的具体大小并没有什么意义，大和小只在于损失的计算方式，并不是接近于0才好。
(3) 如果想要让损失好看点，可以直接到对应的损失函数里面除上10000。
(4) 训练过程中的损失值会保存在logs文件夹下的loss_%Y_%m_%d_%H_%M_%S文件夹中
3 整体代码理解
train.py中代码理解
# -------------------------------------#
# 对数据集进行训练
# -------------------------------------#
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.utils.data import DataLoader # 用来加载数据集
# YoloBody：yolov3网络，详解参考链接：https://blog.csdn.net/weixin_45377629/article/details/124080087
from nets.yolo import YoloBody
# YOLOLoss：损失函数的理解见参考链接：https://blog.csdn.net/weixin_45377629/article/details/124343506
# weights_init：权重初始化不赘述
from nets.yolo_training import YOLOLoss, weights_init
from utils.callbacks import LossHistory # 整log用的
# 整数据集，详解见：https://blog.csdn.net/weixin_45377629/article/details/124116916
from utils.dataloader import YoloDataset, yolo_dataset_collate
from utils.utils import get_anchors, get_classes # 获取先验框和所训练数据集的类别
# 训练一个epoch，详解见下一节
from utils.utils_fit import fit_one_epoch
if __name__ == "__main__":
# -------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
# -------------------------------#
Cuda = True
# --------------------------------------------------------#
# 训练前一定要修改classes_path，使其对应自己的数据集
# --------------------------------------------------------#
classes_path = 'model_data/voc_classes.txt'
# ---------------------------------------------------------------------#
# anchors_path代表先验框对应的txt文件，一般不修改。
# anchors_mask用于帮助代码找到对应的先验框，一般不修改。
# ---------------------------------------------------------------------#
anchors_path = 'model_data/yolo_anchors.txt'
anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
# ------------------------------------------------------------------------------------------------------------------#
# 模型的 预训练权重 对不同数据集是通用的，因为特征是通用的。
# 模型的 预训练权重 比较重要的部分是 主干特征提取网络的权值部分，用于进行特征提取。
# 预训练权重对于99%的情况都必须要用，不用的话主干部分的权值太过随机，特征提取效果不明显，网络训练的结果也不会好
#
# 如果训练过程中存在中断训练的操作，可以将model_path设置成logs文件夹下的权值文件，将已经训练了一部分的权值再次载入。
# 同时修改下方的 冻结阶段 或者 解冻阶段 的参数，来保证模型epoch的连续性。
#
# 当model_path = ''的时候不加载整个模型的权值。
#
# 如果想要让模型从0开始训练，则设置model_path = ''，下面的Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
# 一般来讲，从0开始训练效果会很差，因为权值太过随机，特征提取效果不明显。
#
# 网络一般不从0开始训练，至少会使用主干部分的权值，有些论文提到可以不用预训练，主要原因是他们 数据集较大 且 调参能力优秀。
# 如果一定要训练网络的主干部分，可以了解imagenet数据集，首先训练分类模型，分类模型的 主干部分 和该模型通用，基于此进行训练。
# -----------------------------------------------------------------------------------------------------------------#
model_path = 'model_data/yolo_weights.pth'
# ------------------------------------------------------#
# 输入的shape大小，一定要是32的倍数
# ------------------------------------------------------#
input_shape = [416, 416]
# ----------------------------------------------------#
# 训练分为两个阶段，分别是冻结阶段和解冻阶段。
# 显存不足与数据集大小无关，提示显存不足请调小batch_size。
# 受到BatchNorm层影响，batch_size最小为2，不能为1。
# ----------------------------------------------------#
# ----------------------------------------------------#
# 冻结阶段训练参数
# 此时模型的主干被冻结了，特征提取网络不发生改变
# 占用的显存较小，仅对网络进行微调
# ----------------------------------------------------#
Init_Epoch = 0
Freeze_Epoch = 50
Freeze_batch_size = 8
Freeze_lr = 1e-3
# ----------------------------------------------------#
# 解冻阶段训练参数
# 此时模型的主干不被冻结了，特征提取网络会发生改变
# 占用的显存较大，网络所有的参数都会发生改变
# ----------------------------------------------------#
UnFreeze_Epoch = 100
Unfreeze_batch_size = 4
Unfreeze_lr = 1e-4
# ------------------------------------------------------#
# 是否进行冻结训练，默认先冻结主干训练后解冻训练。
# ------------------------------------------------------#
Freeze_Train = True
# ------------------------------------------------------#
# 用于设置是否使用多线程读取数据
# 开启后会加快数据读取速度，但是会占用更多内存
# 内存较小的电脑可以设置为2或者0
# ------------------------------------------------------#
num_workers = 4
# ----------------------------------------------------#
# 获得图片路径和标签
# ----------------------------------------------------#
train_annotation_path = '2007_train.txt'
val_annotation_path = '2007_val.txt'
# ----------------------------------------------------#
# 获取classes和anchor
# ----------------------------------------------------#
class_names, num_classes = get_classes(classes_path)
anchors, num_anchors = get_anchors(anchors_path)
# ------------------------------------------------------#
# 创建yolo模型
# ------------------------------------------------------#
model = YoloBody(anchors_mask, num_classes)
weights_init(model) # 模型初始化参数
if model_path != '': # 是否加载预训练模型
print('Load weights {}.'.format(model_path))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dict = model.state_dict()
pretrained_dict = torch.load(model_path, map_location=device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model_train = model.train()
if Cuda:
model_train = torch.nn.DataParallel(model)
# 设置这个 flag 可以让内置的 cuDNN 的 auto-tuner 自动寻找最适合当前配置的高效算法，来达到优化运行效率的问题。
cudnn.benchmark = True
model_train = model_train.cuda()
# YOLOLoss见参考链接：
yolo_loss = YOLOLoss(anchors, num_classes, input_shape, Cuda, anchors_mask)
loss_history = LossHistory("logs/")
# ---------------------------#
# 读取数据集对应的txt
# ---------------------------#
with open(train_annotation_path) as f:
train_lines = f.readlines()
with open(val_annotation_path) as f:
val_lines = f.readlines()
num_train = len(train_lines)
num_val = len(val_lines)
# ------------------------------------------------------#
# 主干特征提取网络特征通用，冻结训练可以加快训练速度
# 也可以在训练初期防止权值被破坏。
# Init_Epoch为起始世代
# Freeze_Epoch为冻结训练的轮数
# UnFreeze_Epoch总训练轮数
# 提示显存不足请调小Batch_size
# ------------------------------------------------------#
if True:
batch_size = Freeze_batch_size
lr = Freeze_lr
start_epoch = Init_Epoch # 冻结了训练
end_epoch = Freeze_Epoch
epoch_step = num_train // batch_size
epoch_step_val = num_val // batch_size
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
optimizer = optim.Adam(model_train.parameters(), lr, weight_decay=5e-4) # optimizer优化器
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.94) # 学习率衰减方式
# ---------------------------------#
# 数据集数据读取
# 此处可参考：https://blog.csdn.net/weixin_45377629/article/details/124116916
# ---------------------------------#
train_dataset = YoloDataset(train_lines, input_shape, num_classes, train=True)
val_dataset = YoloDataset(val_lines, input_shape, num_classes, train=False)
gen = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=num_workers, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
gen_val = DataLoader(val_dataset, shuffle=True, batch_size=batch_size, num_workers=num_workers, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
if Freeze_Train: # 冻结backbone参数训练
for param in model.backbone.parameters():
param.requires_grad = False
for epoch in range(start_epoch, end_epoch): # 训练多个epoch，也就是重复多次而已
# 见下面解释！
fit_one_epoch(model_train, model, yolo_loss, loss_history, optimizer, epoch,
epoch_step, epoch_step_val, gen, gen_val, end_epoch, Cuda)
lr_scheduler.step()
if True:
batch_size = Unfreeze_batch_size
lr = Unfreeze_lr
start_epoch = Freeze_Epoch # 解冻后训练
end_epoch = UnFreeze_Epoch
epoch_step = num_train // batch_size
epoch_step_val = num_val // batch_size
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
optimizer = optim.Adam(model_train.parameters(), lr, weight_decay=5e-4)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.94)
train_dataset = YoloDataset(train_lines, input_shape, num_classes, train=True)
val_dataset = YoloDataset(val_lines, input_shape, num_classes, train=False)
# gen一般就是指train_loader
gen = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=num_workers, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
# gen_val一般指val_loader
gen_val = DataLoader(val_dataset, shuffle=True, batch_size=batch_size, num_workers=num_workers, pin_memory=True,
drop_last=True, collate_fn=yolo_dataset_collate)
if Freeze_Train:
for param in model.backbone.parameters():
param.requires_grad = True
for epoch in range(start_epoch, end_epoch):
fit_one_epoch(model_train, model, yolo_loss, loss_history, optimizer, epoch,
epoch_step, epoch_step_val, gen, gen_val, end_epoch, Cuda)
lr_scheduler.step()
4 一轮训练与验证代码
看一个epoch的内部操作
import torch
from tqdm import tqdm # 用来显示进度
from utils.utils import get_lr
def fit_one_epoch(model_train, model, yolo_loss, loss_history, optimizer, epoch, epoch_step, epoch_step_val, gen,
gen_val, Epoch, cuda):
loss = 0
val_loss = 0
model_train.train()
print('Start Train')
with tqdm(total=epoch_step, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3) as pbar:
for iteration, batch in enumerate(gen):
if iteration >= epoch_step:
break
images, targets = batch[0], batch[1]
with torch.no_grad():
if cuda:
images = torch.from_numpy(images).type(torch.FloatTensor).cuda()
targets = [torch.from_numpy(ann).type(torch.FloatTensor).cuda() for ann in targets]
else:
images = torch.from_numpy(images).type(torch.FloatTensor)
targets = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets]
# ----------------------#
# 清零梯度
# ----------------------#
optimizer.zero_grad()
# ----------------------#
# 前向传播，得到yolo net的三层不同size的输出
# ----------------------#
outputs = model_train(images)
loss_value_all = 0
num_pos_all = 0
# ----------------------#
# 计算损失
# ----------------------#
for l in range(len(outputs)):
# ----------------------------------------------#
# 似乎可以是这样yolo_loss.forward(l, outputs[l], targets)，其实不需要
# 因为YOLOLoss继承自nn.Module，Module中定义了__call__()函数，该函数调用了forward()函数，
# 在执行该语句时，会自动调用__call__()函数
# 此处见参考链接：https://blog.csdn.net/weixin_45377629/article/details/124343506
# ----------------------------------------------#
loss_item, num_pos = yolo_loss(l, outputs[l], targets)
loss_value_all += loss_item
num_pos_all += num_pos
loss_value = loss_value_all / num_pos_all
# ----------------------#
# 反向传播
# ----------------------#
loss_value.backward()
optimizer.step()
loss += loss_value.item()
pbar.set_postfix(**{'loss': loss / (iteration + 1),
'lr': get_lr(optimizer)})
pbar.update(1)
print('Finish Train')
model_train.eval()
print('Start Validation')
with tqdm(total=epoch_step_val, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3) as pbar:
for iteration, batch in enumerate(gen_val):
if iteration >= epoch_step_val:
break
images, targets = batch[0], batch[1]
with torch.no_grad():
if cuda:
images = torch.from_numpy(images).type(torch.FloatTensor).cuda()
targets = [torch.from_numpy(ann).type(torch.FloatTensor).cuda() for ann in targets]
else:
images = torch.from_numpy(images).type(torch.FloatTensor)
targets = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets]
# ----------------------#
# 清零梯度
# ----------------------#
optimizer.zero_grad()
# ----------------------#
# 前向传播
# ----------------------#
outputs = model_train(images)
loss_value_all = 0
num_pos_all = 0
# ----------------------#
# 计算损失，每一个特征图计算一次损失
# ----------------------#
for l in range(len(outputs)):
loss_item, num_pos = yolo_loss(l, outputs[l], targets)
loss_value_all += loss_item
num_pos_all += num_pos
loss_value = loss_value_all / num_pos_all
val_loss += loss_value.item()
pbar.set_postfix(**{'val_loss': val_loss / (iteration + 1)})
pbar.update(1)
print('Finish Validation')
loss_history.append_loss(loss / epoch_step, val_loss / epoch_step_val)
print('Epoch:' + str(epoch + 1) + '/' + str(Epoch))
print('Total Loss: %.3f || Val Loss: %.3f ' % (loss / epoch_step, val_loss / epoch_step_val))
torch.save(model.state_dict(),
'logs/ep%03d-loss%.3f-val_loss%.3f.pth' % (epoch + 1, loss / epoch_step, val_loss / epoch_step_val))
5 感谢链接
https://blog.csdn.net/weixin_44791964/article/details/105310627
https://www.bilibili.com/video/BV1Hp4y1y788?p=15&spm_id_from=pageDriver