Caffe2 - C++ API
A deep learning, cross platform ML framework
accumulate_op.h
1 
17 #ifndef CAFFE2_OPERATORS_ACCUMULATE_OP_H_
18 #define CAFFE2_OPERATORS_ACCUMULATE_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 <typename T, class Context>
27 class AccumulateOp final : public Operator<Context> {
28  public:
29  AccumulateOp(const OperatorDef& operator_def, Workspace* ws)
30  : Operator<Context>(operator_def, ws),
31  gamma_(static_cast<T>(
32  OperatorBase::template GetSingleArgument<float>("gamma", 1.0))) {}
33  USE_OPERATOR_CONTEXT_FUNCTIONS;
34 
35  bool RunOnDevice() override {
36  auto& input = Input(0);
37  auto* output = Output(0);
38  if (output->dims() != input.dims()) {
39  LOG(INFO) << "Reshaping and initializing output.";
40  output->ResizeLike(input);
41  math::Set<T, Context>(
42  output->size(), 0, output->template mutable_data<T>(), &context_);
43  }
44  math::Axpby<T, Context>(
45  input.size(),
46  static_cast<T>(1),
47  input.template data<T>(),
48  gamma_,
49  output->template mutable_data<T>(),
50  &context_);
51  return true;
52  }
53 
54  protected:
55  T gamma_;
56 };
57 
58 } // namespace caffe2
59 
60 #endif // CAFFE2_OPERATORS_ACCUMULATE_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.