Caffe2 - C++ API
A deep learning, cross platform ML framework
cosine_embedding_criterion_op.cc
1 
17 #include "caffe2/operators/cosine_embedding_criterion_op.h"
18 
19 #include <algorithm>
20 
21 #include "caffe2/utils/math.h"
22 
23 namespace caffe2 {
24 
25 template <>
26 bool CosineEmbeddingCriterionOp<CPUContext>::RunOnDevice() {
27  auto& S = Input(0);
28  auto& Y = Input(1);
29  auto* output = Output(0);
30  CAFFE_ENFORCE(
31  S.size() == Y.size(),
32  "The embedding and label should have the same size.");
33  output->ResizeLike(S);
34 
35  const float* Sdata = S.data<float>();
36  const int* Ydata = Y.data<int>();
37  float* output_data = output->mutable_data<float>();
38  for (int i = 0; i < S.size(); ++i) {
39  output_data[i] =
40  Ydata[i] == 1 ? (1.f - Sdata[i]) : std::max(0.f, Sdata[i] - margin_);
41  }
42  return true;
43 }
44 
45 template <>
46 bool CosineEmbeddingCriterionGradientOp<CPUContext>::RunOnDevice() {
47  auto& S = Input(0);
48  auto& Y = Input(1);
49  auto& dOutput = Input(2);
50  auto* dS = Output(0);
51 
52  dS->ResizeLike(S);
53 
54  const float* Sdata = S.data<float>();
55  const int* Ydata = Y.data<int>();
56  const float* dOutput_data = dOutput.data<float>();
57  float* dSdata = dS->mutable_data<float>();
58  for (int i = 0; i < S.size(); ++i) {
59  dSdata[i] = dOutput_data[i] *
60  (Ydata[i] == 1 ? -1.f : static_cast<float>(Sdata[i] >= margin_));
61  }
62  return true;
63 }
64 
65 REGISTER_CPU_OPERATOR(
66  CosineEmbeddingCriterion,
67  CosineEmbeddingCriterionOp<CPUContext>);
68 REGISTER_CPU_OPERATOR(
69  CosineEmbeddingCriterionGradient,
70  CosineEmbeddingCriterionGradientOp<CPUContext>);
71 
72 OPERATOR_SCHEMA(CosineEmbeddingCriterion)
73  .NumInputs(2)
74  .NumOutputs(1)
75  .SetDoc(R"DOC(
76 CosineEmbeddingCriterion takes two inputs: the similarity value and
77 the label, and computes the elementwise criterion output as
78 
79  output = 1 - s, if y == 1
80  max(0, s - margin), if y == -1
81 )DOC")
82  .Input(0, "S", "The cosine similarity as a 1-dim TensorCPU.")
83  .Input(1, "Y", "The label as a 1-dim TensorCPU with int value of 1 or -1.")
84  .Output(0, "loss", "The output loss with the same dimensionality as S.");
85 
86 OPERATOR_SCHEMA(CosineEmbeddingCriterionGradient).NumInputs(3).NumOutputs(1);
87 
88 class GetCosineEmbeddingCriterionGradient : public GradientMakerBase {
89  using GradientMakerBase::GradientMakerBase;
90  vector<OperatorDef> GetGradientDefs() override {
91  return SingleGradientDef(
92  "CosineEmbeddingCriterionGradient",
93  "",
94  vector<string>{I(0), I(1), GO(0)},
95  vector<string>{GI(0)});
96  }
97 };
98 REGISTER_GRADIENT(
99  CosineEmbeddingCriterion,
100  GetCosineEmbeddingCriterionGradient);
101 
102 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.