昇思25天学习打卡营第8天|保存与加载

1. 模型权重的保存与加载

保存模型权重

• 首先定义网络结构，例如使用nn.SequentialCell创建一个简单的全连接网络。
• 使用mindspore.save_checkpoint函数保存模型的权重。需要传入模型实例和保存的文件名。

def network():
    model = nn.SequentialCell(
        nn.Flatten(),
        nn.Dense(28*28, 512),
        nn.ReLU(),
        nn.Dense(512, 512),
        nn.ReLU(),
        nn.Dense(512, 10)
    )
    return model

# 保存模型权重
model = network()
mindspore.save_checkpoint(model, "model.ckpt")

加载模型权重

• 重新创建与保存时相同的模型实例。
• 使用mindspore.load_checkpoint函数加载保存的权重文件，返回一个参数字典。
• 使用mindspore.load_param_into_net函数将参数字典中的参数加载到模型中。
• 检查param_not_load，如果为空，则表示所有参数都已成功加载。

# 加载模型权重
model = network()
param_dict = mindspore.load_checkpoint("model.ckpt")
param_not_load, _ = mindspore.load_param_into_net(model, param_dict)
print(param_not_load)  # 应该为空，表示所有参数都已加载

2. MindIR的保存与加载

保存MindIR

• MindIR是MindSpore的中间表示，用于统一云侧训练和端侧推理。
• 使用mindspore.export函数将模型导出为MindIR格式。需要传入模型实例、输入Tensor、文件名和文件格式。

# 保存MindIR
model = network()
inputs = Tensor(np.ones([1, 1, 28, 28]).astype(np.float32))
mindspore.export(model, inputs, file_name="model", file_format="MINDIR")

加载MindIR

• 使用mindspore.load函数加载MindIR文件。
• 设置上下文模式为GRAPH_MODE，因为nn.GraphCell仅支持图模式。
• 将加载的图转换为nn.GraphCell实例，用于模型推理。

# 加载MindIR
mindspore.set_context(mode=mindspore.GRAPH_MODE)
graph = mindspore.load("model.mindir")
model = nn.GraphCell(graph)
outputs = model(inputs)
print(outputs.shape)  # 输出模型推理结果的形状

总结

MindSpore提供了灵活的模型保存与加载机制，支持模型权重的保存与加载，以及MindIR格式的导出与加载。通过这些功能，用户可以在训练过程中保存模型状态，并在需要时进行加载，便于模型的微调和部署。同时，MindIR格式的引入，使得模型在不同平台之间的转换更加方便。