PaddleSpeech TTS 设计要素 — 训练组件

(以下内容搬运自 PaddleSpeech)

主要讲述 PaddleSpeech TTS 的和训练相关的组件，以及我们为何如此设计它。如果你熟悉 chainer, 可以看出我们受到 chianer 的设计风格的影响。虽然这也不是 chainer 独此一家，我们也参考了 torch lightning 等专门帮忙解决训练问题的库，以及领域专用的库如 detectron2 等为了方便自己的模型开发而作出的设计。总体的设计原则是简单直观，可扩展性强，学习难度不高（这里需要斟酌，有些设计上手是有一点难度，但是理解了其设计，用起来将会很好用。）

全局播报员

训练和修改深度学习模型的时候，常常会有很多 logging 的需要，不少时候，logging 甚至成了模型 debug 和调整的关键，我们常用的也是各种可视化的工具，比如说 paddle 的 visualdl, tensorflow 的 tensorboard 以及还有 vidsom, wnb 等等。对这些工具的使用会以一定的方式加入到我们的代码中。除此之外，logging 以及更便捷的 print 也经常被使用于不同的目的。
在这些工具中 print 是最简单的，它没有 logging 的 logger, handler 的概念，也没有 tensorboard 的 summarywriter logdir 之类的概念，打印信息的时候也不需要 global_step 之类的选项，它足够的轻量，可以出现在代码中的任何位置，而且都打印到共同的 stdout. 当然它的可定制性也有限，比如打印少量的 scalar 还好，打印字典，或者打更复杂的对象的时候就不再直观。而且也稍纵即逝，保存下来又需要重定向之类的操作。
而且对于语音合成的模型开发，我们其实希望有一个更 universal 多媒体的 stdout，其实那就是类似 tensorboard 的工具。它允许很多媒体形态，但是使用的时候需要一个 summary writer, 写入信息的时候需要 step, 如果是图像或者语音等数据的时候，还有一些格式控制的参数。
这在一定程度上会破坏模块化的设计。比如说我的模型由多个 sublayers 构成，我想在某些子模块的 forward 方法内记录一些重要的信息。为此我可能需要把 summary writer 传给这个子模块，但是对于子模块来说，它的功能就是计算，它不应该有多余的考虑，在接口上我们也难以忍受比如说一个 nn.Linear 的初始化方法里有一个 optional 的 visualizer. 并且，对于一个计算模块来说，它哪里知道什么 global step 呢？这些都是和训练过程有关的东西。
于是，我们更常见的做法不是把写 log 的代码放在 layer 的定义中，而是将其返回出来，并且在训练过程中去获取它们，然后写到 summarywriter 中去。但是为此就会需要修改返回值，改多了的结果就是大家倾向于 dict everywhere. SummaryWriter 是训练层面的播报员，然后各个模块纷纷通过修改返回值的方式来向它传送信息。
我们觉得这样的方式稍微有点丑陋，我们更倾向只返回必要的信息，而不是为了迁就可视化和记录而改变返回值。我们的修改方法是，当你需要报道一些信息的时候，应该可以毫不费劲地报道。于是我们模仿了 chainer 的设计，使用了全局的播报员。其使用方式犹如高空轨道空投仓或者折跃。
具体的实现原理是利用了 python 的 module 级变量的全局性和 context manager 的效果。如果你们看到类似 tensorflow 的 with graph.as_default(): 也会发现其实是类似的。
paddlespeech/t2s/training/reporter.py 模块里面有一个模块级的变量 OBSERVATIONS，它在设计上是一个字典。用来存储 key-value.

@contextlib.contextmanager
def scope(observations):
    # make `observation` the target to report to.
    # it is basically a dictionary that stores temporary observations
    global OBSERVATIONS
    old = OBSERVATIONS
    OBSERVATIONS = observations

    try:
        yield
    finally:
        OBSERVATIONS = old

然后我们实现了一个上下文管理器 scope, 用来切换 OBSERVATIONS 这个名称绑定的变量（亦即写入目标）。然后定义了一个 getter 函数来获取 OBSERVATIONS 绑定的字典。

def get_observations():
	global OBSERVATIONS
	return OBSERVATIONS

然后定义一个函数，获取当前 OBSERVATIONS，并往里面写入 key-value pair.

def report(name, value):
    # a simple function to report named value
    # you can use it everywhere, it will get the default target and writ to it
    # you can think of it as std.out
    observations = get_observations()
    if observations is None:
        return
    else:
        observations[name] = value

这样子可以实现什么效果呢？不仅可以随时折跃一个 reporter, 还可以栈式切换。我们下面的测试代码显示了使用方式，先使用 first 作为当前的 OBSERVATION, 写入 first_begin=1, 然后开启第二个 OBSERVATION, 写入 second_begin=2, 再开启第三个 OBSERVATION, 写入 third_begin=3, 然后退出第三个 OBSERVATION 之后，自动就回到第二个, 再写入一些内容，退出第二个 OBSERVATION 之后就自动回到第一个 OBSERVATION.

def test_reporter_scope():
    first = {}
    second = {}
    third = {}

    with scope(first):
        report("first_begin", 1)
        with scope(second):
            report("second_begin", 2)
            with scope(third):
                report("third_begin", 3)
                report("third_end", 4)
            report("seconf_end", 5)
        report("first_end", 6)


    assert first == {'first_begin': 1, 'first_end': 6}
    assert second == {'second_begin': 2, 'seconf_end': 5}
    assert third == {'third_begin': 3, 'third_end': 4}
    print(first)
    print(second)
    print(third)

如此，我们平时写一些模块化的组件的时候，可以直接调用 report 即可。至于报道到了什么地方则可以由调用放来决定，只要调用方准备好 OBSERVATION, 然后开启一个 scope, 在 scope 内调用这个组件即可。
PaddleSpeech TTS 内的 Trainer 就是以这种方式进行信息报道的。

while True:
    self.observation = {}
    # set observation as the report target
    # you can use report freely in Updater.update()


    # updating parameters and state
    with scope(self.observation):
        update() # training for a step is defined here

节省硬盘的保存策略

训练模型的时候，我们往往会进行定期保存，这是比较简单的保存策略。一般存在按 iteration 保存和按照 epoch 保存两种方式。但是对于长期的训练，尤其是多人共用的机器上，硬盘容量可能捉襟见肘，我们可能就会选择不断覆写保存的文件，但是如果模型后期过拟合，或者质量变差，这么做可能会错失表现比较好的模型。
所以常见的策略是择优保存，具体的做法是保存和评估使用同一个周期，评估指标用于决定当前模型要不要保存。然后为了解决硬盘使用量，保存 K 份模型。然后这个 K 的大小就量力而行。
但是这么做就带来了另一个协议上的考虑，这些评估指标也需要保存下来，不然从一个 checkpoint 文件夹 resume training 的时候，就没法知道当前文件夹里的前辈们的表现如何，后来训出来的模型有没有超越它们。所以这就会需要 checkpoint 文件夹内保存一个类似 metadata 的文件，它的扩展性和可读性足够的好。而且接受这套协议的心智负担不大。
我们选择的文件格式是 jsonlines. 类似这样，里面记录了 iteration 数，指标和文件路径。每行是一个 json 记录。这个格式表现力和 json 相当，而且专门用于记录 records, 在这里表现为 list of dicts, 也便于流失读写。和用作配置文件的 dict 式的 json 是不太一样的，虽然一般的 json 也能记录列表，但是这么做就不便流失读写了。而且这个文件比 csv 表现力强一些，又不太复杂，所以我们在 parakeet 里面广泛使用 jsonl 作为元数据文件格式。在讲到文件的预定俗成的时候我们也还会再次提及。

{'iteration': 1000, 'valid_loss': 0.05, checkpooint='/checkpoints/step-1000.pdz'}
{'iteration': 2000, 'valid_loss': 0.02, checkpooint='/checkpoints/step-2000.pdz'}

然后解决了元文件的问题之后，我们回到原来的问题。我要如何实现这个保存策略呢？是手动写到每份实验代码里面，每次照抄一个模板？还是把这个功能抽象出来？如何抽象？
总所周知，keras 中保存这个事情是作为一个 Callback 来实现的。不过我们认为这对于保存行为本身做了太多抽象，要保存什么（参数？优化器？状态？还是其他），保存到哪里，用什么格式，什么文件名，周期等等，如何获取全局的 iteration 和 epoch 都需要有很严格的协议，往往这么做的都弄了一套 trainer 之类的，帮用户把整个训练流程代管了。我们觉得这样有点过分，所以我们还是决定抽象尽可能简单，不做过多的预设和约定。
保留 K 个最优这件事，本身就是动态维护一个列表，比最差的好就把它替换掉，比最差的差就进不来，大约如此。然后我们把入队和出队和两个 action 关联在一起，入队关联保存，出队关联删除，即可完成设计。于是我们实现了最简单的 KBest 策略。这个策略完全不知道什么是模型，什么是优化器，只需要两个 action 就可以。
可以参考 paddlespeech/t2s/training/checkpoint.py，核心的代码如下，如果还没有或者比最差的记录要好，那么就应该保存。如果确定了要替换，则删除最差的那个，然后保存新加入的。

def should_save(self, metric: float) -> bool:
    if not self.full():
        return True


    # already full
    worst_record_path = max(self.best_records, key=self.best_records.get)
    worst_metric = self.best_records[worst_record_path]
    return metric < worst_metric


def full(self):
    return (not self._save_all) and len(self.best_records) == self.max_size


def save_checkpoint_and_update(self, metric, path):
        # remove the worst
        if self.full():
            worst_record_path = max(self.best_records,
                                    key=self.best_records.get)
            self.best_records.pop(worst_record_path)
            self.del_fn(worst_record_path)


        # add the new one
        self.save_fn(path)
        self.best_records[path] = metric

不过目前还没有写好一个真正如何使用它，感觉还是有点难。还是需要对模型有一定的了解才能写好这个功能。

Updater：模型训练流程

出于保持功能的纯粹性和代码的可复用性，我们把模型代码做了抽象，都成为 nn.Layer 的子类，并且把核心的计算功能写在了里面。不同作者的写作风格稍微有差别，有的人喜欢把训练时前向过程写在 forward 里，有的人喜欢把推理时前向过程写在 forward 里。对于把训练时前向过程写在 forward 里的人来说，有的人喜欢直接写到 loss, 有的人喜欢写到预测结果，然后把计算 loss 的过程专门留给一个不含参的 Loss 类或者函数。
我们其实还是更倾向于把训练时前向过程写在 forward 里，但是只写到预测结果，不写到 loss, 因为为了这个模块还可能被更大的模块使用，能直接括号调用，比再调用一个方法看起来更方便一点。（尤其是大家常常默认一个 Layer 就只需要一个方法，那就是 forward 的话语下说这个事情，虽然说有的模型的训练和预测天差地别，不仅仅是关一下 dropout 那么简单）
但是只写前向计算，对于写模型来说是简单了，但是这只是转移了问题，最终我们还是要面对把一些需要之事组成一个实验这件事。比如说训练流程，定期评估，定期保存，可视化之类的。而且在这个过程中我们还会遇上一些只会在训练过程中陪伴你的存在，比如说 optimizer, learning rate scheduler, visualizer 之类的。不考虑这些东西是写不了训练过程的，但是这些东西又不是模型的一部分，不应该把它们写到模型里面吧，看起来确实不应该。
但是又总得把它们写在一个什么地方，写在最终的训练脚本里？可以，但是这就几乎意味着绝对的不可复用性，然而一个模型是如何被训练的，这件事其实还是可以有一定的可复用性的，不仅仅是每次跑那个脚本这种可复用性，而是换个数据集也可以用同样的训练流程跑起来的这个流程本身。
所以我们给这种中间过程也做了一个抽象。那就是 Updater, 它一般来说需要接受模型，优化器，数据流作为参数，功能就是训练。而且由于每个模型的训练方式可能存在差异，所以我们倾向于每个模型写一个对应的 Updater. 但是它不同于最终的训练脚本，还是有一定程度的封装，就是为了把定期的保存，可视化，评估等细节从这里抽厘掉，仅保留最基本的功能，亦即如何训练这个模型。
在训练过程中，可以随意调用我们前面所说的全局播报员，来记录训练过程中需要记录的动态。但是又不需要显式地把全局播报员糅合进 Updater 里面，保存了简洁，突出重点。我们常常看见这样的代码，关键的训练流程掩盖在大段大段的 summary_writer.add_scalar 和 summary_writer.add_iamge() 里面，以及一些繁琐的 global_step += 1 和 if global_step % K == 0: do_something 里面。这些繁琐的细节代码往往不是最关键的，但是代码篇幅缺占得最多。而且它们往往零零碎碎地访问很多临时变量，所以抽象成函数也不方便。所以我们设计了全局播报员这样的信息传递机制，以辅助将训练过程的核心（正向计算，反向计算，更新参数）凸显出来。
但是后来我们把保存方法也写进了 Updater 里面，因为 Updater 里面就有完整的需要保存的信息。所以就把这些功能写进了 Updater 里。
甚至我在思考可以去掉 Trainer 这层抽象也不是不可以。

Visualizer

因为我们选择了 observation 这种通信方式，所以 visualizer 其实就可以纯粹把 observation 里面的东西写到 visualizer 里面。
不过我现在的想法是，Updater 里面既然可以访问自己的 iteration，那么其实只要能弄出来一个 global 的 visualizer 也就可以。

InferenceExport

待实现，目前还需要更多的一些实践经验才能把这个写好。

P.S. 欢迎关注我们的 github repo PaddleSpeech, 是基于飞桨 PaddlePaddle 的语音方向的开源模型库，用于语音和音频中的各种关键任务的开发，包含大量基于深度学习前沿和有影响力的模型。