Caffe2 - C++ API
A deep learning, cross platform ML framework
half_float_ops.cc
1 
17 #include "caffe2/operators/half_float_ops.h"
18 
19 namespace caffe2 {
20 OPERATOR_SCHEMA(FloatToHalf)
21  .NumInputs(1)
22  .NumOutputs(1)
23  .TensorInferenceFunction(
24  [](const OperatorDef& def, const vector<TensorShape>& in) {
25  vector<TensorShape> out;
26  const TensorShape& X = in[0];
27  out.push_back(X);
28  out[0].set_data_type(TensorProto_DataType_FLOAT16);
29 
30  return out;
31  });
32 
33 OPERATOR_SCHEMA(HalfToFloat)
34  .NumInputs(1)
35  .NumOutputs(1)
36  .TensorInferenceFunction(
37  [](const OperatorDef& def, const vector<TensorShape>& in) {
38  vector<TensorShape> out;
39  const TensorShape& X = in[0];
40  out.push_back(X);
41  out[0].set_data_type(TensorProto_DataType_FLOAT);
42 
43  return out;
44  });
45 OPERATOR_SCHEMA(Float16ConstantFill)
46  .NumInputs(0)
47  .NumOutputs(1)
48  .TensorInferenceFunction(Float16FillerTensorInference)
49  .Arg("value", "The value for the elements of the output tensor.")
50  .Arg("shape", "The shape of the output tensor.")
51  .Output(
52  0,
53  "output",
54  "Output tensor of constant values specified by 'value'");
55 
57  using GradientMakerBase::GradientMakerBase;
58  vector<OperatorDef> GetGradientDefs() override {
59  return SingleGradientDef(
60  "HalfToFloat", "", vector<string>{GO(0)}, vector<string>{GI(0)});
61  }
62 };
63 REGISTER_GRADIENT(FloatToHalf, GetFloatToHalfGradient);
64 
66  using GradientMakerBase::GradientMakerBase;
67  vector<OperatorDef> GetGradientDefs() override {
68  return SingleGradientDef(
69  "FloatToHalf", "", vector<string>{GO(0)}, vector<string>{GI(0)});
70  }
71 };
72 REGISTER_GRADIENT(HalfToFloat, GetHalfToFloatGradient);
73 NO_GRADIENT(Float16ConstantFill);
74 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.
static vector< OperatorDef > SingleGradientDef(const Args &...args)
a helper function to allow one to create one single operator def, which is usually the case for many ...