Caffe中CKPT流程写操作流程总结

总述

整体来看，Caffe中的CKPT操作包括两个大的部分：保存.caffemodel文件以及保存.solverstate文件。其中第一个文件中包括Net的所有层的所有Blob参数；第二个文件包括当前网络的一些状态（训练轮数等）以及动量信息。根据测试发现，两个文件大小相似，保存时间差别不大，第二个文件保存时间略微大于第一个。

image.png

1 Net数据写入

Caffe框架中CKPT的流程：

• 1 在solver.cpp中Step首先会判断是否添加了ckpt的标志位。

image.png

• 2 之后我们进入Solver<Dtype>::Snapshot()函数

image.png

其中包括三个重要的函数：（1）SnapshotToBinaryProto()；（2）SnapshotToHDF5()；（3）SnapshotSolverState(model_filename);

其中（1）与（2）用于获取model_filename。

其中（1）为：

image.png

model_file - string model_filename = SnapshotFilename(".caffemodel");

image.png

构建文件名：

image.png

并且在net_->ToProto(&net_param, param_.snapshot_diff());中进行net参数的保存工作。

这里首先进入的是net.cpp文件，该文件中存在ToProto函数，如下所示。

image.png

对网络中每一层均进行参数保存，转化为二进制文件。

每一个Layer中包括多个blob，所以我们需要进入layer.cpp文件。

image.png

之后进入blob.cpp文件：

image.png

其中const float* data_vec = cpu_data();表明，如果数据不在memory，会将其从GPU转出。

同理：

HDF5模式下同样

image.png

2 模型状态写入

之后进入SnapshotSolverState(model_filename);

image.png

这里分两个部分，第一为SnapshotSolverStateToBinaryProto(model_filename);第二部分为SnapshotSolverStateToHDF5(model_filename);。

第一部分当写入为二进制模式时：

image.png image.png

这里将state写入文件，而state包括如下内容：

image.png

此外，这里还会存入网络中的历史数据：

  for (int i = 0; i < history_.size(); ++i) {
    // Add history
    BlobProto* history_blob = state.add_history();
    history_[i]->ToProto(history_blob);
  }

而该数据存储过程中非常耗时，所以拖慢了整体的时间，同样使用ToProto的方式进行。

最终，.solverstate文件被写入。

image.png

最终各个部分的占比：

image.png