mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
508f676c50
Summary: Codemod generated with clangr shard mode, 50 files per diff, clangr code(ndim()->dim()): diffusion/FBS/browse/master/fbcode/caffe2/caffe2/fb/codemods/TensorMethodRename.cpp Reviewed By: salexspb Differential Revision: D12935787 fbshipit-source-id: 303d71d3eb050789af2ab9575e5dcc48f6037086
79 lines
2.5 KiB
C++
79 lines
2.5 KiB
C++
#include "caffe2/operators/sparse_normalize_op.h"
|
|
#include "caffe2/core/tensor.h"
|
|
#include "caffe2/utils/eigen_utils.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
template <>
|
|
bool SparseNormalizeOp<float, CPUContext>::RunOnDevice() {
|
|
CAFFE_ENFORCE_EQ(
|
|
Input(PARAM).size_from_dim(1),
|
|
Input(GRAD).size_from_dim(Input(INDICES).dim()));
|
|
|
|
return DispatchHelper<TensorTypes<int32_t, int64_t>>::call(
|
|
this, Input(INDICES));
|
|
}
|
|
|
|
template <>
|
|
template <typename SIndex>
|
|
bool SparseNormalizeOp<float, CPUContext>::DoRunWithType() {
|
|
const auto* indices = Input(INDICES).template data<SIndex>();
|
|
const auto* paramIn = Input(PARAM).template data<float>();
|
|
auto* paramOut = Output(OUTPUT_PARAM)->template mutable_data<float>();
|
|
const float kEps = 1e-12f;
|
|
|
|
// n: number of sparse embeddings to be normalized
|
|
auto n = Input(INDICES).numel();
|
|
if (n == 0) {
|
|
return true;
|
|
}
|
|
|
|
// embedding length, e.g. 32, 64, 128
|
|
auto block_size = Input(GRAD).numel() / n;
|
|
for (int i = 0; i < n; ++i) {
|
|
auto idx = indices[i];
|
|
auto offsetIdx = idx * block_size;
|
|
ConstEigenVectorMap<float> xVec(paramIn + offsetIdx, block_size);
|
|
auto norm = xVec.template lpNorm<2>();
|
|
|
|
if (use_max_norm_ && norm <= norm_) {
|
|
continue;
|
|
}
|
|
|
|
math::Scale(
|
|
block_size,
|
|
norm_ / (norm + kEps),
|
|
paramOut + offsetIdx,
|
|
paramOut + offsetIdx,
|
|
&context_);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
REGISTER_CPU_OPERATOR(SparseNormalize, SparseNormalizeOp<float, CPUContext>);
|
|
OPERATOR_SCHEMA(SparseNormalize)
|
|
.NumInputs(3)
|
|
.NumOutputs(1)
|
|
.Input(0, "param", "Parameters to be normalized")
|
|
.Input(1, "indices", "Sparse indices")
|
|
.Input(2, "grad", "Gradient computed")
|
|
.Output(0, "output_param", "Normalized parameters")
|
|
.EnforceOneToOneInplace()
|
|
.Arg(
|
|
"use_max_norm",
|
|
"A bool variable to control whether to use max norm \
|
|
or constant norm. When use_max_norm = false, constant norm is used so that \
|
|
all the embedding vectors are scaled to have a L2 norm equals to A \
|
|
(see blow arugment norm=A). If use_max_norm = true, \
|
|
max norm is used so that embedding is scaled so that its l2 norm is no larger \
|
|
than A. If an embedding's norm is less than A originally, \
|
|
the embedding is left unchanged.\
|
|
The default is True.")
|
|
.Arg("norm", "L2 norm of the embedding. The default is 1.0.")
|
|
.SetDoc(R"DOC(
|
|
Given a sparse matrix, apply max_norm or constant_norm sparse regularization.
|
|
)DOC");
|
|
|
|
SHOULD_NOT_DO_GRADIENT(SparseNormalize);
|
|
} // namespace caffe2
|