HuggingFace - 当我从检查点加载时 model.generate() 非常慢

Question

我正在尝试使用 Donut 模型（在 HuggingFace 库中提供）使用我的自定义数据集（格式类似于 RVL-CDIP）进行文档分类。当我训练模型并model.generate()在训练循环中运行模型推理（使用方法）进行模型评估时，这是正常的（每张图像的推理大约需要0.2s）。

但是，如果训练后，我使用该方法将模型保存到检查点save_pretrained，然后使用该from_pretrained方法加载检查点，则model.generate()运行速度极慢（6s ~ 7s）。

这是我用于推理的代码（训练循环中的推理代码完全相同）：

model = VisionEncoderDecoderModel.from_pretrained(CKPT_PATH, config=config)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

accs = []
model.eval()
for i, sample in tqdm(enumerate(val_ds), total=len(val_ds)):
    pixel_values = sample["pixel_values"]
    pixel_values = torch.unsqueeze(pixel_values, 0)
    pixel_values = pixel_values.to(device)

    start = time.time()
    task_prompt = "<s_fci>"
    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
    decoder_input_ids = decoder_input_ids.to(device)
    print(f"Tokenize time: {time.time() - start:.4f}s")

    start = time.time()
    outputs = model.generate(
        pixel_values,
        decoder_input_ids=decoder_input_ids,
        max_length=model.decoder.config.max_position_embeddings,
        early_stopping=True,
        pad_token_id=processor.tokenizer.pad_token_id,
        eos_token_id=processor.tokenizer.eos_token_id,
        use_cache=True,
        num_beams=1,
        bad_words_ids=[[processor.tokenizer.unk_token_id]],
        return_dict_in_generate=True,
    )
    print(f"Inference time: {time.time() - start:.4f}s")

    # turn into JSON
    start = time.time()
    seq = processor.batch_decode(outputs.sequences)[0]
    seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
    seq = re.sub(r"<.*?>", "", seq, count=1).strip()  # remove first task start token
    seq = processor.token2json(seq)
    if "class" not in seq.keys():
        seq["class"] = "other"
    print(f"Decoding time: {time.time() - start:.4f}s")

    gt = sample["labels"]
    score = float(seq["class"] == gt["class"])

    accs.append(score)

acc_score = np.mean(accs)

print(f"Accuracy: {acc_score * 100:.4f}%")

我在 NVIDIA A100 40GB GPU 上运行该模型。我使用的 Anaconda 环境具有以下要求：

cudatoolkit==11.7
torch==1.13.1+cu117
torchvision==0.14.1+cu117
datasets==2.10.1
transformers==4.26.1
sentencepiece==0.1.97
onnx==1.12.0
protobuf==3.20.0

如何使用 HuggingFace 库加速 Donut 模型推理？非常感谢。

除了使用检查点测量训练时的推理时间和推理时间之外，我没有尝试过任何其他方法。