代码主要是参考bubbliiing的github YOLOv3的代码：github.com/bubbliiiing…
对于源代码的解读
训练部分
train.py文件
'''
训练自己的目标检测模型一定需要注意以下几点：
1、训练前仔细检查自己的格式是否满足要求，该库要求数据集格式为VOC格式，需要准备好的内容有输入图片和标签
输入图片为.jpg图片，无需固定大小，传入训练前会自动进行resize。
灰度图会自动转成RGB图片进行训练，无需自己修改。
输入图片如果后缀非jpg，需要自己批量转成jpg后再开始训练。
标签为.xml格式，文件中会有需要检测的目标信息，标签文件和输入图片文件相对应。
2、训练好的权值文件保存在logs文件夹中，每个epoch都会保存一次，如果只是训练了几个step是不会保存的，epoch和step的概念要捋清楚一下。
在训练过程中，该代码并没有设定只保存最低损失的，因此按默认参数训练完会有100个权值，如果空间不够可以自行删除。
这个并不是保存越少越好也不是保存越多越好，有人想要都保存、有人想只保存一点，为了满足大多数的需求，还是都保存可选择性高。
3、损失值的大小用于判断是否收敛，比较重要的是有收敛的趋势，即验证集损失不断下降，如果验证集损失基本上不改变的话，模型基本上就收敛了。
损失值的具体大小并没有什么意义，大和小只在于损失的计算方式，并不是接近于0才好。如果想要让损失好看点，可以直接到对应的损失函数里面除上10000。
训练过程中的损失值会保存在logs文件夹下的loss_%Y_%m_%d_%H_%M_%S文件夹中
4、调参是一门蛮重要的学问，没有什么参数是一定好的，现有的参数是我测试过可以正常训练的参数，因此我会建议用现有的参数。
但是参数本身并不是绝对的，比如随着batch的增大学习率也可以增大，效果也会好一些；过深的网络不要用太大的学习率等等。
这些都是经验上，只能靠各位同学多查询资料和自己试试了。
'''
if __name__ == "__main__":
#-------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
#-------------------------------#
Cuda = False
#---------------------------------------------------------------------#
# classes_path 指向model_data下的txt，与自己训练的数据集相关
# 训练前一定要修改classes_path，使其对应自己的数据集
#---------------------------------------------------------------------#
# 在训练的过程中需要修改成我们自己的类标签文件
classes_path = 'model_data/voc_classes.txt'
#---------------------------------------------------------------------#
# anchors_path 代表先验框对应的txt文件，一般不修改。
# anchors_mask 用于帮助代码找到对应的先验框，一般不修改。
#---------------------------------------------------------------------#
# anchor的基本信息，和预测一样，一般不进行修改
anchors_path = 'model_data/yolo_anchors.txt'
anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
#----------------------------------------------------------------------------------------------------------------------------#
# 权值文件的下载请看README，可以通过网盘下载。模型的 预训练权重 对不同数据集是通用的，因为特征是通用的。
# 模型的 预训练权重 比较重要的部分是 主干特征提取网络的权值部分，用于进行特征提取。
# 预训练权重对于99%的情况都必须要用，不用的话主干部分的权值太过随机，特征提取效果不明显，网络训练的结果也不会好
#
# 如果训练过程中存在中断训练的操作，可以将model_path设置成logs文件夹下的权值文件，将已经训练了一部分的权值再次载入。
# 同时修改下方的 冻结阶段 或者 解冻阶段 的参数，来保证模型epoch的连续性。
#
# 当model_path = ''的时候不加载整个模型的权值。
#
# 此处使用的是整个模型的权重，因此是在train.py进行加载的，下面的pretrain不影响此处的权值加载。
# 如果想要让模型从主干的预训练权值开始训练，则设置model_path = ''，下面的pretrain = True，此时仅加载主干。
# 如果想要让模型从0开始训练，则设置model_path = ''，下面的pretrain = Fasle，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
#
# 一般来讲，网络从0开始的训练效果会很差，因为权值太过随机，特征提取效果不明显，因此非常、非常、非常不建议大家从0开始训练！
# 如果一定要从0开始，可以了解imagenet数据集，首先训练分类模型，获得网络的主干部分权值，分类模型的 主干部分 和该模型通用，基于此进行训练。
#----------------------------------------------------------------------------------------------------------------------------#
# 就是我们的模型参数
model_path = 'model_data/yolo_weights.pth'
#------------------------------------------------------#
# input_shape 输入的shape大小，一定要是32的倍数
#------------------------------------------------------#
# 输入大小，输入图片越大，效果越准确，但是会更慢，但是必须是32的倍数
input_shape = [416, 416]
#----------------------------------------------------------------------------------------------------------------------------#
# pretrained 是否使用主干网络的预训练权重，此处使用的是主干的权重，因此是在模型构建的时候进行加载的。
# 如果设置了model_path，则主干的权值无需加载，pretrained的值无意义。
# 如果不设置model_path，pretrained = True，此时仅加载主干开始训练。
# 如果不设置model_path，pretrained = False，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
#----------------------------------------------------------------------------------------------------------------------------#
# 注意这个地方写的是主干网络，就是model_path加载整个网络的权重，model_path有值pretrained设置成啥都没用，model_path = ""的时候，我们设置了pretrained,在我们的darknet53的主干网络中，有个初始话参数的位置，直接进行初始化参数。model_path="",pretrained=False，直接从0开始训练
pretrained = False
#----------------------------------------------------------------------------------------------------------------------------#
# 训练分为两个阶段，分别是冻结阶段和解冻阶段。设置冻结阶段是为了满足机器性能不足的同学的训练需求。
# 冻结训练需要的显存较小，显卡非常差的情况下，可设置Freeze_Epoch等于UnFreeze_Epoch，此时仅仅进行冻结训练。
#
# 在此提供若干参数设置建议，各位训练者根据自己的需求进行灵活调整：
# （一）从整个模型的预训练权重开始训练：
# Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 100，Freeze_Train = True（默认参数）
# Init_Epoch = 0，UnFreeze_Epoch = 100，Freeze_Train = False（不冻结训练）
# 其中：UnFreeze_Epoch可以在100-300之间调整。optimizer_type = 'sgd'，Init_lr = 1e-2。
# （二）从主干网络的预训练权重开始训练：
# Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 300，Freeze_Train = True（冻结训练）
# Init_Epoch = 0，UnFreeze_Epoch = 300，Freeze_Train = False（不冻结训练）
# 其中：由于从主干网络的预训练权重开始训练，主干的权值不一定适合目标检测，需要更多的训练跳出局部最优解。
# UnFreeze_Epoch可以在200-300之间调整，YOLOV5和YOLOX均推荐使用300。optimizer_type = 'sgd'，Init_lr = 1e-2。
# （三）batch_size的设置：
# 在显卡能够接受的范围内，以大为好。显存不足与数据集大小无关，提示显存不足（OOM或者CUDA out of memory）请调小batch_size。
# 受到BatchNorm层影响，batch_size最小为2，不能为1。
# 正常情况下Freeze_batch_size建议为Unfreeze_batch_size的1-2倍。不建议设置的差距过大，因为关系到学习率的自动调整。
#----------------------------------------------------------------------------------------------------------------------------#
#------------------------------------------------------------------#
# 冻结阶段训练参数
# 此时模型的主干被冻结了，特征提取网络不发生改变
# 占用的显存较小，仅对网络进行微调
# Init_Epoch 模型当前开始的训练世代，其值可以大于Freeze_Epoch，如设置：
# Init_Epoch = 60、Freeze_Epoch = 50、UnFreeze_Epoch = 100
# 会跳过冻结阶段，直接从60代开始，并调整对应的学习率。
# （断点续练时使用）
# Freeze_Epoch 模型冻结训练的Freeze_Epoch
# (当Freeze_Train=False时失效)
# Freeze_batch_size 模型冻结训练的batch_size
# (当Freeze_Train=False时失效)
#------------------------------------------------------------------#
# 启用了冻结训练的话，首先我们要明确的是冻结的主干网络，就是我们的darknet53的参数不进行变化，但是另外的5层特征提取 + 2层结果预测的部分还是会参数变化的，因为我们采用冻结训练，部分参数不用计算，占用的缓存比较小，所以我们可以将我们的batch_size设置的大一点
Init_Epoch = 0
Freeze_Epoch = 50
Freeze_batch_size = 16
#------------------------------------------------------------------#
# 解冻阶段训练参数
# 此时模型的主干不被冻结了，特征提取网络会发生改变
# 占用的显存较大，网络所有的参数都会发生改变
# UnFreeze_Epoch 模型总共训练的epoch
# Unfreeze_batch_size 模型在解冻后的batch_size
#------------------------------------------------------------------#
# 解冻阶段我们的darknet53的参数也在不断的计算，发生变化，所以显存的占用也会不断的增大，所以我们要将我们的batch_size设置的小一点
UnFreeze_Epoch = 100
Unfreeze_batch_size = 8
#------------------------------------------------------------------#
# Freeze_Train 是否进行冻结训练
# 默认先冻结主干训练后解冻训练。
#------------------------------------------------------------------#
Freeze_Train = True
#------------------------------------------------------------------#
# 其它训练参数：学习率、优化器、学习率下降有关
#------------------------------------------------------------------#
#------------------------------------------------------------------#
# Init_lr 模型的最大学习率
# Min_lr 模型的最小学习率，默认为最大学习率的0.01
#------------------------------------------------------------------#
# 学习率的设置
Init_lr = 1e-2
Min_lr = Init_lr * 0.01
#------------------------------------------------------------------#
# optimizer_type 使用到的优化器种类，可选的有adam、sgd
# 当使用Adam优化器时建议设置 Init_lr=1e-3
# 当使用SGD优化器时建议设置 Init_lr=1e-2
# momentum 优化器内部使用到的momentum参数
# weight_decay 权值衰减，可防止过拟合
#------------------------------------------------------------------#
# 我们使用的优化方法，优化方法momentum参数和权值衰减参数的设置
optimizer_type = "sgd"
momentum = 0.937
weight_decay = 5e-4
#------------------------------------------------------------------#
# lr_decay_type 使用到的学习率下降方式，可选的有step、cos
#------------------------------------------------------------------#
lr_decay_type = "cos"
#------------------------------------------------------------------#
# save_period 多少个epoch保存一次权值，默认每个世代都保存
#------------------------------------------------------------------#
# 训练多少个epoch后，保存整个模型，默认是训练一个epoch进行保存，我的策略一般是在前期
# loss的波动比较大，我们可以相隔一段epoch然后在进行保存，一般我会在前50个epoch中
# 每5个epoch保存一次，当我们过了50个epoch后，loss下降到一个稳定的位置，我们需要进行微调
# 我们可以使用更小的学习率，每一个epoch进行一个保存，这样的话，可以节约很大的空间
save_period = 1
#------------------------------------------------------------------#
# num_workers 用于设置是否使用多线程读取数据
# 开启后会加快数据读取速度，但是会占用更多内存
# 内存较小的电脑可以设置为2或者0
#------------------------------------------------------------------#
# 使用多线程的方式进行dataset的构建，一般设置0、2、4
num_workers = 4
#----------------------------------------------------#
# 获得图片路径和标签
#----------------------------------------------------#
# 训练的信息主要是图片路径和x1,y1,x2,y2,class_num，...的构建
# 看一条：C:\Users\xxx\yolo3-pytorch-master\VOCdevkit/VOC2007/JPEGImages/000007.jpg 141,50,500,330,6
train_annotation_path = '2007_train.txt'
val_annotation_path = '2007_val.txt'
#----------------------------------------------------#
# 获取classes和anchor
#----------------------------------------------------#
# 获取类信息和anchors信息，在预测的时候介绍过，忘记可以去YOLO类中去查看
class_names, num_classes = get_classes(classes_path)
anchors, num_anchors = get_anchors(anchors_path)
#------------------------------------------------------#
# 创建yolo模型
#------------------------------------------------------#
# 创建我们的模型
model = YoloBody(anchors_mask, num_classes, pretrained=pretrained)
# 假设我们的模型中没有预训练参数的时候我们需要自行进行参数初始化，详细见下解读
if not pretrained:
weights_init(model)
# 我们加载模型参数(值得学习，固定写法)
# 我们先进行判断是否加载主干网络预训练参数，不加载自行初始化，加载使用预训练参数；
# 然后我们在判断时候加载模型参数，若加载模型参数的话，即使设置了预训练参数也会被覆盖掉，
# 没有加载模型参数的话，按照上面的参数加载方式进行初始化
# 重要程度：model_path >> pretrained >> weight_init
if model_path != '':
#------------------------------------------------------#
# 权值文件请看README，百度网盘下载
#------------------------------------------------------#
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)
# yolo的loss类的加载，在(七) Loss篇进行讲解
yolo_loss = YOLOLoss(anchors, num_classes, input_shape, Cuda, anchors_mask)
# 在log中写入我们的每个epoch的loss的信息，我们可以在可视化部分观察数据变化，看是否收敛
loss_history = LossHistory("logs/", model, input_shape=input_shape)
# 模型训练
model_train = model.train()
if Cuda:
model_train = torch.nn.DataParallel(model)
cudnn.benchmark = True
model_train = model_train.cuda()
#---------------------------#
# 读取数据集对应的txt
#---------------------------#
# 读取txt文件主要是为了构建dataset和dataloader
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总训练世代
# 提示OOM或者显存不足请调小Batch_size
#------------------------------------------------------#
if True:
UnFreeze_flag = False
#------------------------------------#
# 冻结一定部分训练
#------------------------------------#
# 进行冻结训练的话，看下面的循环，我们是针对backbone也就是darknet53的参数，将他们都设置成不需要梯度更新的状态
if Freeze_Train:
for param in model.backbone.parameters():
param.requires_grad = False
#-------------------------------------------------------------------#
# 如果不冻结训练的话，直接设置batch_size为Unfreeze_batch_size
#-------------------------------------------------------------------#
# 我们的冻结状态的batch_size和非冻结状态的batch_size不同
batch_size = Freeze_batch_size if Freeze_Train else Unfreeze_batch_size
#-------------------------------------------------------------------#
# 判断当前batch_size与64的差别，自适应调整学习率
#-------------------------------------------------------------------#
# 自动调整学习率
nbs = 64
Init_lr_fit = max(batch_size / nbs * Init_lr, 1e-4)
Min_lr_fit = max(batch_size / nbs * Min_lr, 1e-6)
#---------------------------------------#
# 根据optimizer_type选择优化器
#---------------------------------------#
# 在这个地方实际上是我们构造我们的optimizer，我们的参数中分成3种，一种weight我们使用p1存储，一种是bias我们使用b2进行存储，还有一种就是BN层，没有bias，我们使用p0对其weight进行存储
pg0, pg1, pg2 = [], [], []
for k, v in model.named_modules():
if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias)
if isinstance(v, nn.BatchNorm2d) or "bn" in k:
pg0.append(v.weight)
elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight)
# 我们使用optimizer_type进行optimizer的选用，字典选值的方式可以学习，然后将
optimizer = {
'adam' : optim.Adam(pg0, Init_lr_fit, betas = (momentum, 0.999)),
'sgd' : optim.SGD(pg0, Init_lr_fit, momentum = momentum, nesterov=True)
}[optimizer_type]
# 在optimizer中更新参数
optimizer.add_param_group({"params": pg1, "weight_decay": weight_decay})
optimizer.add_param_group({"params": pg2})
#---------------------------------------#
# 获得学习率下降的公式
#---------------------------------------#
lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
#---------------------------------------#
# 判断每一个世代的长度
#---------------------------------------#
# 获取一个epoch中我们需要进行多少轮的step,整除部分保留，余数部分不够一个step的舍弃
# train_step
epoch_step = num_train // batch_size
# val_step
epoch_step_val = num_val // batch_size
# 数据集不能满足一个step 就直接罢工，抛出异常
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。")
#---------------------------------------#
# 构建数据集加载器。
#---------------------------------------#
# 构建Dataset和DataLoader的加载
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)
#---------------------------------------#
# 开始模型训练
#---------------------------------------#
# 因为我们的非冻结Epoch就是要训练的轮次，所以直接循环从初始epoch到非冻结epoch
for epoch in range(Init_Epoch, UnFreeze_Epoch):
#---------------------------------------#
# 如果模型有冻结学习部分
# 则解冻，并设置参数
#---------------------------------------#
# 前期是冻结训练，后期是非冻结训练，batch_size不同，需要进行切换
if epoch >= Freeze_Epoch and not UnFreeze_flag and Freeze_Train:
batch_size = Unfreeze_batch_size
#-------------------------------------------------------------------#
# 判断当前batch_size与64的差别，自适应调整学习率
#-------------------------------------------------------------------#
# 因为batch_size变化，自适应调整学习率也要变化
nbs = 64
Init_lr_fit = max(batch_size / nbs * Init_lr, 1e-4)
Min_lr_fit = max(batch_size / nbs * Min_lr, 1e-6)
#---------------------------------------#
# 获得学习率下降的公式
#---------------------------------------#
# 获取学习率下降的function
lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
# 因为我们进入了非冻结阶段的训练，所以要将参数的需要梯度下降的属性设置成True
for param in model.backbone.parameters():
param.requires_grad = True
# batch_size发生变化，step的数量也要发生变化
epoch_step = num_train // batch_size
epoch_step_val = num_val // batch_size
# 数据集不能满足一个step 就直接罢工，抛出异常
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。")
# 同理，batch_size的变化，我们的dataloader的加载也发生变化，重新加载
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)
# 非冻结训练标志设置成True
UnFreeze_flag = True
# 设置优化器的学习率，参数(优化器，学习率优化方法，epoch)
set_optimizer_lr(optimizer, lr_scheduler_func, epoch)
# 真正的训练方法(核心)
fit_one_epoch(model_train, model, yolo_loss, loss_history, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val, UnFreeze_Epoch, Cuda, save_period)
# 日志文件管理器关闭
loss_history.writer.close()
复制代码- 对上面调用到的函数进行详解
- weights_init解读
def weights_init(net, init_type='normal', init_gain = 0.02): # 初始化方法 def init_func(m): classname = m.__class__.__name__ # 假设我们的属性中存在weight的话并且是卷积层的话 if hasattr(m, 'weight') and classname.find('Conv') != -1: # 使用给定方差均值的正太分布初始化 if init_type == 'normal': torch.nn.init.normal_(m.weight.data, 0.0, init_gain) # xavier初始化 elif init_type == 'xavier': torch.nn.init.xavier_normal_(m.weight.data, gain=init_gain) # 凯明初始化 elif init_type == 'kaiming': torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') # orthogonal初始化 elif init_type == 'orthogonal': torch.nn.init.orthogonal_(m.weight.data, gain=init_gain) # 只能在上面提供的方法中选择一种初始化，否则引发exception else: raise NotImplementedError('initialization method [%s] is not implemented' % init_type) # 要是BN层的话，我们使用正态分布初始化和常数初始化 elif classname.find('BatchNorm2d') != -1: torch.nn.init.normal_(m.weight.data, 1.0, 0.02) torch.nn.init.constant_(m.bias.data, 0.0) # 初始化完成，应用到网络中 print('initialize network with %s type' % init_type) net.apply(init_func) 复制代码- YoloDataset的解读
- 对于YoloDataset的解读：juejin.cn/post/707962… 详细的分析都放在这个部分
- LossHistory的解读
- 对于LossHistory的解读：juejin.cn/post/707962… 详细的分析都放在这个部分
- get_lr_scheduler的解读
# lr_decay_type = "cos",lr = max(batch_size / nbs * Init_lr, 1e-4) # min_lr = max(batch_size / nbs * Min_lr, 1e-6) total_iters总的迭代epoch数 # 延后进行分析 def get_lr_scheduler(lr_decay_type, lr, min_lr, total_iters, warmup_iters_ratio = 0.1, warmup_lr_ratio = 0.1, no_aug_iter_ratio = 0.3, step_num = 10): # yolox的 def yolox_warm_cos_lr(lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter, iters): if iters <= warmup_total_iters: # lr = (lr - warmup_lr_start) * iters / float(warmup_total_iters) + warmup_lr_start lr = (lr - warmup_lr_start) * pow(iters / float(warmup_total_iters), 2) + warmup_lr_start elif iters >= total_iters - no_aug_iter: lr = min_lr else: lr = min_lr + 0.5 * (lr - min_lr) * ( 1.0 + math.cos(math.pi* (iters - warmup_total_iters) / (total_iters - warmup_total_iters - no_aug_iter)) ) return lr def step_lr(lr, decay_rate, step_size, iters): if step_size < 1: raise ValueError("step_size must above 1.") n = iters // step_size out_lr = lr * decay_rate ** n return out_lr if lr_decay_type == "cos": warmup_total_iters = min(max(warmup_iters_ratio * total_iters, 1), 3) warmup_lr_start = max(warmup_lr_ratio * lr, 1e-6) no_aug_iter = min(max(no_aug_iter_ratio * total_iters, 1), 15) func = partial(yolox_warm_cos_lr ,lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter) else: decay_rate = (min_lr / lr) ** (1 / (step_num - 1)) step_size = total_iters / step_num func = partial(step_lr, lr, decay_rate, step_size) return func 复制代码- set_optimizer_lr的解读
# 设置指定epoch的学习率 def set_optimizer_lr(optimizer, lr_scheduler_func, epoch): # 使用指定的lr_scheduler学习率下降方式完成学习率的下降 lr = lr_scheduler_func(epoch) # 将学习率更新到优化器参数中 for param_group in optimizer.param_groups: param_group['lr'] = lr 复制代码- fit_one_epoch的解读(
重点)- 写法比较固定，可以学习的地方是定制化tqdm的写法
- 重点需要解析的方法：yolo_loss() loss的计算详见：
def fit_one_epoch(model_train, model, yolo_loss, loss_history, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val, Epoch, cuda, save_period): # 定义训练loss和val_loss loss = 0 val_loss = 0 # 开始训练的标志 model_train.train() print('Start Train') # 我们可以对我们的tqdm的使用进行定制输出的样式 with tqdm(total=epoch_step,desc=f'Epoch {epoch +1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar: # 我们开始循环遍历我们的dataloader for iteration, batch in enumerate(gen): # 假设迭代的次数超过了我们之前规定的一个epoch的step的数量 我们就直接跳过 if iteration >= epoch_step: break # 取得图像数据和label数据 images, targets = batch[0], batch[1] with torch.no_grad(): # 假设我们使用gpu,我们将数据转化成tensor后放到cuda上 if cuda: images = torch.from_numpy(images).type(torch.FloatTensor).cuda() targets = [torch.from_numpy(ann).type(torch.FloatTensor).cuda() for ann in targets] # 没有使用gpu,数据转化成tensor即可 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) # 记录每个step的loss loss_value_all = 0 #----------------------# # 计算损失 #----------------------# # 因为我们是有三个输出[[bs,13,13,75],[bs,26,26,75],[bs,52,52,75]]，所以我们循环遍历每一个特征图的输出，和真实的label之间计算损失 for l in range(len(outputs)): # 计算损失的核心就是：yolo_loss loss_item = yolo_loss(l, outputs[l], targets) loss_value_all += loss_item # 得到这个step的损失 loss_value = loss_value_all #----------------------# # 反向传播 #----------------------# # 固定写法，反向传播和优化器优化 loss_value.backward() optimizer.step() loss += loss_value.item() # 对我们的定制化的tqdm进行参数设置，主要输出loss,和学习率 pbar.set_postfix(**{'loss' : loss / (iteration + 1), 'lr' : get_lr(optimizer)}) # tqdm的更新 pbar.update(1) print('Finish Train') # 一个epoch训练完成之后，我们进行eval 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 #----------------------# # 计算损失 #----------------------# for l in range(len(outputs)): loss_item = yolo_loss(l, outputs[l], targets) loss_value_all += loss_item loss_value = loss_value_all val_loss += loss_value.item() pbar.set_postfix(**{'val_loss': val_loss / (iteration + 1)}) pbar.update(1) print('Finish Validation') # 然后将我们的loss信息，记录在我们的log中 loss_history.append_loss(epoch + 1, 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)) # 判断是不是到了最后一个epoch或者当前的epoch对我们的保存周期参数可以整除的话，我们就保留模型 if (epoch + 1) % save_period == 0 or epoch + 1 == Epoch: torch.save(model.state_dict(), 'logs/ep%03d-loss%.3f-val_loss%.3f.pth' % (epoch + 1, loss / epoch_step, val_loss / epoch_step_val)) 复制代码