Caffe2 - C++ API
A deep learning, cross platform ML framework
square_root_divide_op.h
1 
17 #ifndef CAFFE2_OPERATORS_SQUARE_ROOT_DIVIDE_OP_H_
18 #define CAFFE2_OPERATORS_SQUARE_ROOT_DIVIDE_OP_H_
19 
20 #include "caffe2/core/context.h"
21 #include "caffe2/core/operator.h"
22 #include "caffe2/utils/math.h"
23 
24 namespace caffe2 {
25 
26 template <class Context>
27 class SquareRootDivideOp final : public Operator<Context> {
28  public:
29  USE_OPERATOR_CONTEXT_FUNCTIONS;
30  USE_DISPATCH_HELPER;
31 
32  SquareRootDivideOp(const OperatorDef& operator_def, Workspace* ws)
33  : Operator<Context>(operator_def, ws) {}
34 
35  bool RunOnDevice() override {
36  return DispatchHelper<TensorTypes<float>>::call(this, Input(DATA));
37  }
38 
39  private:
40  template <typename TData>
41  bool DoRunWithType() {
43  this, Input(SCALE));
44  }
45 
46  template <typename TData, typename TScale>
47  bool DoRunWithType2() {
48  auto& data = Input(DATA);
49  auto& scale = Input(SCALE);
50  auto* Y = Output(0);
51  Y->ResizeLike(data);
52  size_t batchSize = data.dim(0);
53  size_t exampleSize = data.size_from_dim(1);
54  CAFFE_ENFORCE(batchSize == scale.dim(0), batchSize, " != ", scale.dim(0));
55  auto* scalePtr = scale.template data<TScale>();
56  auto* dataPtr = data.template data<TData>();
57  auto* yPtr = Y->template mutable_data<TData>();
58  for (auto i = 0; i < batchSize; ++i) {
59  auto scale = scalePtr[i];
60  CAFFE_ENFORCE(scale >= 0, scale, " < 0");
61  auto multiplier = scale == 0 ? 1.0 : 1 / std::sqrt(scale);
62  math::Scale<TData, Context>(
63  exampleSize,
64  multiplier,
65  dataPtr + i * exampleSize,
66  yPtr + i * exampleSize,
67  &context_);
68  }
69  return true;
70  }
71 
72  INPUT_TAGS(DATA, SCALE);
73 };
74 
75 } // namespace caffe2
76 
77 #endif // CAFFE2_OPERATORS_SQUARE_ROOT_DIVIDE_OP_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.