c++模型如何加载权重文件的部分参数权重

img

img


C++模型中的参数名在权重文件里面都有，如何从权重文件中选出对应的参数权重加载？

要在C++中加载模型的部分参数权重，您可以使用一些流行的深度学习库，如TensorFlow或PyTorch，并使用它们提供的C++ API。以下是一个使用TensorFlow C++ API加载模型权重文件的示例：

#include <tensorflow/c/c_api.h>

int main() {
  // 创建一个新的TensorFlow会话
  TF_SessionOptions* session_options = TF_NewSessionOptions();
  TF_Session* session = TF_NewSession(session_options, TF_NewStatus());

  // 加载模型的权重文件
  const char* model_path = "path/to/model.pb";
  const char* checkpoint_path = "path/to/checkpoint.ckpt";
  
  TF_Graph* graph = TF_NewGraph();
  TF_Status* status = TF_NewStatus();

  // 从.pb文件中加载图
  TF_Buffer* graph_def = NULL;
  TF_Buffer* checkpoint_bytes = NULL;
  graph_def = TF_NewBufferFromFile(model_path, status);
  TF_ImportGraphDefOptions* options = TF_NewImportGraphDefOptions();
  TF_GraphImportGraphDef(graph, graph_def, options, status);

  // 从.ckpt文件中加载权重
  checkpoint_bytes = TF_NewBufferFromFile(checkpoint_path, status);
  TF_SessionRun(
      session,
      NULL,  // 输入节点
      NULL,  // 输入张量
      0,     // 输入数量
      NULL,  // 输出节点
      NULL,  // 输出张量
      0,     // 输出数量
      NULL,  // 目标操作节点
      checkpoint_bytes->data,  // 权重数据
      checkpoint_bytes->length,  // 权重数据长度
      NULL,  // 运行元数据
      status);

  // 检查是否加载成功
  if (TF_GetCode(status) != TF_OK) {
    fprintf(stderr, "加载权重文件时出错: %s\n", TF_Message(status));
    return 1;
  }

  // 可以使用模型进行预测或其他操作

  // 清理资源
  TF_DeleteBuffer(graph_def);
  TF_DeleteBuffer(checkpoint_bytes);
  TF_DeleteGraph(graph);
  TF_DeleteSession(session, status);
  TF_DeleteStatus(status);
  TF_DeleteImportGraphDefOptions(options);
  TF_DeleteSessionOptions(session_options);

  return 0;
}

请注意，这只是一个简单的示例，您需要根据自己的模型和需求进行适当的修改。此外，您还需要正确安装和配置TensorFlow C++ API，并链接所需的库文件。

可以使用PyTorch的state_dict()方法从模型中获取参数字典，然后使用torch.load()方法加载权重文件。在加载权重文件后，您可以使用state_dict().update()方法将权重与模型中的参数进行匹配。
可参考

import torch
from torch import nn

# 初始化模型和权重文件
model = nn.Linear(10, 5)
weights_files = './test.pt'
weights = torch.load(weights_files)

# 从模型中获取参数字典
model_dict = model.state_dict()

# 将权重与模型中的参数进行匹配
match_dict = {k: v for k, v in weights.items() if k in model_dict}
model_dict.update(match_dict)

# 重新构建模型并加载权重
model.load_state_dict(model_dict)