加快 USE（通用句子编码器）嵌入的生成

Question

我正在使用universal sentence encoder 解决语义相似性问题。该数据集包含学术文章的摘要。平均长度约为 1500。数据中有大约 30 万条记录，为所有这些记录生成 USE 嵌入需要很长时间。我正在寻找优化这一点的方法。目前，为 10k 行数据生成嵌入需要大约 15 分钟。

from tqdm import tqdm    
use_module_url = "https://tfhub.dev/google/universal-sentence-encoder/4"

model = hub.load(use_module_url)
print ("module %s loaded" % use_module_url)

def embed(input):
  return model(input)

def get_features(texts):
        if type(texts) is str:
            texts = [texts]
        return embed(texts)    

def data_iterator(data):
  chunk_list = []
  for x in tqdm(range(0, len(data), 1000)):
    if x+1000 > len(data):
      chunk_list.append(data[x:len(data)])
    else:
      chunk_list.append(data[x:x+1000])
  return chunk_list

data = df['text'][:10000].values
data_processed = list(map(process_text, data))

在这里，我想加快为我的数据生成 USE 嵌入的速度。我正在尝试 kaggle 内核并打开了 GPU。 GPU 利用率不超过 2-3%，CPU 利用率约为 120%

%%time
BASE_VECTORS = []

chunk_list = data_iterator(data_processed)

for i in tqdm(chunk_list):
    BASE_VECTORS_tmp = get_features(i)
    BASE_VECTORS.extend(BASE_VECTORS_tmp)

BASE_VECTORS = np.asarray(BASE_VECTORS)

花费时间 CPU时间：user 16min 48s，sys：2min 59s，总计：19min 47s 挂墙时间：15分13秒

Answer 1

可能你没有安装 tensorflow 的 GPU 版本，或者有一些 CUDNN 版本不匹配。通常 USE 会大量使用 GPU。