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 = 0for i in range(3): # 3个 feature mapconv, 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 加 1if global_steps < warmup_steps:lr = global_steps / warmup_steps *cfg.TRAIN.LR_INITelse: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 cfgtrainset = 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_epochinput_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 processfor 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_lossgradients = 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 rateglobal_steps.assign_add(1)if global_steps < warmup_steps:lr = global_steps / warmup_steps *cfg.TRAIN.LR_INITelse: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 datawith 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")