Caffe2 - C++ API
A deep learning, cross platform ML framework
sparse_normalize_op.cc
1 #include "caffe2/operators/sparse_normalize_op.h"
2 #include "caffe2/core/tensor.h"
3 #include "caffe2/utils/eigen_utils.h"
4 
5 namespace caffe2 {
6 
7 template <>
8 bool SparseNormalizeOp<float, CPUContext>::RunOnDevice() {
9  CAFFE_ENFORCE_EQ(
10  Input(PARAM).size_from_dim(1),
11  Input(GRAD).size_from_dim(Input(INDICES).dim()));
12 
13  return DispatchHelper<TensorTypes<int32_t, int64_t>>::call(
14  this, Input(INDICES));
15 }
16 
17 template <>
18 template <typename SIndex>
19 bool SparseNormalizeOp<float, CPUContext>::DoRunWithType() {
20  const auto* indices = Input(INDICES).template data<SIndex>();
21  const auto* paramIn = Input(PARAM).template data<float>();
22  auto* paramOut = Output(OUTPUT_PARAM)->template mutable_data<float>();
23  const float kEps = 1e-12f;
24 
25  // n: number of sparse embeddings to be normalized
26  auto n = Input(INDICES).numel();
27  if (n == 0) {
28  return true;
29  }
30 
31  // embedding length, e.g. 32, 64, 128
32  auto block_size = Input(GRAD).numel() / n;
33  for (int i = 0; i < n; ++i) {
34  auto idx = indices[i];
35  auto offsetIdx = idx * block_size;
36  ConstEigenVectorMap<float> xVec(paramIn + offsetIdx, block_size);
37  auto norm = xVec.template lpNorm<2>();
38 
39  if (use_max_norm_ && norm <= norm_) {
40  continue;
41  }
42 
43  math::Scale(
44  block_size,
45  norm_ / (norm + kEps),
46  paramOut + offsetIdx,
47  paramOut + offsetIdx,
48  &context_);
49  }
50  return true;
51 }
52 
53 REGISTER_CPU_OPERATOR(SparseNormalize, SparseNormalizeOp<float, CPUContext>);
54 OPERATOR_SCHEMA(SparseNormalize)
55  .NumInputs(3)
56  .NumOutputs(1)
57  .Input(0, "param", "Parameters to be normalized")
58  .Input(1, "indices", "Sparse indices")
59  .Input(2, "grad", "Gradient computed")
60  .Output(0, "output_param", "Normalized parameters")
61  .EnforceOneToOneInplace()
62  .Arg(
63  "use_max_norm",
64  "A bool variable to control whether to use max norm \
65  or constant norm. When use_max_norm = false, constant norm is used so that \
66  all the embedding vectors are scaled to have a L2 norm equals to A \
67  (see blow arugment norm=A). If use_max_norm = true, \
68  max norm is used so that embedding is scaled so that its l2 norm is no larger \
69  than A. If an embedding's norm is less than A originally, \
70  the embedding is left unchanged.\
71  The default is True.")
72  .Arg("norm", "L2 norm of the embedding. The default is 1.0.")
73  .SetDoc(R"DOC(
74 Given a sparse matrix, apply max_norm or constant_norm sparse regularization.
75 )DOC");
76 
77 SHOULD_NOT_DO_GRADIENT(SparseNormalize);
78 } // namespace caffe2
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13