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