1 #include "caffe2/operators/cosine_embedding_criterion_op.h" 5 #include "caffe2/utils/math.h" 10 bool CosineEmbeddingCriterionOp<CPUContext>::RunOnDevice() {
13 auto* output = Output(0);
16 "The embedding and label should have the same size.");
17 output->ResizeLike(S);
19 const float* Sdata = S.data<
float>();
20 const int* Ydata = Y.data<
int>();
21 float* output_data = output->mutable_data<
float>();
22 for (
int i = 0; i < S.size(); ++i) {
24 Ydata[i] == 1 ? (1.f - Sdata[i]) : std::max(0.f, Sdata[i] - margin_);
30 bool CosineEmbeddingCriterionGradientOp<CPUContext>::RunOnDevice() {
33 auto& dOutput = Input(2);
38 const float* Sdata = S.data<
float>();
39 const int* Ydata = Y.data<
int>();
40 const float* dOutput_data = dOutput.data<
float>();
41 float* dSdata = dS->mutable_data<
float>();
42 for (
int i = 0; i < S.size(); ++i) {
43 dSdata[i] = dOutput_data[i] *
44 (Ydata[i] == 1 ? -1.f :
static_cast<float>(Sdata[i] >= margin_));
49 REGISTER_CPU_OPERATOR(
50 CosineEmbeddingCriterion,
51 CosineEmbeddingCriterionOp<CPUContext>);
52 REGISTER_CPU_OPERATOR(
53 CosineEmbeddingCriterionGradient,
54 CosineEmbeddingCriterionGradientOp<CPUContext>);
56 OPERATOR_SCHEMA(CosineEmbeddingCriterion)
60 CosineEmbeddingCriterion takes two inputs: the similarity value and 61 the label, and computes the elementwise criterion output as 63 output = 1 - s, if y == 1 64 max(0, s - margin), if y == -1 66 .Input(0, "S",
"The cosine similarity as a 1-dim TensorCPU.")
67 .Input(1,
"Y",
"The label as a 1-dim TensorCPU with int value of 1 or -1.")
68 .Output(0,
"loss",
"The output loss with the same dimensionality as S.");
70 OPERATOR_SCHEMA(CosineEmbeddingCriterionGradient).NumInputs(3).NumOutputs(1);
72 class GetCosineEmbeddingCriterionGradient :
public GradientMakerBase {
73 using GradientMakerBase::GradientMakerBase;
74 vector<OperatorDef> GetGradientDefs()
override {
75 return SingleGradientDef(
76 "CosineEmbeddingCriterionGradient",
78 vector<string>{I(0), I(1), GO(0)},
79 vector<string>{GI(0)});
83 CosineEmbeddingCriterion,
84 GetCosineEmbeddingCriterionGradient);
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...