unilm遇上对偶学习，模型参数共享的思考

本文记录一下unilm与对偶学习的一些碰撞火花。

【业务情景】：需要使用unilm来进行对偶学习。比如我们有A模型，B模型，需要同时训练这两个，他们的输入和输出是互补的。A的输出可以作为B的输入，B的输出可以作为A的输入。

【unilm说明】：仓库,论文，unilm的训练和推断阶段不太一致，包括模型结构和数据在模型里的处理。仓库提供的代码是分了两个阶段：训练+推断，即不同阶段使用不同的模型结构。

【我遇到的问题】：既然要使用对偶学习，那么我就需要这个模型同时能够进行训练和推断，训练一步之后，所取得的成效应该立即在推断中使用。即训练和推断应该共享参数，而前文说过模型结构略有不同，这该如何共享呢？

【我的思考和探索】：问题的突破点在“模型结构略有不同”，大部分结构甚至参数都是相同的。

想法①：把bert的encoder，embedding部分单独拿出来，使两个阶段共享这些参数。又突然想起来，由于训练和推断的不一样，数据在module里的处理都不一样（比如在self attention部分的数据处理就不同），所以这个不可行。

想法②：每训练一个批次的数据，就把“训练A”的state_dict导出，赋值给“推断A“（即名称相同的地方更新）。感觉是可行，但是这样效率有点过低。

想法③：加载模型的时候把每一层的QKV矩阵都拿出来，直接赋值给“训练A”和“推断A”，相当于是这两个的底层共用矩阵相同，反向传播的时候也共享了参数更新。这样看起来貌似比较合理。

【结果】：网上找到了一篇文章: 提到state_dict函数的返回值是浅拷贝，即返回的模型参数是一个“引用”，那么正好符合我的使用场景，即想法②可行，且只需要加载模型的时候赋值就可以了。

=更新=
貌似不太行，state_dict只是提供了一个引用地址，但是赋值的时候你去load_state_dict的时候不太得行。

【测试】：我拿出之前一项作业，先初始化模型，且用一个变量A保存最后一层的bias并打印，值为[ 0.0367, -0.0233], 训练几步，再用B保存最后一层的bias并打印，值为[ 0.0331, -0.0197],此时再打印A，发现A值已经变为和B一样了。

【结论】：model.state_dict()是浅拷贝，返回的参数仍然会随着网络的训练而变化。（出处）

最后还是采用了第三种方法。
具体代码：

def clone_params(model1 , model2):
    #embeddings
    model2.bert.embeddings.word_embeddings = model1.bert.embeddings.word_embeddings
    model2.bert.embeddings.position_embeddings = model1.bert.embeddings.position_embeddings
    model2.bert.embeddings.token_type_embeddings = model1.bert.embeddings.token_type_embeddings
    model2.bert.embeddings.LayerNorm = model1.bert.embeddings.LayerNorm

    #encoder.layer
    for i in range(len(model2.bert.encoder.layer)):
        #attention.self
        model2.bert.encoder.layer[i].attention.self.query = model1.bert.encoder.layer[i].attention.self.query
        model2.bert.encoder.layer[i].attention.self.key   = model1.bert.encoder.layer[i].attention.self.key
        model2.bert.encoder.layer[i].attention.self.value = model1.bert.encoder.layer[i].attention.self.value

        #attention.output
        model2.bert.encoder.layer[i].attention.output = model1.bert.encoder.layer[i].attention.output

        #layer.intermediate
        model2.bert.encoder.layer[i].intermediate = model1.bert.encoder.layer[i].intermediate

        #layer.output
        model2.bert.encoder.layer[i].output = model1.bert.encoder.layer[i].output
    #cls
    model2.cls.predictions.transform = model1.cls.predictions.transform

Ref:

1.https://www.cnblogs.com/LukeStepByStep/p/11248361.html

Contact me : jianshu(at)std.uestc.edu.cn