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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
|
// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
#ifndef EIGEN_CXX11_NEURAL_NETWORKS_SOFTMAX_H
#define EIGEN_CXX11_NEURAL_NETWORKS_SOFTMAX_H
namespace Eigen {
/** SoftMax
* \ingroup CXX11_NeuralNetworks_Module
*
* \brief Applies a softmax
*
* The input parameter is expected to be a col-major tensor with a rank of 2 (depth and other).
*
* The result can be assigned to a tensor of rank and dimensions equal to that of the input. The result will be laid out in col-major order.
*
*/
namespace {
struct SoftmaxOp {
SoftmaxOp(const float beta) : beta_(beta) { }
template <typename Input>
typename Input::Dimensions dimensions(const Input& input) const {
return input.dimensions();
}
template <typename Input, typename Output, typename Device>
void eval(const Input& input, Output& output, const Device& device) const
{
#if !defined(EIGEN_HAS_INDEX_LIST)
// nvcc doesn't support cxx11
Eigen::array<typename internal::traits<Input>::Index, 1> depth_dim;
depth_dim[0] = 0;
Eigen::array<typename internal::traits<Input>::Index, 2> bcast;
bcast[0] = dimensions(input)[0];
bcast[1] = 1;
DSizes<typename internal::traits<Input>::Index, 2> dims2d;
dims2d[0] = 1;
dims2d[1] = dimensions(input)[1];
#else
// Take advantage of cxx11 to give the compiler information it can use to
// optimize the code.
Eigen::IndexList<Eigen::type2index<0>> depth_dim;
Eigen::IndexList<int, Eigen::type2index<1>> bcast;
bcast.set(0, dimensions(input)[0]);
Eigen::IndexList<Eigen::type2index<1>, typename internal::traits<Input>::Index> dims2d;
dims2d.set(1, dimensions(input)[1]);
#endif
output.device(device) = ((input - input.maximum(depth_dim).eval().reshape(dims2d).broadcast(bcast)) * beta_).exp();
output.device(device) = output / (output.sum(depth_dim).eval().reshape(dims2d).broadcast(bcast));
}
private:
const float beta_;
};
}
template <typename Input>
EIGEN_ALWAYS_INLINE
static const TensorCustomUnaryOp<const SoftmaxOp, const Input>
SoftMax(const Input& input, const float beta)
{
EIGEN_STATIC_ASSERT(internal::traits<Input>::Layout == ColMajor, YOU_MADE_A_PROGRAMMING_MISTAKE);
EIGEN_STATIC_ASSERT(internal::traits<Input>::NumDimensions == 2, YOU_MADE_A_PROGRAMMING_MISTAKE);
const SoftmaxOp op(beta);
return input.customOp(op);
}
} // end namespace Eigen
#endif // EIGEN_CXX11_NEURAL_NETWORKS_SOFTMAX_H
|
