如何设计一个合理、灵活的深度学习训练框架

一、如何设计一个Trainer

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

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

作者：emiya
链接：https://zhuanlan.zhihu.com/p/97326458
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

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)

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

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

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

作者：emiya
链接：https://zhuanlan.zhihu.com/p/97326458
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

class TrainerBase():
    def __init__(self):
        pass
    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 、 beforestep 、after_step、run_step 方法即可。

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 中，你的很多的”组件“都会在这里完成他们的任务。他们会每隔一定的步长执行一下自己的任务。

　　为了进一步的对代码进行一些抽象，把每一个”任务“都定义为一个个独立的”Hook“，每一个hook 都会实现自己的 beforre_train 、after_train 、 before_step 、after_step、run_step 方法：

class HookBase:

    def before_train(self):
        """
        Called before the first iteration.
        """
        pass

    def after_train(self):
        """
        Called after the last iteration.
        """
        pass

    def before_step(self):
        """
        Called before each iteration.
        """
        pass

    def after_step(self):
        """
        Called after each iteration.
        """
        pass

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

 

作者：emiya
链接：https://zhuanlan.zhihu.com/p/97326458
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

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类大概是如下形式：

作者：emiya
链接：https://zhuanlan.zhihu.com/p/97326458
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

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 方法是会随着不同的数据加载、不同的模型定义、损失定义而改变的，因此，有必要再进行一次抽象，让类的功能更加的泛化

二、Trainer设计方法

2.1 TrainerBase

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

作者：emiya
链接：https://zhuanlan.zhihu.com/p/97326458
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

class HookBase:

    def before_train(self):
        pass

    def after_train(self):
        pass

    def before_step(self):
        pass

    def after_step(self):
        pass


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

作者：emiya
链接：https://zhuanlan.zhihu.com/p/97326458
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

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={}!
loss_dict = {}".format(
                    self.iter, loss_dict
                )
            )

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

2.3 DefaultTrainer

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

作者：emiya
链接：https://zhuanlan.zhihu.com/p/97326458
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

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 两个内容，而它们会在之后介绍模型的部分详细展开。至于到这里注册的各种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

    @classmethod
    def build_test_loader(cls, cfg, dataset_name):
        text_mapper = BoundaryTestMapper(cfg)
        test_data_loader = build_detection_test_loader(cfg, dataset_name, mapper=text_mapper)
        return test_data_loader