Transformers中加载预训练模型的过程剖析(二)

接着 Transformers中加载预训练模型的过程剖析(一) 来讲模型初始和载入预训练权重。

初始化模型

经过以上两个步骤，对于下载的 shibing624/text2vec-base-chinese-paraphrase 预训练模型(底层的模型是ernie)，我们从代码 model = AutoModel.from_pretrained(model_path) 成功加载了配置文件，并且最终返回了 model = ErnieModel.from_pretrained(model_path) ，换成以其他模型为基础的预训练模型，对应会得到 model = xxx.from_pretrained(model_path) 。

下面来看看用配置文件来初始化模型和构建模型的各个层和结构的过程。

类 ErnieModel 里是没有 from_pretrained 方法的，它只拿了配置文件和定义了模型的架构(将用于初始化模型)，如图15所示。顺着父类可以发现 from_pretrained 方法是在 transformers/modeling_utils.py 文件中(其他模型的话同样也会找到这里)，且其是一个类方法，cls表示调用它的类 ErnieModel ，如图16所示。

图15

图16

预训练权重文件

加载预训练权重是在 from_pretrained 方法的以下代码中，如图17所示。这里会按框架不同找对应的预训练权重文件，将权重保存在 archive_file 变量里，比如对于本篇文章使用的预训练模型的预训练权重文件是 pytorch_model.bin ，它在 transformers/utils/init.py 文件的 WEIGHTS_NAME 变量里面被定义，如图18所示。

图17

图18

模型初始化

模型的初始化在如下代码判断，如上所述，cls表示调用它的类 ErnieModel ，因此这里会得到 ErnieModel 中定义的模型架构。当然，到这里为止模型权重还未被预训练权重替换。

图19

载入权重

再往下，会调用 _load_pretrained_model 方法来载入预训练权重，如图20。 _load_pretrained_model 方法里会将初始化模型的权重与预训练权重做各种比对，这里不详述，只说两点：一个是参数初始化，一个是载入预训练权重。

图20

参数初始化

参数初始化代码如图21，这里 _initialize_weights 本质上会引用 ErnieModel 的父类 ErniePreTrainedModel 中的 _init_weights 方法定义的初始化，如图22所示。

图21

图22

载入预训练权重

继续往下找到载入预训练权重的代码判断，函数 _load_state_dict_into_model 本质上会用 torch/nn/modules/module.py 里面的 _load_from_state_dict 方法来把初始化模型的权重替换为预训练权重。

图23

创建实例

至此，所有准备工作已完毕，将返回一个 model 实例，使用此实例就可以调用 ErnieModel 里的 forward 方法来计算输入，然后得到模型的输出，继而可以处理一些下游任务了。