caffe源码学习（2）-softmax loss层

回顾softmax回归原理：

$$p_j^i=p(y^{(i)}=j|x^{(i)})=\frac{e^{\theta_j^Tx^{(i)}}}{\sum_{k}e^{\theta_k^Tx^{(i)}}}$$
$$loss=-\sum_{i}y^ilog(p_j^i)$$
$$loss=\frac{loss}{M}$$
　　严格的说，在caffe中，完整的softmax回归由 Inner Product(ip2) layer 和 softmax两个层构成，其中ip2实现内积操作 $z_j=\theta_j^Tx^{(i)}$,作为softmax loss层的输入。softmax loss层实现正向传播计算损失函数，反向传播计算梯度。

正向传播、反向传播主要主要在如下两个文件中：softmax_loss_layer.cpp和softmax_layer.cpp。
softmax_loss_layer.cpp 前向计算后向计算：

//softmax_loss_layer.cpp 前向计算后向计算
template <typename Dtype>
void SoftmaxWithLossLayer<Dtype>::Forward_cpu(
    const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
  // The forward pass computes the softmax prob values.
  softmax_layer_->Forward(softmax_bottom_vec_, softmax_top_vec_);//SoftmaxLayer类型成员变量，内部详细实现softmax计算流程
  const Dtype* prob_data = prob_.cpu_data();
  const Dtype* label = bottom[1]->cpu_data();
  int dim = prob_.count() / outer_num_;
  int count = 0;
  Dtype loss = 0;
  for (int i = 0; i < outer_num_; ++i) {  //遍历所有样本
    for (int j = 0; j < inner_num_; j++) {
      const int label_value = static_cast<int>(label[i * inner_num_ + j]);
      if (has_ignore_label_ && label_value == ignore_label_) {
        continue;
      } 
      DCHECK_GE(label_value, 0);
      DCHECK_LT(label_value, prob_.shape(softmax_axis_));
      //计算loss值
      loss -= log(std::max(prob_data[i * dim + label_value * inner_num_ + j],
                           Dtype(FLT_MIN)));
      ++count;
    }
  }
  top[0]->mutable_cpu_data()[0] = loss / get_normalizer(normalization_, count);
  if (top.size() == 2) {
    top[1]->ShareData(prob_);
  }
}

softmax主要的计算过程在SoftmaxLayer中实现,与原理不同之处，实现过程中减去了max最大值。

template <typename Dtype>
void SoftmaxLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  softmax_axis_ =
      bottom[0]->CanonicalAxisIndex(this->layer_param_.softmax_param().axis());
  top[0]->ReshapeLike(*bottom[0]);
  vector<int> mult_dims(1, bottom[0]->shape(softmax_axis_));
  sum_multiplier_.Reshape(mult_dims);
  Dtype* multiplier_data = sum_multiplier_.mutable_cpu_data();
  caffe_set(sum_multiplier_.count(), Dtype(1), multiplier_data); //此处将sum_multiplier_初始化为元素为1的列向量
  outer_num_ = bottom[0]->count(0, softmax_axis_);
  inner_num_ = bottom[0]->count(softmax_axis_ + 1);
  vector<int> scale_dims = bottom[0]->shape();
  scale_dims[softmax_axis_] = 1;
  scale_.Reshape(scale_dims);
}

template <typename Dtype>
void SoftmaxLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top) {
  const Dtype* bottom_data = bottom[0]->cpu_data();
  Dtype* top_data = top[0]->mutable_cpu_data();
  Dtype* scale_data = scale_.mutable_cpu_data();
  int channels = bottom[0]->shape(softmax_axis_);
  int dim = bottom[0]->count() / outer_num_;
  caffe_copy(bottom[0]->count(), bottom_data, top_data); //ip2的内积操作结果直接拷贝至top层
  // We need to subtract the max to avoid numerical issues, compute the exp,
  // and then normalize.
  for (int i = 0; i < outer_num_; ++i) {
    // initialize scale_data to the first plane
    caffe_copy(inner_num_, bottom_data + i * dim, scale_data); // 初始化scale_data变量，内部用于存储类别概率的最大值
    for (int j = 0; j < channels; j++) {
      for (int k = 0; k < inner_num_; k++) {
        scale_data[k] = std::max(scale_data[k],         /// 保存类别概率的最大值
            bottom_data[i * dim + j * inner_num_ + k]);
      }
    }
    // subtraction  减去最大值，防止溢出；使用矩阵操作，减去最大值top_data-scale_data
    //注意：sum_multiplier_被初始化为1的向量 M=10,N=1,K=1 ,alpha=-1,beta=1
    caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, channels, inner_num_,
        1, -1., sum_multiplier_.cpu_data(), scale_data, 1., top_data);
    // exponentiation  求自然对数
    caffe_exp<Dtype>(dim, top_data, top_data);
    // sum after exp  利用内积操作求和，注意sum_multiplier_内部元素为1
    caffe_cpu_gemv<Dtype>(CblasTrans, channels, inner_num_, 1.,
        top_data, sum_multiplier_.cpu_data(), 0., scale_data);
    // division  除以归一化因子
    for (int j = 0; j < channels; j++) {
      caffe_div(inner_num_, top_data, scale_data, top_data);
      top_data += inner_num_;//top_data地址偏移更新
    }
  }
}
/// 没有被调用
template <typename Dtype>
void SoftmaxLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down,
    const vector<Blob<Dtype>*>& bottom) {
  const Dtype* top_diff = top[0]->cpu_diff();
  const Dtype* top_data = top[0]->cpu_data();
  Dtype* bottom_diff = bottom[0]->mutable_cpu_diff();
  Dtype* scale_data = scale_.mutable_cpu_data();
  int channels = top[0]->shape(softmax_axis_);
  int dim = top[0]->count() / outer_num_;
  caffe_copy(top[0]->count(), top_diff, bottom_diff);
  for (int i = 0; i < outer_num_; ++i) {
    // compute dot(top_diff, top_data) and subtract them from the bottom diff
    for (int k = 0; k < inner_num_; ++k) {
      scale_data[k] = caffe_cpu_strided_dot<Dtype>(channels,
          bottom_diff + i * dim + k, inner_num_,
          top_data + i * dim + k, inner_num_);
    }
    // subtraction
    caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, channels, inner_num_, 1,
        -1., sum_multiplier_.cpu_data(), scale_data, 1., bottom_diff + i * dim);
  }
  // elementwise multiplication
  caffe_mul(top[0]->count(), bottom_diff, top_data, bottom_diff);
}

2 softmax loss反向运算

某个样本，softmax loss 设层的输入为$z，z$为向量。 softmax 层输出为$f(z_j)=\frac{e^{z_j}}{\sum_{k}e^{z_k}}$，设$y$为真实类别标签，则softmax loss层损失函数为：

$$L(y,z)=-logf(z_y)=log\sum_{k}e^{z_k}-{z_y}$$ 对输入$z$求导：
$$\frac{\partial L(y,z)}{\partial z_k}=\begin{cases} & \text{ if } k=y, \frac{e^{z_y}}{\sum_{k}e^{z_k}}-1=f(z_y)-1 \\ & \text{ if } k\neq y, \frac{e^{z_k}}{\sum_{k}e^{z_k}}=f(z_k) \end{cases}$$
其中

//softmax反向求导计算
template <typename Dtype>
void SoftmaxWithLossLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
  if (propagate_down[1]) {
    LOG(FATAL) << this->type()
               << " Layer cannot backpropagate to label inputs.";
  }
  if (propagate_down[0]) {
    Dtype* bottom_diff = bottom[0]->mutable_cpu_diff();
    const Dtype* prob_data = prob_.cpu_data();
    caffe_copy(prob_.count(), prob_data, bottom_diff); //将softmax输出拷贝至softmax loos 层bottom_diff; 即 f(z)
    const Dtype* label = bottom[1]->cpu_data();
    int dim = prob_.count() / outer_num_;
    int count = 0;
    for (int i = 0; i < outer_num_; ++i) { //遍历所有样本，计算deta
      for (int j = 0; j < inner_num_; ++j) {
        const int label_value = static_cast<int>(label[i * inner_num_ + j]);
        if (has_ignore_label_ && label_value == ignore_label_) {
          for (int c = 0; c < bottom[0]->shape(softmax_axis_); ++c) {
            bottom_diff[i * dim + c * inner_num_ + j] = 0;
          }
        } else {
          bottom_diff[i * dim + label_value * inner_num_ + j] -= 1;  /// deta; f(z,y)-1
          ++count;
        }
      }
    }
    // Scale gradient
    Dtype loss_weight = top[0]->cpu_diff()[0] /
                        get_normalizer(normalization_, count);
    caffe_scal(prob_.count(), loss_weight, bottom_diff);
  }
}