每次推理都Loading checkpoint shards:

语言模型推理代码为:

hf_generator = pipeline("text2text-generation", model="aaa"
output = hf_generator(prompt, max_length=len(prompt)+128, do_sample=True)

每次推理都会显示:

Loading checkpoint shards: 100%|██████████| 2/2 [00:09<00:00,  4.62s/it]
Loading checkpoint shards: 100%|██████████| 2/2 [00:09<00:00,  4.62s/it]
Loading checkpoint shards: 100%|██████████| 2/2 [00:09<00:00,  4.64s/it]
Loading checkpoint shards: 100%|██████████| 2/2 [00:09<00:00,  4.74s/it]
Loading checkpoint shards: 100%|██████████| 2/2 [00:08<00:00,  4.17s/it]
Loading checkpoint shards: 100%|██████████| 2/2 [00:09<00:00,  4.67s/it]
Loading checkpoint shards: 100%|██████████| 2/2 [00:09<00:00,  4.69s/it]

导致十分占用时间,如何解决?

  1. 缩减模型大小:考虑使用更小的模型或模型子集来减少加载时间。您可以尝试使用经过压缩或剪枝的模型,或使用具有较少参数的模型。

  2. 持久化加载:将模型加载到内存中,以避免在每次推理时重新加载模型。您可以将加载模型代码放置在循环外部,并在每次推理时重用已加载的模型。这样可以减少重复加载的时间消耗。

下面是示例代码,展示了如何在循环外加载模型并重复使用:

hf_generator = pipeline("text2text-generation", model="aaa")
# 在循环之前加载模型

for i in range(num_iterations):
    output = hf_generator(prompt, max_length=len(prompt)+128, do_sample=True)
    # 进行推理操作