每次推理都Loading checkpoint shards:

Question

语言模型推理代码为:

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]

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

Answer 1

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)
    # 进行推理操作