1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
|
#include "caffe2/operators/softmax_utils.h"
#include "caffe2/core/context.h"
#include "caffe2/utils/eigen_utils.h"
#include "caffe2/utils/math.h"
namespace caffe2 {
namespace softmax_utils {
#define CAFFE2_SPECIALIZED_SOFTMAX_CPU(T) \
template <> \
void SoftmaxCPU<T>( \
const int N, \
const int D, \
const bool logarithmic, \
const T* X, \
T* Y, \
T* scratch, \
CPUContext* context) { \
ConstEigenArrayMap<T> X_arr(X, D, N); \
EigenArrayMap<T> Y_arr(Y, D, N); \
EigenVectorArrayMap<T> scratch_arr(scratch, N); \
scratch_arr = X_arr.colwise().maxCoeff().transpose(); \
Y_arr = X_arr.rowwise() - scratch_arr.transpose(); \
math::Exp<T, CPUContext>(N * D, Y, Y, context); \
if (logarithmic) { \
scratch_arr += Y_arr.colwise().sum().log().transpose(); \
Y_arr = X_arr.rowwise() - scratch_arr.transpose(); \
} else { \
scratch_arr = Y_arr.colwise().sum().inverse().transpose(); \
Y_arr = Y_arr.rowwise() * scratch_arr.transpose(); \
} \
}
CAFFE2_SPECIALIZED_SOFTMAX_CPU(float)
#undef CAFFE2_SPECIALIZED_SOFTMAX_CPU
} // namespace softmax_utils
} // namespace caffe2
|
