yolov3代码详解（六）

train.py

from __future__ import division

from models import *
from utils.logger import *
from utils.utils import *
from utils.datasets import *
from utils.parse_config import *
from test import evaluate

from terminaltables import AsciiTable

import os
import sys
import time
import datetime
import argparse

import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
from torch.autograd import Variable
import torch.optim as optim
##########################################################################
#训练
##########################################################################

"""
train.py的主要工作流程
1.解析输入的各种参数，如没有则使用默认参数
2.打印各种参数
3.初始化日志
4.获得train_path、valid_path和class_names的文件路径
5.创建model，随机初始化权重，也可以加载预训练的参数
6.加载训练图像
7.选择优化器
8.开始epoch轮，反向传播
9.开始训练batch_i批
10.每累积gradient_accumulations批，进行一次梯度下降
11.记录训练日志
12.每训练完evaluation_interval轮输出一次评估结果
13.每训练完opt.checkpoint_interval轮，保存一次checkpoints
"""


if __name__ == "__main__":

    ###############################################################################################
    ################                              初始化                            ################   
    #1.解析输入的各种参数，如没有则使用默认参数
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=100, help="number of epochs")
    parser.add_argument("--batch_size", type=int, default=8, help="size of each image batch")
    #梯度累积  台阶前坡度累计数
    parser.add_argument("--gradient_accumulations", type=int, default=2, help="number of gradient accums before step") 
    parser.add_argument("--model_def", type=str, default="config/yolov3.cfg", help="path to model definition file")
    parser.add_argument("--data_config", type=str, default="config/coco.data", help="path to data config file")
    parser.add_argument("--pretrained_weights", type=str, help="if specified starts from checkpoint model")
    #批处理生成期间要使用的cpu线程数
    parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
    parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
    parser.add_argument("--checkpoint_interval", type=int, default=1, help="interval between saving model weights")
    #验证集上的区间求值
    parser.add_argument("--evaluation_interval", type=int, default=1, help="interval evaluations on validation set")
    #如果为真，则每十批计算一次映射
    parser.add_argument("--compute_map", default=False, help="if True computes mAP every tenth batch")
    #允许多规模培训
    parser.add_argument("--multiscale_training", default=True, help="allow for multi-scale training")
    opt = parser.parse_args()
    #2.打印各种参数
    print(opt)

    #3.初始化日志
    logger = Logger("logs")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    os.makedirs("output", exist_ok=True)
    os.makedirs("checkpoints", exist_ok=True)


    ###############################################################################################
    ################                              加载网络                          ################  
    #4.获得train_path、valid_path和class_names的文件路径
    # Get data configuration  解析数据配置
    data_config = parse_data_config(opt.data_config)
    train_path = data_config["train"]
    valid_path = data_config["valid"]
    class_names = load_classes(data_config["names"])

    #5.创建model，随机初始化权重，也可以加载预训练的参数
    # Initiate model  启动模型
    model = Darknet(opt.model_def).to(device)  
    model.apply(weights_init_normal)  #方法应用把方法应用于每一个module,这里意思是进行初始化  在utils中

    # If specified we start from checkpoint         ？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？？
    if opt.pretrained_weights:
        if opt.pretrained_weights.endswith(".pth"):
            model.load_state_dict(torch.load(opt.pretrained_weights))
        else:
            model.load_darknet_weights(opt.pretrained_weights)


    ###############################################################################################
    ################                           放入dataloader                       ################  
    # Get dataloader
    #6.加载训练图像   
    dataset = ListDataset(train_path, augment=True, multiscale=opt.multiscale_training)
    #DataLoader的collate_fn参数，实现自定义的batch输出
    #- shuffle：设置为True的时候，每个世代都会打乱数据集 
    #- collate_fn：如何取样本的，我们可以定义自己的函数来准确地实现想要的功能 
    #- drop_last：告诉如何处理数据集长度除于batch_size余下的数据。True就抛弃，否则保留
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=opt.batch_size,
        shuffle=True,
        num_workers=opt.n_cpu,
        pin_memory=True,
        collate_fn=dataset.collate_fn,
    )

    #使用优化器  Adam 是一种可以替代传统随机梯度下降过程的一阶优化算法，它能基于训练数据迭代地更新神经网络权重。？？？？？？？？？？？？？？？
     #7.选择优化器
    optimizer = torch.optim.Adam(model.parameters())

    metrics = [
        "grid_size",
        "loss",
        "x",
        "y",
        "w",
        "h",
        "conf",
        "cls",
        "cls_acc",
        "recall50",
        "recall75",
        "precision",
        "conf_obj",
        "conf_noobj",
    ]


    ###############################################################################################
    ################                             开始迭代                           ################ 
    #8.开始训练epoch轮，反向传播
    #训练epoch轮，每轮有batch_i批，每批有imgs个图，每个图有targets个目标
    for epoch in range(opt.epochs):
        model.train()
        start_time = time.time()
        #这里主要要参考ListDataset中的__getitem__和DataLoader中的collate_fn设置。
        #9.开始训练batch_i批
        for batch_i, (_, imgs, targets) in enumerate(dataloader):
            batches_done = len(dataloader) * epoch + batch_i  #训练完了几批，dataloader一次读入batch_size批

            imgs = Variable(imgs.to(device))   #得到图片
            targets = Variable(targets.to(device), requires_grad=False)   #得到标签

           """
           这里进行计算loss。其实这个loss的计算是在yolo层计算的，其实不难理解，yolo层是负责目标检测的层，
           需要输出目标的类别、坐标、大小，所以会在这一层进行loss计算。
           这个代码可以从Darknet类的前向通路中发现（在训练的时候targets是有值的，不等于None）
           yolo层的具体实现是在YOLOLayer中，可查看其forward函数得知loss计算过程，代码（YOLOLayer部分）
           """
            loss, outputs = model(imgs, targets)   #imgs, targets喂给model,得到loss, outputs
            loss.backward()             #反向传播loss

            #10.每累积gradient_accumulations批，进行一次梯度下降
            if batches_done % opt.gradient_accumulations:
                # Accumulates gradient before each step  在每一步之前累积坡度
                optimizer.step()   #只有用了optimizer.step()，模型才会更新
                optimizer.zero_grad()


            #11.记录训练日志
            # ----------------
            #   Log progress
            # ----------------

            log_str = "\n---- [Epoch %d/%d, Batch %d/%d] ----\n" % (epoch, opt.epochs, batch_i, len(dataloader))

            metric_table = [["Metrics", *[f"YOLO Layer {i}" for i in range(len(model.yolo_layers))]]]
            #[["Metrics","YOLO Layer 1","YOLO Layer 2","YOLO Layer 3"]]

            # Log metrics at each YOLO layer
            for i, metric in enumerate(metrics):
                formats = {m: "%.6f" for m in metrics}   #刚开始其他指标类型都设为%.6f
                formats["grid_size"] = "%2d"             #更新grid_size的类型，为%2d
                formats["cls_acc"] = "%.2f%%"            #更新cls_acc的类型，为%.2f%%
                #例：print("%.2f%%"%23)  输出为23.00%
                #将值保存到metric_table中吗
                row_metrics = [formats[metric] % yolo.metrics.get(metric, 0) for yolo in model.yolo_layers]   #得到metric（表示"loss""x""y"等）标签对应的值                    
                metric_table += [[metric, *row_metrics]]    #  *args可以传入多个参数
                #例如是 grid_size的话   结果就是[grid_size,13,26,52]


                # Tensorboard logging
                tensorboard_log = []
                for j, yolo in enumerate(model.yolo_layers):
                    for name, metric in yolo.metrics.items():
                        if name != "grid_size":
                            tensorboard_log += [(f"{name}_{j+1}", metric)]
                tensorboard_log += [("loss", loss.item())]
                logger.list_of_scalars_summary(tensorboard_log, batches_done)

            log_str += AsciiTable(metric_table).table
            log_str += f"\nTotal loss {loss.item()}"

            # Determine approximate time left for epoch  确定epoch的大致剩余时间
            epoch_batches_left = len(dataloader) - (batch_i + 1)
            time_left = datetime.timedelta(seconds=epoch_batches_left * (time.time() - start_time) / (batch_i + 1))
            log_str += f"\n---- ETA {time_left}"

            print(log_str)

            model.seen += imgs.size(0)


        #12.每个evaluation_interval轮输出一次评估结果
        if epoch % opt.evaluation_interval == 0:
            print("\n---- Evaluating Model ----")
            # Evaluate the model on the validation set
            precision, recall, AP, f1, ap_class = evaluate(
                model,
                path=valid_path,
                iou_thres=0.5,
                conf_thres=0.5,
                nms_thres=0.5,
                img_size=opt.img_size,
                batch_size=8,
            )
            evaluation_metrics = [
                ("val_precision", precision.mean()),
                ("val_recall", recall.mean()),
                ("val_mAP", AP.mean()),
                ("val_f1", f1.mean()),
            ]
            logger.list_of_scalars_summary(evaluation_metrics, epoch)

            # Print class APs and mAP
            ap_table = [["Index", "Class name", "AP"]]
            for i, c in enumerate(ap_class):
                ap_table += [[c, class_names[c], "%.5f" % AP[i]]]
            print(AsciiTable(ap_table).table)
            print(f"---- mAP {AP.mean()}")

        if epoch % opt.checkpoint_interval == 0:
            torch.save(model.state_dict(), f"checkpoints/yolov3_ckpt_%d.pth" % epoch)