Caffe2 - C++ API
A deep learning, cross platform ML framework
channel_stats_op.h
1 
17 #ifndef CAFFE2_OPERATORS_CHANNEL_STATS_OP_H
18 #define CAFFE2_OPERATORS_CHANNEL_STATS_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 ChannelStatsOp : public Operator<Context> {
28  public:
29  USE_OPERATOR_CONTEXT_FUNCTIONS;
30  ChannelStatsOp(const OperatorDef& operator_def, Workspace* ws)
31  : Operator<Context>(operator_def, ws) {}
32  ~ChannelStatsOp() {}
33 
34  bool RunOnDevice() override {
35  return true;
36  }
37 
38  protected:
39  INPUT_TAGS(INPUT);
40  OUTPUT_TAGS(SUM, SUMSQ);
41 
42  Tensor<Context> sumScratch_;
43  Tensor<Context> sumsqScratch_;
44 };
45 
46 } // namespace caffe2
47 
48 #endif // CAFFE2_OPERATORS_CHANNEL_STATS_OP_H
Tensor is the basic class in Caffe2 that stores a contiguous memory with its shape information...
Definition: tensor.h:109
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.