Caffe源码，caffe::Solver<Dtype>::Snapshot运行分析

╰+攻爆jí腚メ 2022-02-16 04:49 285阅读 0赞

# 1. 概述 #

在进行网络训练的时候会保存网络权值信息到文件中，用来后序部署等。在Caffe中实现这个功能是使用Snapshot()函数实现的，在Caffe中权值文件的保存形式有2中，它们是BinaryProto格式文件和hdf5格式文件，一般未指明的情况下，缺省为BinaryProto格式的形式。本篇文章介绍的存储方式是BinaryProto，hdf5形式的存储方式也是具有类似性。

**存储的总体流程：**  
在网络训练次数达到规定的权值文件保存的次数时，就会调用`Solver::Snapshot()`函数保存权值文件，这个函数会调用到net的权值文件保存函数`ToProto()`，net的这个函数接下来就会调用Layer里面的函数`ToProto()`，最后会调用到Blob里面的`ToProto()`函数。正是这样层级的调用方式将网络的权值保存到文件中。

# 2. 流程分析 #

## 2.1 caffemodel文件的保存 ##

这个文件保存的是网络中学习到的权值文件，在需要保存权值文件的时候调用的接口是：

//solver.cpp
    template <typename Dtype>
    void Solver<Dtype>::Snapshot() { 
      CHECK(Caffe::root_solver());
      string model_filename;
      switch (param_.snapshot_format()) { 
      case caffe::SolverParameter_SnapshotFormat_BINARYPROTO: //BinaryProto文件
        model_filename = SnapshotToBinaryProto();
        break;
      case caffe::SolverParameter_SnapshotFormat_HDF5: //hdf5文件
        model_filename = SnapshotToHDF5();
        break;
      default:
        LOG(FATAL) << "Unsupported snapshot format.";
      }
    
      SnapshotSolverState(model_filename); //保存solver_state
    }

这个函数会调用`SnapshotToBinaryProto()`来使用net的`ToProto()`接口：

//net.cpp
    template <typename Dtype>
    string Solver<Dtype>::SnapshotToBinaryProto() { 
      string model_filename = SnapshotFilename(".caffemodel");  //得到保存文件的路径
      LOG(INFO) << "Snapshotting to binary proto file " << model_filename;
      NetParameter net_param;	//存放需要保存的所有信息
      net_->ToProto(&net_param, param_.snapshot_diff());	//调用net.cpp中的ToProto()函数
      WriteProtoToBinaryFile(net_param, model_filename);
      return model_filename;
    }

在net的内部`ToProto()`是逐层调用Layer的`ToProto()`接口的：

//layer.hpp
    // Serialize LayerParameter to protocol buffer
    template <typename Dtype>
    void Layer<Dtype>::ToProto(LayerParameter* param, bool write_diff) { 
      param->Clear();
      param->CopyFrom(layer_param_);
      param->clear_blobs();
      for (int i = 0; i < blobs_.size(); ++i) { 
        blobs_[i]->ToProto(param->add_blobs(), write_diff);
      }
    }

最后就调用到了Blob的`ToProto()`接口，实现具体权值数据的存储：

//blob.cpp
    template <>
    void Blob<float>::ToProto(BlobProto* proto, bool write_diff) const { 
      proto->clear_shape();
      for (int i = 0; i < shape_.size(); ++i) { 
        proto->mutable_shape()->add_dim(shape_[i]);
      }
      proto->clear_data();
      proto->clear_diff();
      const float* data_vec = cpu_data();
      for (int i = 0; i < count_; ++i) { 
        proto->add_data(data_vec[i]);
      }
      if (write_diff) { 
        const float* diff_vec = cpu_diff();
        for (int i = 0; i < count_; ++i) { 
          proto->add_diff(diff_vec[i]);
        }
      }
    }

## 2.2 solverstate文件的保存 ##

这个文件保存的的是在训练过程中的梯度信息，首先来看下梯度信息的更新过程吧，在某个具体的solver在进行`Step()`函数的时候，没进行一次解算都会保存一次网络中的梯度信息。

template <typename Dtype>
    void Solver<Dtype>::Step(int iters) { 
        ......
    
      while (iter_ < stop_iter) { 
        ......
        // accumulate the loss and gradient
        Dtype loss = 0;
        for (int i = 0; i < param_.iter_size(); ++i) { 
          loss += net_->ForwardBackward();
        }
        loss /= param_.iter_size();
        ......
        ApplyUpdate(); //保存梯度信息
        ......
      }
    }

在`ApplyUpdate()`函数中会调用`ComputeUpdateValue()`函数用来保存当前迭代次数的梯度信息

template <typename Dtype>
    void SGDSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) { 
      const vector<Blob<Dtype>*>& net_params = this->net_->learnable_params();
      const vector<float>& net_params_lr = this->net_->params_lr();
      Dtype momentum = this->param_.momentum();
      Dtype local_rate = rate * net_params_lr[param_id];
      // Compute the update to history, then copy it to the parameter diff.
      //保存结果算过程中的梯度信息
      switch (Caffe::mode()) { 
      case Caffe::CPU: { 
        caffe_cpu_axpby(net_params[param_id]->count(), local_rate,
                  net_params[param_id]->cpu_diff(), momentum,
                  history_[param_id]->mutable_cpu_data());
        caffe_copy(net_params[param_id]->count(),
            history_[param_id]->cpu_data(),
            net_params[param_id]->mutable_cpu_diff());
        break;
      }
      case Caffe::GPU: { 
    #ifndef CPU_ONLY
        sgd_update_gpu(net_params[param_id]->count(),
            net_params[param_id]->mutable_gpu_diff(),
            history_[param_id]->mutable_gpu_data(),
            momentum, local_rate);
    #else
        NO_GPU;
    #endif
        break;
      }
      default:
        LOG(FATAL) << "Unknown caffe mode: " << Caffe::mode();
      }
    }

上面讲到了网络的梯度信息是保存在一个blob数组中的，那么在调用`SnapshotSolverState()`函数的时候就会将梯度信息保存到文件中，以供后面网络进行恢复。

SnapshotSolverState(model_filename); //保存solver_state

这里保存的形式也是一样分为两种类型的情况，其实现为：

template <typename Dtype>
    void SGDSolver<Dtype>::SnapshotSolverState(const string& model_filename) { 
      switch (this->param_.snapshot_format()) { 
        case caffe::SolverParameter_SnapshotFormat_BINARYPROTO:
          SnapshotSolverStateToBinaryProto(model_filename);
          break;
        case caffe::SolverParameter_SnapshotFormat_HDF5:
          SnapshotSolverStateToHDF5(model_filename);
          break;
        default:
          LOG(FATAL) << "Unsupported snapshot format.";
      }
    }
    
    template <typename Dtype>
    void SGDSolver<Dtype>::SnapshotSolverStateToBinaryProto(
        const string& model_filename) { 
      SolverState state;
      state.set_iter(this->iter_);
      state.set_learned_net(model_filename);
      state.set_current_step(this->current_step_);
      state.clear_history();
      for (int i = 0; i < history_.size(); ++i) { 
        // Add history
        BlobProto* history_blob = state.add_history();
        history_[i]->ToProto(history_blob);
      }
      string snapshot_filename = Solver<Dtype>::SnapshotFilename(".solverstate");
      LOG(INFO)
        << "Snapshotting solver state to binary proto file " << snapshot_filename;
      WriteProtoToBinaryFile(state, snapshot_filename.c_str());
    }

在上面讲完了有关权值文件保存的过程，对于权值文件加载到网络中的流程在之前的文章中有所介绍，请参考：[Caffe源码，训练流程分析][Caffe]中的第三节。

[Caffe]: https://blog.csdn.net/m_buddy/article/details/89297483