如何保存Simple Transformer的pool层权重并加载到自定义自编码器中

步骤1：提取Simple Transformer的pool层权重并保存为pickle格式 #

首先确认你微调后的Simple Transformer模型实例的pool层路径，默认情况下Simple Transformer封装的预训练模型pool层在.model.pooler路径下，你可以先执行print(list(trained_st_model.model.named_modules()))确认层的命名，避免路径错误。
保存权重的代码如下：

import pickle
import torch

# 此处trained_st_model为你微调完成的Simple Transformer模型实例
pool_state_dict = trained_st_model.model.pooler.state_dict()

# 以pickle格式保存权重
with open("st_pool_weights.pkl", "wb") as f:
    pickle.dump(pool_state_dict, f)

注：PyTorch内置的torch.save()方法默认使用pickle协议序列化，和上述写法效果完全一致，如果你有跨设备加载需求，用torch.save(pool_state_dict, "st_pool_weights.pkl")后续加载更方便。

步骤2：将权重加载到自定义自编码器的pool层 #

加载前必须保证自编码器的pool层结构和Simple Transformer的pool层结构完全一致，包括层的类型、维度、是否带偏置、激活函数配置等，否则会触发维度不匹配报错。
加载代码如下：

import pickle
import torch

# 实例化你的自定义自编码器
your_autoencoder = CustomAutoEncoder()

# 加载保存的pool层权重
with open("st_pool_weights.pkl", "rb") as f:
    loaded_pool_weights = pickle.load(f)

# 将权重加载到自编码器的对应pool层，此处your_autoencoder.pool_layer替换为你自编码器中pool层的实际属性名
load_result = your_autoencoder.pool_layer.load_state_dict(loaded_pool_weights, strict=True)
print("权重加载不匹配的键：", load_result)

如果打印结果中缺失键和多余键都为空，说明权重加载完全成功。

可选操作 #

• 如果你不需要后续微调pool层，加载完成后可以固定权重：

  for param in your_autoencoder.pool_layer.parameters():
      param.requires_grad = False
  

• 验证权重正确性：对比保存前和加载后的权重数值即可，例如：

  # 保存前计算权重和
  print(trained_st_model.model.pooler.dense.weight.sum())
  # 加载后计算权重和
  print(your_autoencoder.pool_layer.dense.weight.sum())
  # 两者数值完全一致则说明加载无误
  

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