tensorflow2.x加载预训练模型

tensorflow的预训练模型网址为：https://hub.tensorflow.google.cn

已经下载好的预训练模型：bert_zh_L-12_H-768_A-12的v4版本（链接：https://pan.baidu.com/s/1Ld_mvXf8Es4lwCaW1PgKsQ 
提取码：3xba）和bert_zh_preprocess的v3版本（链接：https://pan.baidu.com/s/1jseOQKlVj88-Hpak2PXsCw 
提取码：s9eq）

模型离线加载：

import tensorflow_hub as hub
# A dependency of the preprocessing for BERT inputs
#模型预处理时内部使用到了对应的模块，所以必须导入
import tensorflow_text as text

preprocessor = hub.KerasLayer(
    "data/bert_preprocess_module/3")
encoder = hub.KerasLayer(
    "data/bert_zh_encoder_L-12_H-768_A-12/4",
    trainable=False)
text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
# 模型线上地址：https://hub.tensorflow.google.cn/tensorflow/bert_zh_preprocess/3
#无参数，主要用来对输入的字符串文本进行分词处理
encoder_inputs = preprocessor(text_input)
# https://hub.tensorflow.google.cn/tensorflow/bert_zh_L-12_H-768_A-12/4

outputs = encoder(encoder_inputs)
pooled_output = outputs["pooled_output"]  # [batch_size, 768].
sequence_output = outputs["sequence_output"]  # [batch_size, seq_length, 768].

embedding_model = tf.keras.Model(text_input, pooled_output)
while True:
    mysentence = input()
    sentences = tf.constant([mysentence])
    print(embedding_model(sentences))

基于预训练模型创建新模型：

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text as text


class Mybert(tf.keras.models.Model):
    def __init__(self):
        super(Mybert, self).__init__()
        self.preprocessor = hub.KerasLayer(
            "data/bert_preprocess_module/3")
        self.encoder = hub.KerasLayer(
            "data/bert_zh_encoder_L-12_H-768_A-12/4",
            trainable=False)

    def call(self, inputs):
        # ***************word token embedding begin***************
        inputs = tf.convert_to_tensor(inputs)
        encoder_inputs = self.preprocessor(inputs)
        # https://hub.tensorflow.google.cn/tensorflow/bert_zh_L-12_H-768_A-12/4

        outputs = self.encoder(encoder_inputs)

        return outputs["pooled_output"]

if __name__ == '__main__':
    mybert = Mybert()
    while True:
        mysentence = input()
        sentences = tf.constant([mysentence])
        print(mybert(sentences))