如何使经过训练的 Torch 模型与 Transformeres 兼容？

Question

我已经使用

torch.save

训练并保存了 PyTorch 模型。现在我想将其加载为

GPT2LMHeadModel

。

from_pretrained

方法查找目录，然后查找 HuggingFace 注册表，但这些注册表都不存在。我只有一个序列化的 PyTorch 模型和一个

nn.Module

类。如何将它们与 Transformers 库集成？从头开始训练不是一个选择（需要太长时间）。

Answer 1

根本原因

您的本地路径需要在路径模型的注册表中注册。

解决方案

我建议您使用save_pretrained方法。它注册本地目录以供进一步使用 .from_pretrained 方法。

import torch
from transformers import GPT2LMHeadModel, GPT2Config

your_model = torch.load('path/to/your/model.pth')
state_dict = your_model.state_dict()
config = GPT2Config.from_pretrained('gpt2')
model = GPT2LMHeadModel(config)
model.load_state_dict(state_dict)

# Save local directory
model.save_pretrained('path/to/save/transformers_model')

# Now, you can load the Transformers model using from_pretrained
loaded_model = GPT2LMHeadModel.from_pretrained('path/to/save/transformers_model')

额外球：

GitHub 上的类似讨论：https://github.com/huggingface/transformers/issues/7849