Caffe2 - C++ API
A deep learning, cross platform ML framework
cast_op.h
1 #pragma once
2 
3 #include "caffe2/core/context.h"
4 #include "caffe2/core/operator.h"
5 #include "caffe2/utils/cast.h"
6 #include "caffe2/utils/conversions.h"
7 #include "caffe2/utils/math.h"
8 #include "caffe2/core/logging.h"
9 #include "caffe2/core/types.h"
10 
11 namespace caffe2 {
12 
13 template <class Context>
14 class CastOp : public Operator<Context> {
15  public:
16  USE_OPERATOR_CONTEXT_FUNCTIONS;
17 
18  explicit CastOp(const OperatorDef& operator_def, Workspace* ws)
19  : Operator<Context>(operator_def, ws) {
20  const ArgumentHelper helper(operator_def);
21  TensorProto_DataType to = cast::GetCastDataType(helper, "to");
22  TensorProto_DataType from = cast::GetCastDataType(helper, "from_type");
23 
24  SetBody(to);
25  }
26 
27  bool RunOnDevice() override {
28  return (this->*body_)();
29  }
30 
31  // Allow for Context-specific implementations
32  void SetBody(TensorProto_DataType to);
33 
34  template <typename DstType>
35  bool DoRunWithDstType();
36 
37  template <typename DstType, typename SrcType>
38  bool DoRunWithType() {
39  auto& input = Input(0);
40  auto* output = Output(0);
41  output->ResizeLike(input);
42  const auto* data = input.template data<SrcType>();
43  auto* out = output->template mutable_data<DstType>();
44  auto N = input.size();
45  for (int64_t i = 0; i < N; ++i) {
46  out[i] = static_cast<DstType>(data[i]);
47  }
48  return true;
49  }
50 
51  private:
52  bool (CastOp::*body_)();
53 };
54 
55 } // namespace caffe2
A helper class to index into arguments.
Definition: proto_utils.h:200
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:47
const Tensor & Input(int idx, DeviceType type=Context::GetDeviceType())
Retrieve a non-owning reference to the input at position &#39;idx&#39; for this operator. ...
Definition: operator.h:702
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13