train.py

#!/usr/bin/env python
import os
import json
import torch
import numpy as np
import queue
import pprint
import random
import argparse
import importlib
import threading
import traceback
import torch.distributed as dist
import torch.multiprocessing as mp

from tqdm import tqdm
from torch.multiprocessing import Process, Queue, Pool

from core.dbs import datasets
from core.utils import stdout_to_tqdm
from core.config import SystemConfig
from core.sample import data_sampling_func
from core.nnet.py_factory import NetworkFactory

torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

# 命令行参数
def parse_args():
    parser = argparse.ArgumentParser(description="Training Script")
    parser.add_argument("cfg_file", help="config file", type=str) #训练用的配置文件
    parser.add_argument("--iter", dest="start_iter",
                        help="train at iteration i",
                        default=0, type=int)  #指定训练从第i次迭代开始
    parser.add_argument("--workers", default=4, type=int)
    parser.add_argument("--initialize", action="store_true")

    parser.add_argument("--distributed", action="store_true")  # 分布式训练
    parser.add_argument("--world-size", default=-1, type=int,
                        help="number of nodes of distributed training")  # 分布式节点的数量
    parser.add_argument("--rank", default=0, type=int,
                        help="node rank for distributed training")  # 分布式训练节点的等级
    parser.add_argument("--dist-url", default=None, type=str,
                        help="url used to set up distributed training")
    parser.add_argument("--dist-backend", default="nccl", type=str)

    args = parser.parse_args()
    return args


def prefetch_data(system_config, db, queue, sample_data, data_aug):
    ind = 0
    print("start prefetching data...")
    # 用pid[进程id]做随机数种子
    np.random.seed(os.getpid())
    while True:
        try:
            data, ind = sample_data(system_config, db, ind, data_aug=data_aug)
            queue.put(data)
        except Exception as e:
            traceback.print_exc()
            raise e


def _pin_memory(ts):
    if type(ts) is list:
        return [t.pin_memory() for t in ts]
    return ts.pin_memory()


def pin_memory(data_queue, pinned_data_queue, sema):
    while True:
        data = data_queue.get()

        data["xs"] = [_pin_memory(x) for x in data["xs"]]
        data["ys"] = [_pin_memory(y) for y in data["ys"]]

        pinned_data_queue.put(data)

        if sema.acquire(blocking=False):
            return


def init_parallel_jobs(system_config, dbs, queue, fn, data_aug):
    tasks = [Process(target=prefetch_data, args=(system_config, db, queue, fn, data_aug)) for db in dbs]
    for task in tasks:
        task.daemon = True
        task.start()
    return tasks


def terminate_tasks(tasks):
    for task in tasks:
        task.terminate()


# 训练函数（训练数据，验证数据，系统参数，模型，参数）
def train(training_dbs, validation_db, system_config, model, args):
    # reading arguments from command
    start_iter = args.start_iter
    distributed = args.distributed
    world_size = args.world_size
    initialize = args.initialize
    gpu = args.gpu
    rank = args.rank

    # reading arguments from json file
    batch_size = system_config.batch_size
    learning_rate = system_config.learning_rate
    max_iteration = system_config.max_iter
    pretrained_model = system_config.pretrain
    stepsize = system_config.stepsize
    snapshot = system_config.snapshot
    val_iter = system_config.val_iter
    display = system_config.display
    decay_rate = system_config.decay_rate
    stepsize = system_config.stepsize

    print("Process {}: building model...".format(rank))
    nnet = NetworkFactory(system_config, model, distributed=distributed, gpu=gpu)  # 构建模型
    if initialize:
        nnet.save_params(0)
        exit(0)

    # 开4个队列去存数据
    # queues storing data for training
    training_queue = Queue(system_config.prefetch_size)
    validation_queue = Queue(5)

    # queues storing pinned data for training
    pinned_training_queue = queue.Queue(system_config.prefetch_size)
    pinned_validation_queue = queue.Queue(5)

    # allocating resources for parallel reading
    training_tasks = init_parallel_jobs(system_config, training_dbs, training_queue, data_sampling_func, True)
    if val_iter:
        validation_tasks = init_parallel_jobs(system_config, [validation_db], validation_queue, data_sampling_func,
                                              False)

    training_pin_semaphore = threading.Semaphore()
    validation_pin_semaphore = threading.Semaphore()
    training_pin_semaphore.acquire()
    validation_pin_semaphore.acquire()

    training_pin_args = (training_queue, pinned_training_queue, training_pin_semaphore)
    training_pin_thread = threading.Thread(target=pin_memory, args=training_pin_args)
    training_pin_thread.daemon = True
    training_pin_thread.start()

    validation_pin_args = (validation_queue, pinned_validation_queue, validation_pin_semaphore)
    validation_pin_thread = threading.Thread(target=pin_memory, args=validation_pin_args)
    validation_pin_thread.daemon = True
    validation_pin_thread.start()

    # 看是否有预训练的模型
    if pretrained_model is not None:
        if not os.path.exists(pretrained_model):
            raise ValueError("pretrained model does not exist")
        print("Process {}: loading from pretrained model".format(rank))
        nnet.load_pretrained_params(pretrained_model)

    # 看有没有开始的迭代器
    if start_iter:
        nnet.load_params(start_iter)
        learning_rate /= (decay_rate ** (start_iter // stepsize))
        nnet.set_lr(learning_rate)
        print("Process {}: training starts from iteration {} with learning_rate {}".format(rank, start_iter + 1,
                                                                                           learning_rate))
    else:
        nnet.set_lr(learning_rate)

    # 训练模型
    if rank == 0:
        print("training start...")
    nnet.cuda()
    nnet.train_mode()

    #
    with stdout_to_tqdm() as save_stdout:
        # Tqdm 是一个快速，可扩展的Python进度条，可以在 Python 长循环中添加一个进度提示信息
        for iteration in tqdm(range(start_iter + 1, max_iteration + 1), file=save_stdout, ncols=80):
            training = pinned_training_queue.get(block=True)
            training_loss = nnet.train(**training)

            # 如果设置了display的步长，我们在步长整数倍时展示损失函数的值
            if display and iteration % display == 0:
                print("Process {}: training loss at iteration {}: {}".format(rank, iteration, training_loss.item()))
            del training_loss

            # 如果设置了变量迭代器步长[这边是验证集了]
            if val_iter and validation_db.db_inds.size and iteration % val_iter == 0:
                nnet.eval_mode()
                validation = pinned_validation_queue.get(block=True)
                validation_loss = nnet.validate(**validation)
                print("Process {}: validation loss at iteration {}: {}".format(rank, iteration, validation_loss.item()))
                nnet.train_mode()

            # 快照步长
            if iteration % snapshot == 0 and rank == 0:
                nnet.save_params(iteration)

            # 学习率更新步长
            if iteration % stepsize == 0:
                learning_rate /= decay_rate
                nnet.set_lr(learning_rate)

    # sending signal to kill the thread
    # 杀掉进程的消息
    training_pin_semaphore.release()
    validation_pin_semaphore.release()

    # terminating data fetching processes
    # 结束任务的消息
    terminate_tasks(training_tasks)
    terminate_tasks(validation_tasks)


# 主函数main
def main(gpu, ngpus_per_node, args):
    # 将gpu复制到args里
    args.gpu = gpu
    if args.distributed:
        args.rank = args.rank * ngpus_per_node + gpu
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                world_size=args.world_size, rank=args.rank)

    rank = args.rank

    # 读取配置文件
    # cfg_file是必须传入的要训练模型名字
    cfg_file = os.path.join("./configs", args.cfg_file + ".json") #配置文件的路径！
    with open(cfg_file, "r") as f:
        config = json.load(f)

    # 上面的读取的单独模型配置的"system"里添加一项
    config["system"]["snapshot_name"] = args.cfg_file

    # 更新参数配置
    system_config = SystemConfig().update_config(config["system"])

    # 模型的名字作为模型导入
    model_file = "core.models.{}".format(args.cfg_file)  # 模型路径
    model_file = importlib.import_module(model_file)     # importlib.import_module，导入模型。
    model = model_file.model()

    # 从i系统配置里取出的参数
    train_split = system_config.train_split
    val_split = system_config.val_split

    print("Process {}: loading all datasets...".format(rank))

    # 用了几个worker
    dataset = system_config.dataset
    workers = args.workers
    print("Process {}: using {} workers".format(rank, workers))

    training_dbs = [datasets[dataset](config["db"], split=train_split, sys_config=system_config) for _ in
                    range(workers)]   ## 涉及图片读取，标注文件读取
    validation_db = datasets[dataset](config["db"], split=val_split, sys_config=system_config) ## 打印1次“start prefetching data...”

    if rank == 0:
        print("system config...")
        pprint.pprint(system_config.full)  # 打印system config

        print("db config...")
        pprint.pprint(training_dbs[0].configs)  # 打印db config

        print("len of db: {}".format(len(training_dbs[0].db_inds)))  #训练集的图片数
        print("distributed: {}".format(args.distributed))   #是否分布式训练

    # 调用train训练函数
    train(training_dbs, validation_db, system_config, model, args)


if __name__ == "__main__":
    # 先parse参数
    args = parse_args()

    # 分布式的选项以及节点数量
    distributed = args.distributed
    world_size = args.world_size

    # 分布式节点必须大于0
    if distributed and world_size < 0:
        raise ValueError("world size must be greater than 0 in distributed training")

    # 每个节点上的gpu的数量
    ngpus_per_node = torch.cuda.device_count()

    # torch 下分布式运行
    if distributed:
        args.world_size = ngpus_per_node * args.world_size
        mp.spawn(main, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
    else:
        main(None, ngpus_per_node, args)