bert推理速度太慢的解决方案[用于生成表征向量]_python bert推理加速方法-优快云博客

本文链接：https://blog.youkuaiyun.com/junjie20082008/article/details/104280864

本文介绍如何优化BERT模型以快速生成表征向量。采用python生成器、tf.data.Dataset.from_generator，结合类的实例变量实现数据源源不断地输入，确保BERT模型能高效地为文本生成表征。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

思路借鉴：https://blog.youkuaiyun.com/hezhefly/article/details/98877796，再次表示感谢

本文的用途是基于bert的表征向量的生成

基于上面的文章的分类预测思路，修改成了基于bert 的表征向量生成，其中表征向量以最后一层CLS的向量表示，也可以取最后一层的所有token的平均，或者取倒数第二层的所有tokens的平均(倒数第二层不那么接近任务，但是又能学习到句子的较高层的语义）https://www.zhihu.com/question/334510491/answer/772077284，有别的需求可以按照这个进行修改。

解决方案

1.使用python的生成器，让程序“误以为”有很多序列需要预测，这里构造yield形式即可；
2.利用tf.data.Dataset.from_generator，加载生成器，声明好数据结构和类型；
3.利用class类的实例变量self的全局性，通过self.inputs把数据“喂给到”生成器内部，这样就保证了数据的“源源不断”；
4.程序需要close的机制，用于保证停掉生成器的工作。

"""Extract pre-computed feature vectors from BERT."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tokenization import FullTokenizer, validate_case_matches_checkpoint
from modeling import BertConfig
import tensorflow as tf
from tensorflow.python.estimator.estimator import Estimator
from tensorflow.python.estimator.run_config import RunConfig


import codecs
import collections
import json
import re

import modeling
import tokenization
import tensorflow as tf
import os
os.environ["CUDA_VISI