根据层次结构将pytorch模型参数导出到不同的文件中

import torch
import pickle

# 定义一个示例模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = torch.nn.Linear(10, 20)
        self.fc2 = torch.nn.Linear(20, 30)
        self.fc3 = torch.nn.Linear(30, 40)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

# 创建模型实例
model = MyModel()

# 获取模型参数字典
params = model.state_dict()

# 遍历参数字典，将参数导出到不同的文件中
for layer_name, layer_params in params.items():
    # 构造文件名
    file_name = f"{layer_name}_params.pkl"
    
    # 导出参数到文件
    with open(file_name, 'wb') as f:
        pickle.dump(layer_params, f)

# 加载参数文件并设置为模型的state_dict
for layer_name, layer_params in params.items():
    # 构造文件名
    file_name = f"{layer_name}_params.pkl"
    
    # 加载参数文件
    with open(file_name, 'rb') as f:
        loaded_params = pickle.load(f)
    
    # 设置为模型的state_dict
    model.state_dict()[layer_name].copy_(loaded_params)