Ускорение генерации вложений USE (универсальный кодировщик предложений)

Question

Я работаю над проблемой подобия semanti c, используя универсальный кодировщик предложений . Набор данных содержит аннотации научных статей. Средняя длина составляет около 1500. В данных ~ 300 тыс. Записей, и создание встраивания USE для всех из них займет довольно много времени. Я ищу способы оптимизировать это. В настоящее время создание вложений для 10 тыс. Строк данных заняло ~ 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 и включил графический процессор. Использование графического процессора go не превышает 2–3%, а загрузка ЦП составляла ~ 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)

*

Затраченное время Время ЦП: пользователь 16 минут 48 секунд, системный: 2 минуты 59 секунд, всего: 19 минут 47 секунд Время на стене : 15мин 13с