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")