如何将微调后的PyTorch版Musicgen模型转为Hugging Face格式并上传Hub？

问题核心 #

你当前代码的错误在于：load_state_dict返回的是包含missing_keys和unexpected_keys的元组，而非权重字典本身；且Audiocraft原生模型的权重结构和Transformers库的MusicgenModel不完全匹配，直接加载会导致结构不兼容。

正确转换与上传步骤 #

1. 权重格式适配与本地保存 #

先通过Audiocraft加载微调权重，再转换为Transformers兼容的格式：

import torch
from transformers import MusicgenConfig, MusicgenModel
from audiocraft.models import musicgen
import os

# 创建本地保存目录
save_dir = "./musicgen-finetuned"
os.makedirs(save_dir, exist_ok=True)

# 加载Audiocraft版Musicgen medium模型（用于匹配权重结构）
audiocraft_model = musicgen.MusicGen.get_pretrained('facebook/musicgen-medium', device='cuda')
# 加载你的微调权重文件
audiocraft_model.lm.load_state_dict(torch.load('NEW_MODEL.pt'))
# 提取适配Transformers的权重（Audiocraft的lm模块对应Transformers的模型主体）
target_state_dict = audiocraft_model.lm.state_dict()

# 初始化Transformers版模型与配置
config = MusicgenConfig.from_pretrained('facebook/musicgen-medium')
hf_model = MusicgenModel(config)
# 加载权重时关闭严格匹配，忽略Audiocraft特有的无关键
hf_model.load_state_dict(target_state_dict, strict=False)

# 保存模型与配置到本地
hf_model.save_pretrained(save_dir)
config.save_pretrained(save_dir)

2. 验证本地加载 #

执行以下代码确认模型可正常读取：

loaded_model = MusicgenModel.from_pretrained(save_dir)
print("模型本地加载成功")

3. 上传至Hugging Face Hub #

先完成依赖安装与账号登录：

pip install huggingface_hub
huggingface-cli login

再通过代码上传：

from huggingface_hub import upload_folder

upload_folder(
    folder_path=save_dir,
    repo_id="你的用户名/自定义模型名",  # 示例：zeke/musicgen-finetuned-medium
    repo_type="model",
)

关键注意事项 #

• 版本兼容：确保transformers、audiocraft、torch的版本与训练时一致，避免权重解析错误。
• strict参数：必须设置strict=False，跳过Audiocraft模型中Transformers不需要的额外权重键。

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