Tensorflow版本yolo v3源码阅读笔记(4)

在前面3篇文章中，我们已经了解了yolov3的模型以及损失函数的计算方法，在上一篇中又了解到了数据集的制作方法。

在这一篇中，我们就开始来分析训练yolov3神经网络的源代码吧。

说起来，训练神经网络的方法就这么几步：

-1. 输入图片数据，求出预测值。
-2. 根据预测值、真实值和损失函数来计算损失值。
-3. 计算梯度信息，利用梯度下降方法来优化神经网络。
-4. 循环上面的步骤若干次。

下面让我们来分析代码的具体实现过程吧。

导入需要的库

import os
import time
import shutil
import numpy as np
import tensorflow as tf
import core.utils as utils
from tqdm import tqdm#python中的进度条库
from core.dataset import Dataset
from core.yolov3 import YOLOv3, decode, compute_loss
from core.config import cfg

导入数据集

trainset = Dataset('train')

一些参数

logdir = "./data/log"#日志文件的目录
steps_per_epoch = len(trainset)#每个回合训练的步数(即数据集的长度)
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)#用来记录训练到第几步了
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch#训练过程中，在warmup_steps步数之前学习率是一种变化趋势，在之后又是另一种变化趋势
total_steps = cfg.TRAIN.EPOCHS * steps_per_epoch#神经网络训练的总的步数

构建网络

#神经网络的输入张量
input_tensor = tf.keras.layers.Input([416, 416, 3])
#神经网络的输出特征图
conv_tensors = YOLOv3(input_tensor)

output_tensors = []
for i, conv_tensor in enumerate(conv_tensors):
    pred_tensor = decode(conv_tensor, i)#处理神经网络的输出特征图
    output_tensors.append(conv_tensor)#保存未处理过的特征图信息
    output_tensors.append(pred_tensor)#保存处理过的特征图信息

#定义模型，这个模型非常简洁，一个输入input_tensor，一个输出output_tensors
model = tf.keras.Model(input_tensor, output_tensors)
#定义优化器
optimizer = tf.keras.optimizers.Adam()
#如果日志文件存在，就循环删除这个目录，然后重新建一个日志目录
if os.path.exists(logdir): shutil.rmtree(logdir)
writer = tf.summary.create_file_writer(logdir)

定义训练函数

def train_step(image_data, target):
    with tf.GradientTape() as tape:#构建一个梯度环境
        pred_result = model(image_data, training=True)#得到3个尺度的特征图信息
        #初始化框回归损失、置信度损失、分类损失
        giou_loss=conf_loss=prob_loss=0

        # 循环遍历3个尺度的特征图
        for i in range(3):
            #取出未处理的特征图和处理过的特征图
            conv, pred = pred_result[i*2], pred_result[i*2+1]
            #把这两个特征图信息送入损失函数去计算损失，注意，*target[i]这个参数有2个信息，一个是真实框的特征图信息，一个是这个尺度下的所有真实框
            loss_items = compute_loss(pred, conv, *target[i], i)
            giou_loss += loss_items[0]#框回归损失
            conf_loss += loss_items[1]#置信度损失
            prob_loss += loss_items[2]#类别损失

        #计算总的损失值
        total_loss = giou_loss + conf_loss + prob_loss

        #根据损失值计算出所有参数的梯度信息
        gradients = tape.gradient(total_loss, model.trainable_variables)
        #通过优化器对所有参数进行优化
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        tf.print("=> STEP %4d   lr: %.6f   giou_loss: %4.2f   conf_loss: %4.2f   "
                 "prob_loss: %4.2f   total_loss: %4.2f" %(global_steps, optimizer.lr.numpy(),
                                                          giou_loss, conf_loss,
                                                          prob_loss, total_loss))
        # 记录当前训练的总步数
        global_steps.assign_add(1)
        # 根据情况对学习率做动态改变
        if global_steps < warmup_steps:
            lr = global_steps / warmup_steps *cfg.TRAIN.LR_INIT
        else:
            lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (
                (1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi))
            )
        #重新设置优化器的学习率
        optimizer.lr.assign(lr.numpy())

        # 将训练过程中的学习率、步数、各种损失值写入日志目录，方便可视化
        with writer.as_default():
            tf.summary.scalar("lr", optimizer.lr, step=global_steps)
            tf.summary.scalar("loss/total_loss", total_loss, step=global_steps)
            tf.summary.scalar("loss/giou_loss", giou_loss, step=global_steps)
            tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps)
            tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps)
        #写文件的一个常用操作，把缓存中的数据强行写入文件
        writer.flush()

最后一步

#神经网络训练cfg.TRAIN.EPOCHS个回合
for epoch in range(cfg.TRAIN.EPOCHS):
    for image_data, target in trainset:
        train_step(image_data, target)
    #每个回合训练完成后，保存yolov3的参数信息，以后可以直接加载训练好的参数，就不用从零开始一步步训练了。
    model.save_weights("./yolov3")

以上就是训练神经网络的具体过程，借助tensorflow这个框架，训练过程非常的清晰明了。

至此，yolov3神经网络总算是训练好了，下一步我们就可以来测试一下这个训练好的神经网络效果怎么样，但是别急，在测试之前，我们还要搞清楚测试函数中出现的几个函数。