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
83
84
85
86
87
88
89
90
91
92
93
94
| #ifndef CAFFE_OPERATORS_BATCH_BOX_COX_OPS_H_
| #define CAFFE_OPERATORS_BATCH_BOX_COX_OPS_H_
|
| #include "caffe2/core/context.h"
| #include "caffe2/core/logging.h"
| #include "caffe2/core/operator.h"
| #include "caffe2/utils/math.h"
|
| namespace caffe2 {
|
| template <class Context>
| class BatchBoxCoxOp final : public Operator<Context> {
| public:
| USE_OPERATOR_CONTEXT_FUNCTIONS;
| template <class... Args>
| explicit BatchBoxCoxOp(Args&&... args)
| : Operator<Context>(std::forward<Args>(args)...),
| min_block_size_(
| this->template GetSingleArgument<int>("min_block_size", 256)) {}
|
| bool RunOnDevice() override {
| return DispatchHelper<TensorTypes<float, double>>::call(this, Input(DATA));
| }
|
| template <typename T>
| bool DoRunWithType();
|
| protected:
| template <typename T>
| void BoxCoxNaive(
| int64_t N,
| int64_t D,
| const T* data_ptr,
| const T* lambda1_ptr,
| const T* lambda2_ptr,
| T k_eps,
| T* output_ptr);
|
| #ifdef CAFFE2_USE_MKL
| template <typename T>
| void BoxCoxNonzeroLambda(
| int64_t D,
| const T* data_ptr,
| const T* lambda1,
| const T* lambda2,
| T k_eps,
| T* output_ptr);
|
| template <typename T>
| void BoxCoxZeroLambda(
| int64_t D,
| const T* data_ptr,
| const T* lambda2,
| T k_eps,
| T* output_ptr);
|
| template <typename T>
| void BoxCoxMixedLambda(
| const T* data_ptr,
| const vector<int>& nonzeros,
| const vector<int>& zeros,
| const T* lambda1,
| const T* lambda2,
| const T* lambda2_z,
| T k_eps,
| T* buffer,
| T* output_ptr);
|
| vector<int> nonzeros_, zeros_;
|
| // Buffers used by the MKL version are cached across calls.
| struct CachedBuffers {
| virtual ~CachedBuffers() {}
| int type_;
| };
| template <typename T>
| struct TypedCachedBuffers : public CachedBuffers {
| vector<T> lambda1_, lambda2_, lambda2_z_;
| vector<T> accumulator_;
| };
| template <typename T>
| TypedCachedBuffers<T>& GetBuffers();
| unique_ptr<CachedBuffers> buffers_;
|
| #endif // CAFFE2_USE_MKL
|
| int min_block_size_;
|
| INPUT_TAGS(DATA, LAMBDA1, LAMBDA2);
| };
|
| } // namespace caffe2
|
| #endif // CAFFE_OPERATORS_BATCH_BOX_COX_OPS_H_
|
|
