能否将PyTorch模型的架构与权重一同保存为JSON文件？

核心思路 #

PyTorch没有原生支持将模型架构（Python类定义）直接序列化为JSON，因此需要手动提取架构的可序列化描述，再与转换后的权重数据合并为JSON文件。用户端则通过解析JSON中的架构描述重建模型，再加载权重。这是当前满足你需求的唯一可行方案——必须自行编写架构的JSON描述，让用户据此构建模型。

具体实现步骤 #

1. 提取架构描述 #

把模型的关键参数（如输入维度、层数、每层类型及参数）整理成JSON可序列化的字典。对于复杂模型（如带分支的ResNet），需补充分支结构、自定义层的参数信息，确保用户能完全复现模型结构。

2. 转换权重为JSON格式 #

state_dict()返回的权重是张量，需将其转为列表并记录形状，确保后续能反向恢复为张量。

3. 合并并保存为JSON #

将架构描述和权重数据合并为一个大字典，用json.dump写入文件。

4. 用户端重建模型并加载权重 #

用户根据JSON中的架构描述手动构建模型，再将权重列表转回张量，通过load_state_dict加载。

代码示例 #

保存模型到JSON #

import torch
import json
from collections import OrderedDict

# 示例训练好的模型
class SimpleCNN(torch.nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(in_channels, 16, kernel_size=3, padding=1)
        self.relu = torch.nn.ReLU()
        self.pool = torch.nn.MaxPool2d(2)
        self.fc1 = torch.nn.Linear(16*16*16, num_classes)  # 假设输入为32x32

# 初始化并训练模型（此处省略训练逻辑）
model = SimpleCNN(3, 10)
# model.train() ... 训练步骤

# 提取架构描述
arch_desc = {
    "model_name": "SimpleCNN",
    "init_params": {"in_channels": 3, "num_classes": 10},
    "layers": [
        {"name": "conv1", "type": "Conv2d", "in_channels": 3, "out_channels":16, "kernel_size":3, "padding":1},
        {"name": "relu", "type": "ReLU"},
        {"name": "pool", "type": "MaxPool2d", "kernel_size":2},
        {"name": "fc1", "type": "Linear", "in_features": 4096, "out_features":10}
    ]
}

# 处理权重数据
weight_data = {}
for key, tensor in model.state_dict().items():
    weight_data[key] = {
        "shape": list(tensor.shape),
        "data": tensor.cpu().numpy().tolist()
    }

# 合并为JSON结构
model_json = {
    "architecture": arch_desc,
    "weights": weight_data
}

# 保存到JSON文件
with open("trained_model.json", "w") as f:
    json.dump(model_json, f, indent=4)

用户端加载模型 #

import torch
import json
from collections import OrderedDict

# 加载JSON文件
with open("trained_model.json", "r") as f:
    model_json = json.load(f)

arch_desc = model_json["architecture"]
weight_data = model_json["weights"]

# 根据架构描述构建模型
def build_model(arch_desc):
    model = torch.nn.Sequential()
    for idx, layer_info in enumerate(arch_desc["layers"]):
        layer_type = layer_info["type"]
        if layer_type == "Conv2d":
            layer = torch.nn.Conv2d(
                in_channels=layer_info["in_channels"],
                out_channels=layer_info["out_channels"],
                kernel_size=layer_info["kernel_size"],
                padding=layer_info["padding"]
            )
        elif layer_type == "ReLU":
            layer = torch.nn.ReLU()
        elif layer_type == "MaxPool2d":
            layer = torch.nn.MaxPool2d(kernel_size=layer_info["kernel_size"])
        elif layer_type == "Linear":
            layer = torch.nn.Linear(
                in_features=layer_info["in_features"],
                out_features=layer_info["out_features"]
            )
        model.add_module(f"{layer_info['name']}", layer)
    return model

# 重建模型并加载权重
model = build_model(arch_desc)
state_dict = OrderedDict()
for key, info in weight_data.items():
    tensor = torch.tensor(info["data"]).view(info["shape"])
    state_dict[key] = tensor

model.load_state_dict(state_dict)
model.eval()  # 切换到推理模式

注意事项 #

• 对于自定义层，需在架构描述中补充层的所有必要参数，用户端需实现对应的层类才能重建。
• 复杂模型（如带残差分支、注意力机制的模型）需要更细致的架构描述，例如记录分支的连接关系、层的顺序等。
• 权重转列表会增加文件体积，大模型可能会导致JSON文件过大，需权衡需求与存储成本。

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