如何加载非Torch库的data2vec2预训练PyTorch模型？

1. 补全依赖项 #

先解决Data2VecMultiModel导入时的依赖缺失问题，执行以下命令安装核心依赖：

pip install omegaconf fairseq torchvision transformers

2. 加载视觉模型（base_imagenet.pt） #

直接用torch.load只能拿到权重字典，无法得到可用模型。需先实例化对应架构的data2vec视觉模型，再加载本地权重：

import torch
from fairseq.models.data2vec.data2vec_vision import Data2VecVisionModel

# 初始化base规格的data2vec视觉模型
model = Data2VecVisionModel.from_pretrained(
    "",
    checkpoint_file="base_imagenet.pt",
    arch="data2vec_vision_base",
    task="image_classification"
)
model.eval()

# 生成图像嵌入示例
from torchvision import transforms
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# 假设img是PIL格式图像
img_tensor = preprocess(img).unsqueeze(0)
with torch.no_grad():
    img_embeddings = model.extract_features(img_tensor)

3. 加载NLP模型（nlp_base.pt） #

文本模型同样需要先实例化对应架构，再加载本地权重：

from fairseq.models.data2vec.data2vec_text import Data2VecTextModel

# 初始化base规格的data2vec文本模型
model = Data2VecTextModel.from_pretrained(
    "",
    checkpoint_file="nlp_base.pt",
    arch="data2vec_text_base",
    task="language_modeling"
)
model.eval()

# 生成文本嵌入示例（需搭配对应模型的字典文件）
from fairseq.data import Dictionary
# 字典文件(dict.txt)需从data2vec2官方发布包中获取
dictionary = Dictionary.load("path/to/dict.txt")
tokenized_text = dictionary.encode_line("输入文本内容", append_eos=True).unsqueeze(0)
with torch.no_grad():
    text_embeddings = model.extract_features(tokenized_text)

关键注意点 #

• 确保实例化的模型架构（base）与下载的权重文件规格完全匹配
• 文本模型必须搭配对应版本的字典文件，否则无法正确 tokenize 输入
• 所有预测操作前需调用model.eval()，关闭训练相关的随机层（如dropout）

内容的提问来源于stack exchange，提问作者Katya Willard