Tensorflow2.0 实现 YOLOv3（七）：train.py

文章说明

本系列文章旨在对 Github 上 malin9402 提供的代码进行说明，在这篇文章中，我们会对 YOLOv3 项目中的 train.py 文件进行说明。

如果只是想运行 Github 上的代码，可以参考对 YOLOv3 代码的说明一文。

导入需要的库

import os
import time
import shutil
import numpy as np
import tensorflow as tf
from tqdm import tqdm  # 显示进度条功能

import core.utils as utils
from core.config import cfg
from core.yolov3 import YOLOv3, decode, compute_loss
from core.dataset import Dataset

构建数据集

trainset = Dataset('train')

设置参数

• steps_per_epoch：每个 epoch 要训练 len(trainset) 次；
• global_steps：是一个变量，用来记录现在是第几次训练；
• warmup_steps：在训练过程中我们希望在训练 warmup_steps 次之前学习率有一种变化趋势，在之后有另一种变化趋势；
• total_steps：训练总次数。

steps_per_epoch = len(trainset)
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = cfg.TRAIN.EPOCHS * steps_per_epoch

构建网络

# 确定输入张量的shape
input_tensor = tf.keras.layers.Input([416, 416, 3])
# 确定输出张量
conv_tensors = YOLOv3(input_tensor)  # 3个张量（feature map）
output_tensors = []
for i, conv_tensor in enumerate(conv_tensors):
    pred_tensor = decode(conv_tensor, i)
    output_tensors.append(conv_tensor)   # 未处理的[batch_size, output_size, output_size,255]，表示在 feature map 上的检测框信息。
    output_tensors.append(pred_tensor)   # 处理的[batch_size, output_size, output_size,255]，表示在原始图像上的检测框信息。
# 构建模型
model = tf.keras.Model(input_tensor, output_tensors)

初始化优化器

optimizer = tf.keras.optimizers.Adam()

设置保存文件

logdir = "./data/log"
# 删除logdir路径下的文件
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) # 将图片输入模型
        giou_loss = conf_loss = prob_loss = 0
        for i in range(3): # 3个 feature map
            conv, pred = pred_result[i*2], pred_result[i*2+1] # 包括未经解码处理的输出和已解码处理的输出
            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) # global_steps 加 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()

这里学习率的设置在论文Bag of Tricks for Image Classification with Convolutional Neural Networks中被提出，称为 consine decay stage，其变化趋势为：

训练

for epoch in range(cfg.TRAIN.EPOCHS):
    for image_data, target in trainset:
        train_step(image_data, target)
    model.save_weights("./yolov3")

注意，如果用 model.save("./yolov3") 则保存全部模型，加载时用 model = load_model(‘yolov3.h5’) 即可。这里用的 model.save_weights("./yolov3") 只保存模型的参数，不保存结构，加载时用 model.load_weights(‘yolov3.h5’) 之前，得先构建一个关于 model 的参数结构才能加载。

完整代码

import os
import time
import shutil
import numpy as np
import tensorflow as tf
import core.utils as utils
from tqdm import tqdm
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
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)

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)
        giou_loss=conf_loss=prob_loss=0

        # optimizing process
        for i in range(3):
            conv, pred = pred_result[i*2], pred_result[i*2+1]
            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))
        # update learning rate
        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())

        # writing summary data
        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()


for epoch in range(cfg.TRAIN.EPOCHS):
    for image_data, target in trainset:
        train_step(image_data, target)
    model.save_weights("./yolov3")