Press "Enter" to skip to content

在Hugging Face Transformers中，更快的TensorFlow模型

Published July 18, 2023 by 四海吧

在Hugging Face Transformers中，更快的TensorFlow模型 四海 第1张

在过去的几个月里，Hugging Face团队一直在努力改进Transformers的TensorFlow模型，使其更加健壮和快速。最近的改进主要集中在两个方面：

  1. 计算性能：BERT、RoBERTa、ELECTRA和MPNet已经得到改进，以使计算时间更快。这种计算性能的提升在所有计算方面都是显著的：图/急切模式、TF Serving以及CPU/GPU/TPU设备。
  2. TensorFlow Serving：这些TensorFlow模型中的每个模型都可以使用TensorFlow Serving进行部署，以从中受益于这种计算性能的提升。

计算性能

为了展示计算性能的改进，我们进行了一项彻底的基准测试，在测试中我们将BERT的性能与TensorFlow Serving 4.2.0版本与Google的官方实现进行了比较。这项基准测试在使用序列长度为128的GPU V100上运行（时间以毫秒为单位）：

与Google的实现相比，v4.2.0中的BERT实现更快，速度提高了约10%。除此之外，它还比4.1.1版本中的实现快两倍。

TensorFlow Serving

前一节演示了全新的Bert模型在Transformers的最新版本中在计算性能方面取得了显著的提升。在本节中，我们将逐步展示如何使用TensorFlow Serving将Bert模型部署到生产环境中，以从计算性能的提升中受益。

什么是TensorFlow Serving？

TensorFlow Serving属于TensorFlow Extended（TFX）提供的一组工具，使得将模型部署到服务器变得比以往更加容易。TensorFlow Serving提供了两个API，一个可以通过HTTP请求调用，另一个可以通过gRPC在服务器上运行推断。

什么是SavedModel？

SavedModel包含一个独立的TensorFlow模型，包括其权重和架构。它不需要运行模型的原始源代码，这使得它非常适合与支持读取SavedModel的任何后端（如Java、Go、C++或JavaScript等）共享或部署。SavedModel的内部结构如下所示：

savedmodel
    /assets
        -> 这里是模型所需的资源（如果有的话）
    /variables
        -> 包含权重的模型检查点
   saved_model.pb -> 代表模型图的protobuf文件

如何安装TensorFlow Serving？

有三种安装和使用TensorFlow Serving的方法：

  • 通过Docker容器安装
  • 通过apt包安装
  • 使用pip安装

为了简化操作并与所有现有的操作系统兼容，本教程将使用Docker。

如何创建SavedModel？

SavedModel是TensorFlow Serving所期望的格式。从Transformers v4.2.0开始，创建SavedModel具有以下三个附加功能：

  1. 序列长度可以在每次运行时自由修改。
  2. 所有模型输入都可以用于推断。
  3. 当使用output_hidden_states=True或output_attentions=True返回hidden states或attention时，它们现在被组合为一个单独的输出。

下面是将TFBertForSequenceClassification保存为TensorFlow SavedModel的输入和输出表示：

给定的SavedModel SignatureDef包含以下输入：
  inputs['attention_mask'] tensor_info:
      dtype: DT_INT32
      shape: (-1, -1)
      name: serving_default_attention_mask:0
  inputs['input_ids'] tensor_info:
      dtype: DT_INT32
      shape: (-1, -1)
      name: serving_default_input_ids:0
  inputs['token_type_ids'] tensor_info:
      dtype: DT_INT32
      shape: (-1, -1)
      name: serving_default_token_type_ids:0
给定的SavedModel SignatureDef包含以下输出：
  outputs['attentions'] tensor_info:
      dtype: DT_FLOAT
      shape: (12, -1, 12, -1, -1)
      name: StatefulPartitionedCall:0
  outputs['logits'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1, 2)
      name: StatefulPartitionedCall:1
方法名称为：tensorflow/serving/predict

为了直接将inputs_embeds（标记嵌入）作为输入而不是input_ids（标记ID），我们需要对模型进行子类化以定义一个新的serving签名。以下代码片段展示了如何实现：

from transformers import TFBertForSequenceClassification
import tensorflow as tf

# 创建一个子类以定义一个新的serving签名
class MyOwnModel(TFBertForSequenceClassification):
    # 使用新的input_signature修饰serving方法
    # input_signature表示预期输入的名称、数据类型和形状
    @tf.function(input_signature=[{
        "inputs_embeds": tf.TensorSpec((None, None, 768), tf.float32, name="inputs_embeds"),
        "attention_mask": tf.TensorSpec((None, None), tf.int32, name="attention_mask"),
        "token_type_ids": tf.TensorSpec((None, None), tf.int32, name="token_type_ids"),
    }])
    def serving(self, inputs):
        # 调用模型处理输入
        output = self.call(inputs)

        # 返回格式化的输出
        return self.serving_output(output)

# 使用新的serving方法实例化模型
model = MyOwnModel.from_pretrained("bert-base-cased")
# 使用saved_model=True保存模型以获得SavedModel版本和h5权重
model.save_pretrained("my_model", saved_model=True)

通过tf.function装饰器的新input_signature参数覆盖了serving方法。请参阅官方文档了解更多关于input_signature参数的信息。serving方法用于定义在使用TensorFlow Serving部署SavedModel时的行为。现在，SavedModel看起来符合预期，看下新的inputs_embeds输入：

给定的SavedModel SignatureDef包含以下输入：
  inputs['attention_mask'] tensor_info：
      dtype: DT_INT32
      shape: (-1, -1)
      name: serving_default_attention_mask:0
  inputs['inputs_embeds'] tensor_info：
      dtype: DT_FLOAT
      shape: (-1, -1, 768)
      name: serving_default_inputs_embeds:0
  inputs['token_type_ids'] tensor_info：
      dtype: DT_INT32
      shape: (-1, -1)
      name: serving_default_token_type_ids:0
给定的SavedModel SignatureDef包含以下输出：
  outputs['attentions'] tensor_info：
      dtype: DT_FLOAT
      shape: (12, -1, 12, -1, -1)
      name: StatefulPartitionedCall:0
  outputs['logits'] tensor_info：
      dtype: DT_FLOAT
      shape: (-1, 2)
      name: StatefulPartitionedCall:1
方法名称为：tensorflow/serving/predict

如何部署和使用SavedModel？

让我们逐步了解如何部署和使用一个用于情感分类的BERT模型。

步骤1

创建一个SavedModel。使用Transformers库加载一个在IMDB数据集上训练的名为nateraw/bert-base-uncased-imdb的PyTorch模型，并将其转换为TensorFlow Keras模型：

from transformers import TFBertForSequenceClassification

model = TFBertForSequenceClassification.from_pretrained("nateraw/bert-base-uncased-imdb", from_pt=True)
# saved_model参数是一个标志，表示在与h5权重一起创建模型的同时创建SavedModel版本
model.save_pretrained("my_model", saved_model=True)

步骤2

创建一个包含SavedModel的Docker容器并运行它。首先，拉取用于CPU的TensorFlow Serving Docker镜像（对于GPU，请将serving替换为serving:latest-gpu）：

docker pull tensorflow/serving

接下来，将serving镜像作为守护进程命名为serving_base运行：

docker run -d --name serving_base tensorflow/serving

将新创建的SavedModel复制到serving_base容器的models文件夹中：

docker cp my_model/saved_model serving_base:/models/bert

通过更改MODEL_NAME来匹配模型的名称（这里是bert），提交用于服务模型的容器，名称（bert）对应于我们要给我们的SavedModel起的名称：

docker commit --change "ENV MODEL_NAME bert" serving_base my_bert_model

并杀死以守护进程运行的serving_base镜像，因为我们不再需要它：

docker kill serving_base

最后，以守护进程的形式运行映像，并将容器中的端口8501（REST API）和8500（gRPC API）映射到主机上，并将容器命名为bert。

docker run -d -p 8501:8501 -p 8500:8500 --name bert my_bert_model

第三步

通过REST API查询模型：

from transformers import BertTokenizerFast, BertConfig
import requests
import json
import numpy as np

sentence = "我喜欢transformers中的新TensorFlow更新。"

# 加载我们SavedModel对应的tokenizer
tokenizer = BertTokenizerFast.from_pretrained("nateraw/bert-base-uncased-imdb")

# 加载我们SavedModel的模型配置
config = BertConfig.from_pretrained("nateraw/bert-base-uncased-imdb")

# 对句子进行分词
batch = tokenizer(sentence)

# 将batch转换为适当的字典
batch = dict(batch)

# 将示例放入大小为1的列表中，对应于批处理大小
batch = [batch]

# REST API需要一个包含关键字instances的JSON，以声明要处理的示例
input_data = {"instances": batch}

# 查询REST API，路径对应于http://host:port/model_version/models_root_folder/model_name:method
r = requests.post("http://localhost:8501/v1/models/bert:predict", data=json.dumps(input_data))

# 解析JSON结果。结果包含在一个名为“predictions”的根键的列表中
# 由于只有一个示例，取列表的第一个元素
result = json.loads(r.text)["predictions"][0]

# 返回的结果是概率，可以是正数也可以是负数，因此我们取它们的绝对值
abs_scores = np.abs(result)

# 取argmax对应于最大概率的索引
label_id = np.argmax(abs_scores)

# 打印正确的LABEL及其索引
print(config.id2label[label_id])

这应该返回POSITIVE。也可以通过gRPC（google远程过程调用）API获得相同的结果：

from transformers import BertTokenizerFast, BertConfig
import numpy as np
import tensorflow as tf
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
import grpc

sentence = "我喜欢transformers中的新TensorFlow更新。"
tokenizer = BertTokenizerFast.from_pretrained("nateraw/bert-base-uncased-imdb")
config = BertConfig.from_pretrained("nateraw/bert-base-uncased-imdb")

# 对句子进行分词，但这次使用TensorFlow张量作为输出，已经批处理大小为1。例如：
# {
#    'input_ids': <tf.Tensor: shape=(1, 3), dtype=int32, numpy=array([[  101, 19082,   102]])>,
#    'token_type_ids': <tf.Tensor: shape=(1, 3), dtype=int32, numpy=array([[0, 0, 0]])>,
#    'attention_mask': <tf.Tensor: shape=(1, 3), dtype=int32, numpy=array([[1, 1, 1]])>
# }
batch = tokenizer(sentence, return_tensors="tf")

# 创建一个通道，将连接到容器的gRPC端口
channel = grpc.insecure_channel("localhost:8500")

# 创建一个用于预测的存根。该存根将用于将gRPC请求发送到TF服务器。
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)

# 创建一个用于预测的gRPC请求
request = predict_pb2.PredictRequest()

# 设置模型的名称，对于此用例，它是bert
request.model_spec.name = "bert"

# 设置用于格式化gRPC查询的签名，这里使用默认签名
request.model_spec.signature_name = "serving_default"

# 从tokenizer给定的input_ids中设置input_ids
# tf.make_tensor_proto将TensorFlow张量转换为Protobuf张量
request.inputs["input_ids"].CopyFrom(tf.make_tensor_proto(batch["input_ids"]))

# attention mask也一样
request.inputs["attention_mask"].CopyFrom(tf.make_tensor_proto(batch["attention_mask"]))

# token type ids也一样
request.inputs["token_type_ids"].CopyFrom(tf.make_tensor_proto(batch["token_type_ids"]))

# 发送gRPC请求到TF服务器
result = stub.Predict(request)

# 输出是一个Protobuf，其中唯一的输出是一个分配给键logits的概率列表
# 由于概率是浮点数，因此将列表转换为浮点数的numpy数组
output = result.outputs["logits"].float_val

# 打印正确的LABEL及其索引
print(config.id2label[np.argmax(np.abs(output))])

结论

由于在transformers中应用的最新更新，现在可以使用TensorFlow Serving轻松地将模型部署到生产环境中。我们正在考虑的下一步之一是直接将预处理部分集成到SavedModel中，以使事情更加简单。

Published in 四海

Leave a Reply

Your email address will not be published. Required fields are marked *