详解softmax与softmax loss的前后向推导及caffe源码实现

本文地址：http://blog.youkuaiyun.com/isMarvellous/article/details/78735208，转载请注明出处。

Softmax层的作用是将输入的预测向量转化为概率值，也就是每个元素介于0和1之间，其和为1。而Softmax loss是基于Softmax的输出，使用多元交叉熵损失函数得到的loss。下面我们来讨论一下他们其中的正向和反向导数推导，以及caffe中的源码实现。为了更好地将推导和代码相结合，以加深理解，本文将会在每个推导部分直接紧跟其代码实现。

1. Softmax

1.1 前向计算

1.1.1 公式推导

假设有K个类别，前面已得出每个类别的分值为$z_i$，则Softmax通过下式计算出相应的概率值：

$$Softmax(z_i) = \frac{exp(z_i)}{\sum_j{exp(z_j)}}, i = 0,1,2,\dots, K-1$$
这样就将 $z_i$映射到了[0,1]，且和为1，即为输入被预测到每个类别的概率。
前向过程比较简单，下面我们来看一下具体实现。

1.1.2 源码实现

我们主要分析Softmax层的Forward_cpu函数，该函数的实现位于caffe的src/caffe/layers/softmax_layer.cpp中。需要说明的是，在caffe的实现中，输入值$z_i$首先减去了最大值，这样避免了后续的exp()计算中可能出现的因数值过大而造成的溢出问题。
首先来解释一下下面代码里几个不太好理解的变量：
scale_data：
是个中间变量，用来存放计算的中间结果。
inner_num_：
在softmax_layer.hpp的声明中为inner_num_ = bottom[0]->count(softmax_axis_ + 1);也就是所有表示概率值的维度的像素点总数。
outer_num_：
在softmax_layer.hpp的声明中为outer_num_ = bottom[0]->count(0, softmax_axis_);可以理解为样本的个数。

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();    // scale_data 是个中间变量，用来存放计算的中间结果
  int channels = bottom[0]->shape(softmax_axis_);
  int dim = bottom[0]->count() / outer_num_;
  // 输出数据初始化为输入数据
  caffe_copy(bottom[0]->count(), bottom_data, top_data);
  // 我们需要减去最大值，计算exp，然后归一化。
  for (int i = 0; i < outer_num_; ++i) {
    // 将中间变量scale_data初始化为输入值的第一个样本平面
    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]);
      }
    }
    // 减去最大值
    caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, channels, inner_num_,
        1, -1., sum_multiplier_.cpu_data(), scale_data, 1., top_data);
    // 计算exp()
    caffe_exp<Dtype>(dim, top_data, top_data);
    // 求和
    caffe_cpu_gemv<Dtype>(CblasTrans, channels, inner_num_, 1.,
        top_data, sum_multiplier_.cpu_data(), 0., scale_data);
    // 除以前面求到的和
    for (int j = 0; j < channels; j++) {
      caffe_div(inner_num_, top_data, scale_data, top_data);
      top_data += inner_num_;    // 指针后移
    }
  }
}

1.2 反向传播

1.2.1 公式推导

如前，设Softmax的输入为$z_i$，输出为$a_i$，那么由链式法则，损失loss对其输入$z_i$的偏导可以如下计算：

$$\frac{\partial loss}{\partial z} = \frac{\partial loss}{\partial a} \cdot \frac{\partial a}{\partial z}$$
其中 $\frac{\partial loss}{\partial a}$是上面的层传回来的梯度，对本层来说是已知的，所以我们只需计算 $\frac{\partial a}{\partial z}$。
由 $a_i = \frac{e^{z_i}}{\sum_j{e^{z_j}}}$：
当 $i=j$，
$$\frac{\partial a_i}{\partial z_i} = \frac{e^{z_i}\sum_j{e^{z_j}} - e^{z_i}e^{z_i}}{(\sum_j{e^{z_j}})^2} = a_i - a_i * a_i$$
这里 $*$表示标量算数乘法。
当$i \neq j$，
$$\frac{\partial a_i}{\partial z_j} = \frac{ - e^{z_i}e^{z_j}}{(\sum_l{e^{z_l}})^2} = a_i * a_j$$
所以，
$$\begin{align} \frac{\partial loss}{\partial z_k} & = \sum_i \frac{\partial loss}{\partial a_i} * \frac{\partial a_i}{\partial z_k} \\ & = \sum_{i \neq k} \frac{\partial loss}{\partial a_i} * \frac{\partial a_i}{\partial z_k} + \frac{\partial loss}{\partial a_k} * \frac{\partial a_k}{\partial z_k} \\ & = \sum_{i \neq k} \frac{\partial loss}{\partial a_i} * ( - a_i * a_k ) + \frac{\partial loss}{\partial a_k} * ( a_k - a_k * a_k ) \\ & = \sum_i \frac{\partial loss}{\partial a_i} * ( - a_i * a_k ) + \frac{\partial loss}{\partial a_k} * a_k \\ & = - \frac{\partial loss}{\partial a} \cdot a * a_k + \frac{\partial loss}{\partial a_k} * a_k \\ & = ( \frac{\partial loss}{\partial a_k} - \frac{\partial loss}{\partial a} \cdot a ) * a_k \end{align}$$
这里的 $\cdot$ 表示向量点乘。
上式写成向量形式，即为：
$$\frac{\partial loss}{\partial z} = ( \frac{\partial loss}{\partial a} - \frac{\partial loss}{\partial a} \cdot a ) \cdot a$$
即为 ( top_diff - top_diff $\cdot$ top_data) $\cdot$ top_data。

1.2.2 源码实现

下面我们主要分析Softmax层的Backward_cpu函数，该函数的实现位于caffe的src/caffe/layers/softmax_layer.cpp中。

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);    // 将bottom_diff初始化为top_diff的值
  for (int i = 0; i < outer_num_; ++i) {
    // 计算开始
    for (int k = 0; k < inner_num_; ++k) {
    // 计算dot(top_diff, top_data)
      scale_data[k] = caffe_cpu_strided_dot<Dtype>(channels,
          bottom_diff + i * dim + k, inner_num_,
          top_data + i * dim + k, inner_num_);
    }
    // 相减
    caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, channels, inner_num_, 1,
        -1., sum_multiplier_.cpu_data(), scale_data, 1., bottom_diff + i * dim);
  }
  // 对应元素相乘
  caffe_mul(top[0]->count(), bottom_diff, top_data, bottom_diff);
}

2. Softmax Loss

2.1 前向计算

2.1.1 公式推导

Softmax Loss就是用Softmax的输出概率作为预测概率值，与真实label做交叉熵损失，在caffe中也是调用了Softmax layer来实现前向传播。先留个坑，以后有空慢慢写。