自定义MixModel调用save_pretrained报错，如何保存模型配置与权重？

save_pretrained是HuggingFace PreTrainedModel类的专属方法，你的MixModel继承的是PyTorch原生nn.Module，因此没有该方法。以下提供两种可行的解决方案：

方案一：让模型继承HuggingFace的PreTrainedModel（推荐） #

通过让模型继承PreTrainedModel并自定义配置类，可直接复用HuggingFace的模型保存/加载生态：

1. 定义自定义配置类
   继承BertConfig，添加模型特有的参数（卷积层、分类头、dropout等）：

   from transformers import BertConfig, PreTrainedModel
   
   class MixModelConfig(BertConfig):
       def __init__(self, conv_in_channels=3072, conv_out_channels=256, conv_kernel_size=5,
                    dropout_rate=0.3, clf_hidden_dim=256, num_labels=6, **kwargs):
           super().__init__(**kwargs)
           # 存储自定义模型参数
           self.conv_in_channels = conv_in_channels
           self.conv_out_channels = conv_out_channels
           self.conv_kernel_size = conv_kernel_size
           self.dropout_rate = dropout_rate
           self.clf_hidden_dim = clf_hidden_dim
           self.num_labels = num_labels
   

2. 修改MixModel继承关系
   调整模型继承PreTrainedModel，用自定义配置初始化各层：

   class MixModel(PreTrainedModel):
       config_class = MixModelConfig
   
       def __init__(self, config):
           super().__init__(config)
           self.bert = BertModel(config)
           self.hidden_size = self.bert.config.hidden_size
           # 用配置参数初始化自定义层
           self.conv = nn.Conv1d(in_channels=config.conv_in_channels, 
                                 out_channels=config.conv_out_channels, 
                                 kernel_size=config.conv_kernel_size, stride=1)
           self.relu = nn.ReLU()
           self.pool = nn.MaxPool1d(kernel_size=64 - config.conv_kernel_size + 1)
           self.dropout = nn.Dropout(config.dropout_rate)
           self.flat = nn.Flatten()
           self.clf1 = nn.Linear(config.conv_out_channels, config.clf_hidden_dim)
           self.clf2 = nn.Linear(config.clf_hidden_dim, config.num_labels)
   
       def forward(self, inputs, mask, labels=None):
           # 注意：不要在forward中转换tensor类型，应在数据加载阶段处理
           x = self.bert(input_ids=inputs, attention_mask=mask, return_dict=True)
           # 修复原代码错误：取BERT最后4层隐藏状态拼接，转置适配Conv1d输入格式
           concat_hidden = torch.cat(x.hidden_states[-4:], dim=-1).transpose(1, 2)
           x = self.conv(concat_hidden)
           x = self.relu(x)
           x = self.pool(x)
           x = self.dropout(x)
           x = self.flat(x)
           x = self.clf1(x)
           x = self.clf2(x)
           return x
   

3. 保存与加载
   现在可直接使用save_pretrained完成全量保存：

   # 初始化模型
   config = MixModelConfig.from_pretrained('bert-base-uncased', output_hidden_states=True)
   model = MixModel(config)
   # 训练模型...
   # 保存模型、权重、配置到指定目录
   model.save_pretrained("./mix_model_save")
   # 加载模型
   loaded_model = MixModel.from_pretrained("./mix_model_save")
   

方案二：手动保存配置与权重（无需修改继承关系） #

如果不想改动模型继承结构，可手动收集配置并保存：

1. 保存配置
   将模型关键参数整理为字典，保存为JSON文件：

   import json
   
   config_dict = {
       "pre_trained": "bert-base-uncased",
       "conv_in_channels": 3072,
       "conv_out_channels": 256,
       "conv_kernel_size": 5,
       "dropout_rate": 0.3,
       "clf_hidden_dim": 256,
       "num_labels": 6,
       "output_hidden_states": True
   }
   
   with open("./mix_model_config.json", "w") as f:
       json.dump(config_dict, f, indent=4)
   

2. 保存模型权重
   使用PyTorch原生方法保存权重：

   torch.save(model.state_dict(), "./mix_model_weights.bin")
   

3. 加载模型
   先加载配置初始化模型，再加载权重：

   # 加载配置
   with open("./mix_model_config.json", "r") as f:
       config_dict = json.load(f)
   # 初始化模型
   model = MixModel(pre_trained=config_dict["pre_trained"])
   # 加载权重
   model.load_state_dict(torch.load("./mix_model_weights.bin"))
   model.eval()
   

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