Caffe2 - C++ API
A deep learning, cross platform ML framework
scale_op.h
1 
17 #ifndef CAFFE2_OPERATORS_SCALE_OP_H_
18 #define CAFFE2_OPERATORS_SCALE_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 ScaleOp final : public Operator<Context> {
28  public:
29  USE_OPERATOR_CONTEXT_FUNCTIONS;
30  ScaleOp(const OperatorDef& operator_def, Workspace* ws)
31  : Operator<Context>(operator_def, ws),
32  scale_(OperatorBase::GetSingleArgument<float>("scale", 1.0)) {}
33 
34  template <typename T>
35  bool DoRunWithType() {
36  auto& X = Input(0);
37  auto* Y = Output(0);
38  Y->ResizeLike(X);
39  math::Scale<T, Context>(
40  X.size(),
41  scale_,
42  X.template data<T>(),
43  Y->template mutable_data<T>(),
44  &context_);
45  return true;
46  }
47 
48  bool RunOnDevice() override {
49  return DispatchHelper<TensorTypes<float>>::call(this, Input(0));
50  }
51 
52  protected:
53  float scale_;
54 };
55 
56 } // namespace caffe2
57 
58 #endif // CAFFE2_OPERATORS_SCALE_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.