求助：使用Hugging Face from_pretrained加载自定义ViT模型失败

我想微调ViT模型，让它预测9个总和为1的连续尺寸比例值，定义了继承自ViTPretrainedModel的模型类，方便调用from_pretrained()：

class ViTForRegression(ViTPreTrainedModel):
    def __init__(self, model_name_or_path, num_labels=NB_SIEVING_SIZE):
        config = ViTConfig.from_pretrained(model_name_or_path)
        super().__init__(config)
        self.model = ViTModel.from_pretrained(model_name_or_path)
        self.regressor = torch.nn.Linear(self.model.config.hidden_size, num_labels)
        self.loss_fn = torch.nn.KLDivLoss(reduction='batchmean')

    def forward(self, pixel_values, labels=None):
        outputs = self.model(pixel_values=pixel_values)
        logits = self.regressor(outputs.last_hidden_state[:, 0])
        logits = torch.nn.functional.softmax(logits, dim=-1)
        log_probs = torch.log(logits) # 为KLDivLoss对softmax结果取log
        if labels is not None:
            loss = self.loss_fn(log_probs, labels)
        return { 'logits': logits, 'loss': loss }

用预训练模型初始化后通过Trainer训练：

model = ViTForRegression.from_pretrained('google/vit-base-patch16-224-in21k')

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=10)],
)

train_results = trainer.train()

训练后用save_pretrained()保存：

model.save_pretrained(os.path.join(SAVED_MODELS_PATH, 'regression_vit')

但在另一个脚本中，定义相同模型类后加载时出错：

model = ViTForRegression.from_pretrained(os.path.join(SAVED_MODELS_PATH, 'regression_vit'))

错误信息：

HFValidationError                         Traceback (most recent call last)
File ~/anaconda3/envs/pa_orcademo_torch/lib/python3.10/site-packages/transformers/configuration_utils.py:629, in PretrainedConfig._get_config_dict(cls, pretrained_model_name_or_path, **kwargs)
    627 try:
    628     # 从本地文件夹、缓存加载，或从模型Hub下载并缓存
--> 629     resolved_config_file = cached_file(
    630         pretrained_model_name_or_path,
    631         configuration_file,
    632         cache_dir=cache_dir,
    633         force_download=force_download,
    634         proxies=proxies,
    635         resume_download=resume_download,
    636         local_files_only=local_files_only,
    637         use_auth_token=use_auth_token,
    638         user_agent=user_agent,
    639         revision=revision,
    640         subfolder=subfolder,
    641         _commit_hash=commit_hash,
    642     )
    643     commit_hash = extract_commit_hash(resolved_config_file, commit_hash)

File ~/anaconda3/envs/pa_orcademo_torch/lib/python3.10/site-packages/transformers/utils/hub.py:417, in cached_file(path_or_repo_id, filename, cache_dir, force_download, resume_download, proxies, use_auth_token, revision, local_files_only, subfolder, repo_type, user_agent, _raise_exceptions_for_missing_entries, _raise_exceptions_for_connection_errors, _commit_hash)
    415 try:
    416     # 从URL加载或使用已缓存文件
--> 417     resolved_file = hf_hub_download(
...
  "qkv_bias": true,
  "torch_dtype": "float32",
  "transformers_version": "4.29.2"
}
' is the correct path to a directory containing a config.json file

已尝试：

• 用model.save_pretrained()保存后加载自定义ViT模型

预期结果：

• 模型成功加载并可用于预测

实际结果：

• 模型加载失败

1. 修复保存代码的语法错误 #

你保存模型的代码缺少一个右括号，会导致保存过程中断，配置或权重文件不完整：

# 原错误代码
model.save_pretrained(os.path.join(SAVED_MODELS_PATH, 'regression_vit')
# 修正后
model.save_pretrained(os.path.join(SAVED_MODELS_PATH, 'regression_vit'))

2. 调整自定义模型类的__init__方法以符合Hugging Face规范 #

自定义模型继承ViTPretrainedModel时，__init__应优先接收config参数，而非直接传入模型路径，避免和from_pretrained()的加载逻辑冲突：

class ViTForRegression(ViTPreTrainedModel):
    def __init__(self, config, num_labels=NB_SIEVING_SIZE):
        super().__init__(config)
        self.model = ViTModel(config)  # 用config初始化，无需重复调用from_pretrained
        self.regressor = torch.nn.Linear(config.hidden_size, num_labels)
        self.loss_fn = torch.nn.KLDivLoss(reduction='batchmean')

    def forward(self, pixel_values, labels=None):
        outputs = self.model(pixel_values=pixel_values)
        logits = self.regressor(outputs.last_hidden_state[:, 0])
        logits = torch.nn.functional.softmax(logits, dim=-1)
        log_probs = torch.log(logits)
        loss = None
        if labels is not None:
            loss = self.loss_fn(log_probs, labels)
        return { 'logits': logits, 'loss': loss }

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        # 加载配置文件
        config = ViTConfig.from_pretrained(pretrained_model_name_or_path)
        # 实例化模型
        model = cls(config, **kwargs)
        # 加载预训练权重
        model.load_state_dict(torch.load(os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin')))
        return model

3. 加载前的验证步骤 #

• 检查保存路径下的regression_vit文件夹，确认存在config.json和pytorch_model.bin两个核心文件
• 打开config.json，确保内容无语法错误（比如末尾无多余逗号）
• 加载脚本中必须定义完全相同的ViTForRegression类，且导入所有依赖库（ViTPretrainedModel、ViTConfig等）

内容的提问来源于stack exchange，提问作者Ben-Jy