如何将PyTorch训练的自定义图像分类模型导出为HuggingFace格式

HuggingFace要求的模型格式必须包含模型权重文件、config.json（记录模型结构与核心参数），图像任务还需额外的preprocessor_config.json（记录预处理逻辑，比如归一化参数、输入图像尺寸）。以下是具体操作步骤：

1. 重新初始化模型实例 #

你训练时用的是Transformers库的图像分类模型（比如AutoModelForImageClassification或ViTForImageClassification），先按照训练时的参数重建模型：

import torch
from transformers import AutoModelForImageClassification, AutoImageProcessor

# 替换成你训练时用的预训练模型基座，比如"google/vit-base-patch16-224"
base_model_name = "你的预训练模型名称"
# 填写训练时的类别总数，以及类别与id的映射关系
num_labels = 3  # 示例：3分类任务
id2label = {0: "类别A", 1: "类别B", 2: "类别C"}
label2id = {"类别A": 0, "类别B": 1, "类别C": 2}

# 初始化模型结构
model = AutoModelForImageClassification.from_pretrained(
    base_model_name,
    num_labels=num_labels,
    id2label=id2label,
    label2id=label2id
)

# 加载你训练好的权重文件
model.load_state_dict(torch.load("best_model.pth"))

2. 保存模型与配置到本地 #

用Transformers内置的save_pretrained方法，它会自动生成符合要求的config.json和权重文件（pytorch_model.bin）：

# 指定保存目录，比如"./my_hf_compatible_model"
save_directory = "./my_hf_compatible_model"
model.save_pretrained(save_directory)

3. 保存图像预处理配置 #

图像任务必须同步保存预处理逻辑，确保部署时的输入与训练时一致：

# 初始化和训练时完全相同的处理器
processor = AutoImageProcessor.from_pretrained(base_model_name)
# 保存到同一个模型目录
processor.save_pretrained(save_directory)

4. 验证导出格式 #

此时你的save_directory目录下应该包含这些文件：

• pytorch_model.bin：模型权重文件
• config.json：模型结构与参数配置
• preprocessor_config.json：图像预处理配置
• .gitattributes：自动生成的版本控制配置

5. 推送到HuggingFace Hub（可选） #

如果要直接部署到Hub，先安装huggingface_hub，然后登录账号推送：

from huggingface_hub import login, push_to_hub

# 登录你的HuggingFace账号
login()
# 推送模型到Hub，替换为你的用户名和自定义模型名
model.push_to_hub("你的用户名/你的模型名称")
processor.push_to_hub("你的用户名/你的模型名称")

针对你的train.py补充说明 #

你当前的训练代码仅保存了模型的state_dict，必须先重建与训练时结构完全一致的模型，再加载state_dict才能导出为HuggingFace格式。务必保证num_labels、id2label等参数与训练时完全匹配，否则会出现加载错误。

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