[Detectron2]04-Trainer/Hooks

emiya

student

本篇文章中将详细的介绍 detectron2 中和模型训练相关的代码。

通过阅读这篇文章，你将可以了解到：

1、如何设计一个合理、灵活的深度学习训练框架
2、如何理解detectron2的训练类Trainer
3、如何自定义训练类用于训练自己的数据

写这篇文章的目的，一方面是希望能够整体的介绍和学习一下detectron2中设计的非常良好的训练流程，另一方面也是希望读者可以通过这篇文章迅速定位到代码中的相关函数，避免在代码阅读中迷失方向。因为我自己在修改detectron2的代码过程中，经常会出现改着改着就不知道某个函数在代码的什么位置了，所以在本篇文章中，我也会有侧重的列出一些大家可能会需要进行修改的代码在文件目录中的具体位置，方便大家进行快速的定位和修改。

本文分为以下三个部分：

1、如何设计一个自己的Trainer 类：介绍如何逐步设计出一个类似于detectron2 的Trainer类
2、detectron2 中的Trainer 相关代码解析：展示detectron2 中的Trainer 类的相关细节
3、如何基于detectron2提供的Trainer，自定义自己的数据加载过程、日志写入过程，等等

对于只希望了解detectron2相关细节的读者，欢迎直接阅读本文的第二部分内容

对于希望基于detectron2，加载自己的数据集，或者想进行更多自定义操作，需要魔改Trainer 的同学，可以阅读第三部分

一、如何设计一个Trainer

1.1 深度学习的整体训练流程

依旧遵从本专栏惯有的叙事风格，在开始介绍detectron2的代码之前，我们还是首先来思考一个问题， ”如果是你自己手动设计一个深度学习的训练流程，你会设计成什么样？“

基本上而言，你大概会写出如下的代码：

class Trainer():
    def __init__():
        # 定义基本参数
        self.max_iter = 10000
        self.save_iter = 2000
        self.log_iter = 1000
        # 首先是定义深度学习训练四件套
        self.model = create_model()
        self.optimizer = create_optimizer()
        self.data_loader = create_dataloader()
        self.learning_rate_adjuster = create_lr_adjuster()
        # 为了保存模型，你可能还会定义 saver， 用于模型的存储
        self.saver = create_saver()
        # 为了记录训练模型过程中的相关信息，可能你还需要定义一个 tensorboard writter
        self.writer = create_tensorboard_writer()
    def train(self):
        iteration = 0
        for self.iter in range(0, self.max_iter):
            # 首先来完成训练三部曲
            # step1 数据加载
            data = next(self.data_loader)
            # step2 loss 计算
            loss , acc , other_info = self.model(data)
            # step3 反向传播
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            # 然后来完成一些训练后的清理工作 
            iteration += 1
            if iteration % self.save_iter == 0:
                self.saver.save(model)
            if iteration % self.log_iter == 0:
                self.writer.log('loss',loss)
                self.writer.log('acc',acc)
                self.writer.log('other_info',other_info)

以上就是一个最基本的深度学习训练流程了，你可以在各种各样的项目中看到上面这样一套代码，他们可能不叫这个名字，但是一定是做类似的事情。而今天要介绍的，就是如何如何一步一步把它拆成更细的粒度，并且进行合理的封装，从而转化为detectron2中的训练流程的

在展开之前，多说一句就是，事实上，detectron2之所以需要把代码设计的非常的抽象，更多的原因是他希望把自己弄成一个”保罗万象的东西“,比方说在每次训练后，你不只想简单的用writer 记录记录loss ，还同时希望将很多的detectron2 训练过程中间的图片信息也保存下来，然后也写入到tensorboard当中，如果代码设计的不合理，要改动的地方就会有很多，这是不利于做多种实验的（当你想在一套代码里实验各种模型时，你会尽量希望每个训练过程对代码进行的修改尽可能的分开），所以就需要有一套设计的更好的”Trainer“

1.2 流程拆分Step1：抽象训练流程

要拆分上面的代码，咱们当然先拆大头 Trainer.train 方法了。在很多项目当中，train 方法都会写的非常的冗长，因为会把所有和训练，和记录相关的代码都放在train当中。那么在这里，我们可以把Trainer 进行如下的抽象：

class TrainerBase():
    def __init__(self):
    def train(self):
        iteration = 0
        self.before_train()
        for self.iter in range(0, max_iter):
            # 首先来完成训练三部曲
            # step1 数据加载
            self.before_step()
            self.run_step()
            self.after_step()
        self.after_train()
class Trainer(TrainerBase):
    def __init__(self):
        # 定义基本参数
        self.Epoch = 100
        self.save_iter = 2000
        self.log_iter = 1000
        # 首先是定义深度学习训练四件套
        self.model = create_model()
        self.optimizer = create_optimizer()
        self.data_loader = create_dataloader()
        self.learning_rate_adjuster = create_lr_adjuster()
        # 为了保存模型，你可能还会定义 saver， 用于模型的存储
        self.saver = create_saver()
        # 为了记录训练模型过程中的相关信息，可能你还需要定义一个 tensorboard writter
        self.writer = create_tensorboard_writer()
    def run_step(self):
        data = next(self.data_loader)
        # step2 loss 计算
        loss , acc , other_info = self.model(data)   
        # step3 反向传播
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
    def after_step(self):
        # 然后来完成一些训练后的清理工作 
        iteration += 1
        if iteration % self.save_iter == 0:
            self.saver.save(model)
        if iteration % self.log_iter == 0:
            self.writer.log('loss',loss)
            self.writer.log('acc',acc)
            self.writer.log('other_info',other_info)

在上述的代码中，我们主要做了两点变动：

1、将 TrainerBase 抽象出来，并在TrainerBase中定义了一个通用的训练流程
2、在TrainerBase 的子类Trainer中，定义了具体的 run_step 方法和 after_step 方法

那么，在这样的定义下，当用户希望定义自己的Trainer的时候，他就只需要继承TrainerBase ，并且实现自己的 beforre_train 、after_train 、 before step 、 after_step、run_step 方法即可。而在此基础之上，detectron2进行了更加进一步的抽象。

1.3 流程拆分Step2：使用Hook执行任务，进一步抽象代码

在上述拆分的基础之上，让我们进一步的来关注一下 after_train 中的这段代码：

class Trainer(TrainerBase):
    def __init__(self):
    def after_step(self):
        # 然后来完成一些训练后的清理工作 
        iteration += 1
        if iteration % self.save_iter == 0:
            self.saver.save(model)
        if iteration % self.log_iter == 0:
            self.writer.log('loss',loss)
            self.writer.log('acc',acc)
            self.writer.log('other_info',other_info)

在after_step 中，你的很多的”组件“都会在这里完成他们的任务。他们会每隔一定的步长执行一下自己的任务。而在detectron2 中，为了进一步的对代码进行一些抽象，把每一个”任务“都定义为一个个独立的”Hook“，每一个hook 都会实现自己的 beforre_train 、after_train 、 before_ step 、 after_step、run_step 方法：

class HookBase:
    def before_train(self):
        Called before the first iteration.
    def after_train(self):
        Called after the last iteration.
    def before_step(self):
        Called before each iteration.
    def after_step(self):
        Called after each iteration.
        pass

随后，基于Hook，就可以对训练代码进行进一步的抽象：

class Saver(HookBase):
    def __init__(self, save_iter):
        self.save_iter = save_iter
    def after_step(self):
        if self.trainer.iter % self.save_iter == 0:
            save_model(self.trainer.model)
class Writer(HookBase):
    def __init__(self,write_iter):
        self._debug_info = {}
        self.write_iter = write_iter
        self.writer = TensorboardWriter(...)
    def before_step(self):
        self._debug_info = {}
    def after_step(self):
        loss = self._debug_info['loss']
        self.writer.write(loss)
class Trainer(TrainerBase):
    def __init__(self):
        self.hooks : List[HookBase] = self.register_hooks()
    def register_hooks(self):
        self.hooks = []
        self.hooks.append(Saver(save_iter))
        self.hooks.append(Writer(write_iter))
        for h in hooks:
            assert isinstance(h, HookBase)
            h.trainer = weakref.proxy(self)
    def before_step(self):
        for hook in self.hooks:
             hook.before_step()
    def run_step(self):
        self.iter += 1
        data = next(self.data_loader)
        # step2 loss 计算
        loss , acc , other_info = self.model(data)   
        # step3 反向传播
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
    def after_step(self):
        for hook in self.hooks:
             hook.after_step()

现在再让我们梳理一下代码中关于hooks的逻辑：

1、Trainer在初始化时注册一系列的hooks ，每个hook 可以完成一个工作
2、注册hooks 的时候，通过 h.trainer = weakref.proxy(self) 把自身变为 hooks的属性，使得hook中可以通过 h.trainer.iter 获取trainer内部记录的一些训练状态相关的信息
3、每个hook都会有自己的一些列参数，这样，如 save_iter , write_iter 这样的信息就不是直接注册在trainer 中，而是记录在每个hook类自己的内容，保证了Trainer代码具有高度的扩展性

这里值得一提的是，上述伪代码中Writer 为了在每一个step中的行为保持独立，在每一轮开始前的before_step 中，都会将内部记录数据的变量self._debug_info清空，但是在after_step 中又可以从self._debug_info 中获取信息，而在上述代码中没有详细介绍信息是如何获取的，而这一部分信息会在后续介绍EventStorage详细介绍，再这里先不做展开。

1.4 流程拆分Step3：细化每个类的功能

在上一步中，我们已经将所有非训练相关的任务通过Hook的方式从Trainer 的代码中分离出去了，那么此时的Trainer类大概是如下形式：

class Trainer(TrainerBase):
    def __init__(self):
        self.hooks : List[HookBase] = self.register_hooks()
    def register_hooks(self):
        self.hooks = []
        self.hooks.append(Saver(save_iter))
        self.hooks.append(Writer(write_iter))
        for h in hooks:
            assert isinstance(h, HookBase)
            h.trainer = weakref.proxy(self)
    def before_step(self):
        for hook in self.hooks:
             hook.before_step()
    def run_step(self):
        self.iter += 1
        data = next(self.data_loader)
        # step2 loss 计算
        loss , acc , other_info = self.model(data)   
        # step3 反向传播
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
    def after_step(self):
        for hook in self.hooks:
             hook.after_step()

但是在上述写法中，注意到 run_step 方法是会随着不同的数据加载、不同的模型定义、损失定义而改变的，因此，有必要再进行一次抽象，让类的功能更加的泛化，于是，我们就终于能够看到detectron2设计的Trainer了：

二、detectron2 的Trainer设计方法

2.1 TrainerBase

基于第一节的分析，当我们 将所有非训练的任务通过Hook 分离 , 当我们 将我们针对不同任务设计的run_step 脚本也分离出去 ，我们就能够获得一个最基本的训练的基类，这个基类定义在 detectron2-engine-defaults.py 文件中,同样的HookBase 也定义在这个文件当中：

在下面这份代码中，只有EventStorage 我还没有介绍，在下一讲中，我将详细的介绍EventStorage 代码中的细节

class HookBase:
    def before_train(self):
    def after_train(self):
    def before_step(self):
    def after_step(self):
class TrainerBase:
    def __init__(self):
        self._hooks = []
    def register_hooks(self, hooks):
        hooks = [h for h in hooks if h is not None]
        for h in hooks:
            assert isinstance(h, HookBase)
            h.trainer = weakref.proxy(self)
        self._hooks.extend(hooks)
    def train(self, start_iter: int, max_iter: int):
        self.iter = self.start_iter = start_iter
        self.max_iter = max_iter
        with EventStorage(start_iter) as self.storage:
            try:
                self.before_train()
                for self.iter in range(start_iter, max_iter):
                    self.before_step()
                    self.run_step()
                    self.after_step()
            finally:
                self.after_train()
    def before_train(self):
        for h in self._hooks:
            h.before_train()
    def after_train(self):
        for h in self._hooks:
            h.after_train()
    def before_step(self):
        for h in self._hooks:
            h.before_step()
    def after_step(self):
        for h in self._hooks:
            h.after_step()
        # this guarantees, that in each hook's after_step, storage.iter == trainer.iter
        self.storage.step()
    def run_step(self):
        raise NotImplementedError

2.2 SimpleTrainer

注意到在上面那个类中，虽然为我们定义了训练流程，当我们回忆起训练类的三大件：模型、数据加载，损失函数和优化器定义时，我们仍然没有在上述文件中看到相关的定义。同时，我们也还没有实现 run_step 方法。而在detectron2的设计中，在这里又进行了两次的拆分，首先，detectron2 在 detectron2-engine-defaults.py 文件中，为我们设计了一套已知三大件后跑run_step 一个标准模板： SimpleTrainer

class SimpleTrainer(TrainerBase):
    def __init__(self, model, data_loader, optimizer):
        super().__init__()
        # 注意到为了灵活性，这里仍然没有定义 data_loader , model 和 optimizer 
        # 仍然是采用了 加载的方式，而真正定义这些的类，会在下一节中介绍
        model.train()
        self.model = model
        self.data_loader = data_loader
        self._data_loader_iter = iter(data_loader)
        self.optimizer = optimizer
    def run_step(self):
        # 通过next 方法获取数据
        data = next(self._data_loader_iter)
        data_time = time.perf_counter() - start
        # 执行前向代码
        loss_dict = self.model(data)
        losses = sum(loss for loss in loss_dict.values())
        self._detect_anomaly(losses, loss_dict)
        metrics_dict = loss_dict
        metrics_dict["data_time"] = data_time
        self._write_metrics(metrics_dict)
        # 进行反向传播
        self.optimizer.zero_grad()
        losses.backward()
        self.optimizer.step()
    def _detect_anomaly(self, losses, loss_dict):
        if not torch.isfinite(losses).all():
            raise FloatingPointError(
                "Loss became infinite or NaN at iteration={}!\nloss_dict = {}".format(
                    self.iter, loss_dict
            )

其次，在对应到具体项目，具体任务的时候，再定义一个继承SimpleTrainer的类 DefaultTrainer，来实现模型的创建、数据的加载等基本方法

2.3 DefaultTrainer

具体到目标检测任务时，真正的Trainer其实是定义在同一目录下的 DefaultTrainer

class DefaultTrainer(SimpleTrainer):
    def __init__(self, cfg):
        model = self.build_model(cfg)
        optimizer = self.build_optimizer(cfg, model)
        data_loader = self.build_train_loader(cfg)
        # For training, wrap with DDP. But don't need this for inference.
        if comm.get_world_size() > 1:
            model = DistributedDataParallel(
                model, device_ids=[comm.get_local_rank()], broadcast_buffers=False
        super().__init__(model, data_loader, optimizer)
        self.start_iter = 0
        self.max_iter = cfg.SOLVER.MAX_ITER
        self.cfg = cfg
        self.register_hooks(self.build_hooks())
    def build_hooks(self): 
        cfg = self.cfg.clone()
        cfg.defrost()
        cfg.DATALOADER.NUM_WORKERS = 0  # save some memory and time for PreciseBN
        ret = [
            hooks.IterationTimer(),
            hooks.LRScheduler(self.optimizer, self.scheduler),
        if comm.is_main_process():
            ret.append(hooks.PeriodicCheckpointer(self.checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD))
        if comm.is_main_process():
            # run writers in the end, so that evaluation metrics are written
            ret.append(hooks.PeriodicWriter(self.build_writers()))
        return ret
    def build_writers(self):
        return [
            # It may not always print what you want to see, since it prints "common" metrics only.
            CommonMetricPrinter(self.max_iter),
            JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
            TensorboardXWriter(self.cfg.OUTPUT_DIR),
    def train(self):
        super().train(self.start_iter, self.max_iter)
        if hasattr(self, "_last_eval_results") and comm.is_main_process():
            verify_results(self.cfg, self._last_eval_results)
            return self._last_eval_results

在这份代码中,最核心的部分就是build_model, build_dataloader 两个内容，而build_model 会在之后介绍模型的部分详细展开，而至于dataloader，在前面介绍dataset 的时候已经详尽的介绍过，所以这里不做展开，至于到这里注册的各种Hooks 以及其能够实现的功能，我们也会在接下来的章节中进行详尽的介绍，这里篇幅不宜过长，就不做展开，喜欢的话，不如点个收藏？

三、如何自定义自己的Trainer

首先需要说的是，为什么你需要自己的一个trainer呢？基本上而言，我自己实践下来可能是基于如下的几个需求：

1、需要定义自己的data_loader，以加载自己的数据集
2、针对自己的训练任务，有自己的评价方法，而希望边训边测试，所以需要定义自己的EvalHook，并在Trainer中调用，这时候就需要自定义自己的 register_hook 方法
3、在数据加载时，需要对数据进行debug，因此需要自定义自己的run_step 方法
4、因为进行简单调试的时候，经常容易训练出nan，因此需要定义自己的_detect_anomaly 方法

如果你也有上述的需求，就可以继承DefaultTrainer，并重写相关的方法，进行自定义即可。上述的后三点会和几个特定的Hook强相关，因此在这里先不做展开，等到时候hook的坑填完了再来补充，而关于第一点，其实也很简单，只需要在自己的训练类中实现自己的 build_data_loader 方法就好:

class TextTrainer(DefaultTrainer):
    def build_train_loader(self, cfg):
        # 重写这个方法就好
        text_mapper = BoundaryMapper(cfg)
        data_loader = build_detection_train_loader(cfg, text_mapper)
        return data_loader