softmax_loss.cu 和 softmax_loss.cpp源码

  #include <algorithm>
  #include <cfloat>
  #include <vector>

  #include "caffe/layers/softmax_loss_layer.hpp"
  #include "caffe/util/math_functions.hpp"

  namespace caffe {

  template <typename Dtype>
  __global__ void SoftmaxLossForwardGPU(const int nthreads,
            const Dtype* prob_data, const Dtype* label, Dtype* loss,
            const int num, const int dim, const int spatial_dim,
            const bool has_ignore_label_, const int ignore_label_,
            Dtype* counts) {
    CUDA_KERNEL_LOOP(index, nthreads) {
      const int n = index / spatial_dim;
      const int s = index % spatial_dim;
      const int label_value = static_cast<int>(label[n * spatial_dim + s]);
      if (has_ignore_label_ && label_value == ignore_label_) {
        loss[index] = ;
        counts[index] = ;
      } else {
        loss[index] = -log(max(prob_data[n * dim + label_value * spatial_dim + s],
                        Dtype(FLT_MIN)));
        counts[index] = ;
      }
    }
  }

  template <typename Dtype>
  void SoftmaxWithLossLayer<Dtype>::Forward_gpu(
      const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
    softmax_layer_->Forward(softmax_bottom_vec_, softmax_top_vec_);
    const Dtype* prob_data = prob_.gpu_data();
    const Dtype* label = bottom[]->gpu_data();
    const int dim = prob_.count() / outer_num_;
    const int nthreads = outer_num_ * inner_num_;
    // Since this memory is not used for anything until it is overwritten
    // on the backward pass, we use it here to avoid having to allocate new GPU
    // memory to accumulate intermediate results in the kernel.
    Dtype* loss_data = bottom[]->mutable_gpu_diff();
    // Similarly, this memory is never used elsewhere, and thus we can use it
    // to avoid having to allocate additional GPU memory.
    Dtype* counts = prob_.mutable_gpu_diff();
    // NOLINT_NEXT_LINE(whitespace/operators)
    SoftmaxLossForwardGPU<Dtype><<<CAFFE_GET_BLOCKS(nthreads),
        CAFFE_CUDA_NUM_THREADS>>>(nthreads, prob_data, label, loss_data,
        outer_num_, dim, inner_num_, has_ignore_label_, ignore_label_, counts);
    Dtype loss;
    caffe_gpu_asum(nthreads, loss_data, &loss);
    Dtype valid_count = -;
    // Only launch another CUDA kernel if we actually need the count of valid
    // outputs.
    if (normalization_ == LossParameter_NormalizationMode_VALID &&
        has_ignore_label_) {
      caffe_gpu_asum(nthreads, counts, &valid_count);
    }
    top[]->mutable_cpu_data()[] = loss / get_normalizer(normalization_,
                                                          valid_count);
    if (top.size() == ) {
      top[]->ShareData(prob_);
    }
  }

  template <typename Dtype>
  __global__ void SoftmaxLossBackwardGPU(const int nthreads, const Dtype* top,
            const Dtype* label, Dtype* bottom_diff, const int num, const int dim,
            const int spatial_dim, const bool has_ignore_label_,
            const int ignore_label_, Dtype* counts) {
    const int channels = dim / spatial_dim;

    CUDA_KERNEL_LOOP(index, nthreads) {
      const int n = index / spatial_dim;
      const int s = index % spatial_dim;
      const int label_value = static_cast<int>(label[n * spatial_dim + s]);

      if (has_ignore_label_ && label_value == ignore_label_) {
        for (int c = ; c < channels; ++c) {
          bottom_diff[n * dim + c * spatial_dim + s] = ;
        }
        counts[index] = ;
      } else {
        bottom_diff[n * dim + label_value * spatial_dim + s] -= ;
        counts[index] = ;
      }
    }
  }

  template <typename Dtype>
  void SoftmaxWithLossLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
    if (propagate_down[]) {
      LOG(FATAL) << this->type()
                 << " Layer cannot backpropagate to label inputs.";
    }
    if (propagate_down[]) {
      Dtype* bottom_diff = bottom[]->mutable_gpu_diff();
      const Dtype* prob_data = prob_.gpu_data();
     const Dtype* top_data = top[]->gpu_data();
     caffe_gpu_memcpy(prob_.count() * sizeof(Dtype), prob_data, bottom_diff);
     const Dtype* label = bottom[]->gpu_data();
     const int dim = prob_.count() / outer_num_;
     const int nthreads = outer_num_ * inner_num_;
     // Since this memory is never used for anything else,
     // we use to to avoid allocating new GPU memory.
     Dtype* counts = prob_.mutable_gpu_diff();
     // NOLINT_NEXT_LINE(whitespace/operators)
     SoftmaxLossBackwardGPU<Dtype><<<CAFFE_GET_BLOCKS(nthreads),
         CAFFE_CUDA_NUM_THREADS>>>(nthreads, top_data, label, bottom_diff,
         outer_num_, dim, inner_num_, has_ignore_label_, ignore_label_, counts);

     Dtype valid_count = -;
     // Only launch another CUDA kernel if we actually need the count of valid
     // outputs.
     if (normalization_ == LossParameter_NormalizationMode_VALID &&
         has_ignore_label_) {
       caffe_gpu_asum(nthreads, counts, &valid_count);
     }
     const Dtype loss_weight = top[]->cpu_diff()[] /
         (get_normalizer(normalization_, valid_count) * Caffe::getThreadNum());
     caffe_gpu_scal(prob_.count(), loss_weight , bottom_diff);
   }
 }

 INSTANTIATE_LAYER_GPU_FUNCS_DISABLE_FP16(SoftmaxWithLossLayer);

 }  // namespace caffe

outer_num_：相当于batch_size

dim: c*w*h

spatial_dim(inner_num_):w*h

softmax_loss.cpp的代码:

outer_num_ = bottom[]->count(, softmax_axis_);
inner_num_ = bottom[]->count(softmax_axis_ + );

其实可以看出来count的只取前,不取后,(0, softmax_axis_)只取了0这一个轴