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
|
// See docs in ../ops/nn_ops.cc.
#define EIGEN_USE_THREADS
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/public/tensor_shape.h"
#include "tensorflow/core/kernels/softmax_op.h"
#include "tensorflow/core/public/tensor.h"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
template <typename Device, typename T>
class SoftmaxOp : public OpKernel {
public:
explicit SoftmaxOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* context) override {
const Tensor& logits_in = context->input(0);
OP_REQUIRES(context, TensorShapeUtils::IsMatrix(logits_in.shape()),
errors::InvalidArgument("logits must be 2-dimensional"));
Tensor* softmax_out = nullptr;
OP_REQUIRES_OK(
context, context->allocate_output(0, logits_in.shape(), &softmax_out));
functor::SoftmaxFunctor<Device, T> functor;
functor(context->eigen_device<Device>(), logits_in.matrix<T>(),
softmax_out->matrix<T>());
}
};
// Partial specialization for a CPUDevice, that uses the Eigen implementation
// from SoftmaxEigenImpl.
namespace functor {
template <typename T>
struct SoftmaxFunctor<CPUDevice, T> {
void operator()(const CPUDevice& d, typename TTypes<T>::ConstMatrix logits,
typename TTypes<T>::Matrix softmax) {
SoftmaxEigenImpl<CPUDevice, T>::Compute(d, logits, softmax);
}
};
} // namespace functor
REGISTER_KERNEL_BUILDER(Name("Softmax")
.Device(DEVICE_CPU)
.TypeConstraint<float>("T"),
SoftmaxOp<CPUDevice, float>);
REGISTER_KERNEL_BUILDER(Name("Softmax")
.Device(DEVICE_CPU)
.TypeConstraint<double>("T"),
SoftmaxOp<CPUDevice, double>);
#if GOOGLE_CUDA
REGISTER_KERNEL_BUILDER(Name("Softmax")
.Device(DEVICE_GPU)
.TypeConstraint<float>("T"),
SoftmaxOp<GPUDevice, float>);
#endif // GOOGLE_CUDA
} // namespace tensorflow
|
