/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/sgd
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/sgd/adagrad_op.h (17993B)
#pragma once
#include "caffe2/core/operator.h"
#include "caffe2/perfkernels/adagrad.h"
#if defined(USE_FBGEMM) && !defined(__NVCC__)
#include "fbgemm/FbgemmEmbedding.h"
#endif
namespace caffe2 {
template
void adagrad_update(
int N,
const float* w,
const float* g,
const float* h,
float* nw,
float* nh,
float epsilon,
float decay,
const float* lr,
Context* /*context*/,
float weight_decay = 0.f) {
return adagrad_update(
N, w, g, h, nw, nh, epsilon, decay, lr[0], weight_decay);
}
template
void adagrad_update_output_effective_lr(
int N,
const float* paramIn,
const float* gradIn,
const float* momentIn,
float* paramOut,
float* momentOut,
float* effectiveLROut,
float epsilon,
float decay,
const float* lr,
Context* /*context*/,
float weight_decay = 0.f) {
for (auto i = 0; i < N; ++i) {
float grad = std::fma(weight_decay, paramIn[i], gradIn[i]);
float moment = momentOut[i] = decay * momentIn[i] + grad * grad;
float effective_lr = effectiveLROut[i] =
lr[0] / (std::sqrt(moment) + epsilon);
paramOut[i] = paramIn[i] + effective_lr * grad;
}
}
template
void adagrad_update_output_effective_lr_and_update(
int N,
const float* paramIn,
const float* gradIn,
const float* momentIn,
float* paramOut,
float* momentOut,
float* effectiveLROut,
float* updateOut,
float epsilon,
float decay,
const float* lr,
Context* /*context*/,
float weight_decay = 0.f) {
for (auto i = 0; i < N; ++i) {
float grad = std::fma(weight_decay, paramIn[i], gradIn[i]);
float moment = momentOut[i] = decay * momentIn[i] + grad * grad;
float effective_lr = effectiveLROut[i] =
lr[0] / (std::sqrt(moment) + epsilon);
float update = updateOut[i] = effective_lr * grad;
paramOut[i] = paramIn[i] + update;
}
}
template
class AdagradOp final : public Operator {
public:
USE_OPERATOR_CONTEXT_FUNCTIONS;
AdagradOp(const OperatorDef& operator_def, Workspace* ws)
: Operator(operator_def, ws),
epsilon_(this->template GetSingleArgument("epsilon", 1e-5f)),
decay_(this->template GetSingleArgument("decay", 1.0f)),
weight_decay_(
this->template GetSingleArgument("weight_decay", 0.f)) {
VLOG(1) << "gradient optimization operator in use: "
<< "AdagradOp"
<< " weight_decay_=" << weight_decay_;
}
bool RunOnDevice() override {
CAFFE_ENFORCE_EQ(
Input(GRAD).numel(),
Input(MOMENT_1).numel(),
"PARAM size: ",
Input(PARAM).numel(),
", GRAD size: ",
Input(GRAD).numel(),
", MOMENT_1 size: ",
Input(MOMENT_1).numel(),
", LR size: ",
Input(LR).numel());
CAFFE_ENFORCE_EQ(Input(GRAD).numel(), Input(PARAM).numel());
Output(OUTPUT_PARAM)->ResizeLike(Input(PARAM));
Output(OUTPUT_MOMENT_1)->ResizeLike(Input(MOMENT_1));
if (OutputSize() == 2) {
adagrad_update(
Input(GRAD).numel(),
Input(PARAM).template data(),
Input(GRAD).template data(),
Input(MOMENT_1).template data(),
Output(OUTPUT_PARAM)->template mutable_data(),
Output(OUTPUT_MOMENT_1)->template mutable_data(),
epsilon_,
decay_,
Input(LR).template data(),
&context_,
weight_decay_);
} else if (OutputSize() == 3) {
Output(OUTPUT_EFFECTIVE_LR)->ResizeLike(Input(GRAD));
adagrad_update_output_effective_lr(
Input(GRAD).numel(),
Input(PARAM).template data(),
Input(GRAD).template data(),
Input(MOMENT_1).template data(),
Output(OUTPUT_PARAM)->template mutable_data(),
Output(OUTPUT_MOMENT_1)->template mutable_data(),
Output(OUTPUT_EFFECTIVE_LR)->template mutable_data(),
epsilon_,
decay_,
Input(LR).template data(),
&context_,
weight_decay_);
} else {
Output(OUTPUT_EFFECTIVE_LR)->ResizeLike(Input(GRAD));
Output(OUTPUT_UPDATE)->ResizeLike(Input(GRAD));
adagrad_update_output_effective_lr_and_update(
Input(GRAD).numel(),
Input(PARAM).template data(),
Input(GRAD).template data(),
Input(MOMENT_1).template data(),
Output(OUTPUT_PARAM)->template mutable_data(),
Output(OUTPUT_MOMENT_1)->template mutable_data(),
Output(OUTPUT_EFFECTIVE_LR)->template mutable_data(),
Output(OUTPUT_UPDATE)->template mutable_data(),
epsilon_,
decay_,
Input(LR).template data(),
&context_,
weight_decay_);
}
return true;
}
protected:
float epsilon_;
float decay_;
float weight_decay_;
INPUT_TAGS(PARAM, MOMENT_1, GRAD, LR);
OUTPUT_TAGS(
OUTPUT_PARAM,
OUTPUT_MOMENT_1,
OUTPUT_EFFECTIVE_LR,
OUTPUT_UPDATE);
};
class SparseAdagradOp final : public Operator {
public:
SparseAdagradOp(const OperatorDef& operator_def, Workspace* ws)
: Operator(operator_def, ws),
epsilon_(this->template GetSingleArgument("epsilon", 1e-5f)),
weight_decay_(
this->template GetSingleArgument("weight_decay", 0.f)) {
VLOG(1) << "gradient optimization operator in use: "
<< "SparseAdagradOp"
<< " weight_decay_=" << weight_decay_;
const float decay = this->template GetSingleArgument("decay", 1.0);
CAFFE_ENFORCE_EQ(
decay, 1.0, "Decay is not supported for SparseSimdAdagradOp");
}
bool RunOnDevice() override {
// Enforce shapes
// input(embedding/momentum) == outputs(embedding/momentum)
CAFFE_ENFORCE_EQ(
Input(PARAM).numel(),
Input(MOMENT_1).numel(),
"Input Param size: ",
Input(PARAM).numel(),
" Input Moment size: ",
Input(MOMENT_1).numel());
CAFFE_ENFORCE_EQ(Input(LR).numel(), 1);
CAFFE_ENFORCE_EQ(
Input(PARAM).size_from_dim(1),
Input(GRAD).size_from_dim(Input(INDICES).dim()));
return DispatchHelper>::call(
this, Input(INDICES));
}
template
bool DoRunWithType() {
const auto* lr = Input(LR).template data();
auto n = Input(INDICES).numel();
const auto* indices = Input(INDICES).template data();
const auto* gradIn = Input(GRAD).template data();
auto* paramOut = Output(OUTPUT_PARAM)->template mutable_data();
auto* momentOut = Output(OUTPUT_MOMENT_1)->template mutable_data();
if (n == 0) {
return true;
}
auto block_size = Input(GRAD).numel() / n;
// input(grad) is compatible with size of indexes
CAFFE_ENFORCE_EQ(
Input(GRAD).numel() % n,
0,
"Incorrect gradient size:",
Input(GRAD).numel(),
" size of indexes:",
n);
#if defined(USE_FBGEMM) && !defined(__NVCC__)
VLOG(1) << "using fbgemm::GenerateSparseAdaGrad in SparseAdagradOp";
if (block_size != last_block_size_) {
last_block_size_ = block_size;
if (std::is_same::value) {
kernel_i32_ = fbgemm::GenerateSparseAdaGrad(
block_size,
/*rowwise=*/false,
/*prefetch=*/16,
weight_decay_ != 0.0f);
} else {
CAFFE_ENFORCE((std::is_same::value));
kernel_i64_ = fbgemm::GenerateSparseAdaGrad(
block_size,
/*rowwise=*/false,
/*prefetch=*/16,
weight_decay_ != 0.0f);
}
}
int num_rows_processed;
if (std::is_same::value) {
num_rows_processed = kernel_i32_(
n,
Input(PARAM).numel(),
paramOut,
gradIn,
momentOut,
reinterpret_cast(indices),
epsilon_,
lr[0],
weight_decay_,
/*counter=*/nullptr,
/*counter_halflife=*/0);
} else {
num_rows_processed = kernel_i64_(
n,
Input(PARAM).numel(),
paramOut,
gradIn,
momentOut,
reinterpret_cast(indices),
epsilon_,
lr[0],
weight_decay_,
/*counter=*/nullptr,
/*counter_halflife=*/0);
}
if (num_rows_processed < n) {
CAFFE_ENFORCE_GE(
Input(PARAM).numel(),
(indices[num_rows_processed] + 1) * block_size,
this->debug_def().input(PARAM),
", out of bound, idx:",
indices[num_rows_processed],
" for input i:",
num_rows_processed,
" and block_size:",
block_size,
" max size:",
Input(PARAM).numel());
return false;
} else {
return true;
}
#endif
VLOG(1)
<< "using internal::adagrad_update_prefetch_inlined in SparseAdagradOp";
const auto* paramIn = Input(PARAM).template data();
const auto* momentIn = Input(MOMENT_1).template data();
std::vector grad(block_size);
for (auto i = 0; i < n; ++i) {
auto idx = indices[i];
auto offsetI = i * block_size;
auto offsetIdx = idx * block_size;
// Enforce:
// access within range
// gradient access within range
CAFFE_ENFORCE_GE(
Input(PARAM).numel(),
block_size + offsetIdx,
this->debug_def().input(PARAM),
", out of bound, idx:",
idx,
" for input i:",
i,
" and block size:",
block_size,
" max size:",
Input(PARAM).numel());
if (block_size == 1) {
float gi = std::fma(weight_decay_, paramIn[idx], gradIn[i]);
float hi = momentOut[idx] = momentIn[idx] + gi * gi;
paramOut[idx] = paramIn[idx] + lr[0] * gi / (std::sqrt(hi) + epsilon_);
} else {
// prefetching
const int prefdist_T0 = 16;
int i_pref = (i < n - prefdist_T0) ? i + prefdist_T0 : i;
std::size_t idx_pref = indices[i_pref];
internal::adagrad_update_prefetch_inlined(
block_size,
paramIn + offsetIdx,
¶mIn[idx_pref * block_size],
gradIn + offsetI,
momentIn + offsetIdx,
&momentIn[idx_pref * block_size],
paramOut + offsetIdx,
¶mOut[idx_pref * block_size],
momentOut + offsetIdx,
&momentOut[idx_pref * block_size],
epsilon_,
lr[0],
weight_decay_);
}
}
return true;
}
protected:
float epsilon_;
const float weight_decay_;
#if defined(USE_FBGEMM) && !defined(__NVCC__)
fbgemm::SparseAdaGradSignature::Type kernel_i32_;
fbgemm::SparseAdaGradSignature::Type kernel_i64_;
std::int64_t last_block_size_{-1};
#endif
INPUT_TAGS(PARAM, MOMENT_1, INDICES, GRAD, LR);
OUTPUT_TAGS(OUTPUT_PARAM, OUTPUT_MOMENT_1);
};
template
class RowWiseSparseAdagradOp final : public Operator {
public:
USE_OPERATOR_CONTEXT_FUNCTIONS;
RowWiseSparseAdagradOp(const OperatorDef& operator_def, Workspace* ws)
: Operator(operator_def, ws),
epsilon_(this->template GetSingleArgument("epsilon", 1e-5f)),
weight_decay_(
this->template GetSingleArgument("weight_decay", 0.f)),
counter_halflife_(
this->template GetSingleArgument("counter_halflife", -1)) {
VLOG(1) << "gradient optimization operator in use: "
<< "RowWiseSparseAdagradOp"
<< " weight_decay_=" << weight_decay_
<< " counter_halflife=" << counter_halflife_;
}
bool RunOnDevice() override {
// Enforce shapes
CAFFE_ENFORCE_EQ(Input(PARAM).sizes()[0], Input(MOMENT_1).numel());
CAFFE_ENFORCE_EQ(Input(LR).numel(), 1);
CAFFE_ENFORCE_EQ(
Input(PARAM).size_from_dim(1),
Input(GRAD).size_from_dim(Input(INDICES).dim()));
return DispatchHelper>::call(
this, Input(INDICES));
}
template
bool DoRunWithType() {
const auto* lr = Input(LR).template data();
auto* param = Output(OUTPUT_PARAM)->template mutable_data();
auto* moment = Output(OUTPUT_MOMENT_1)->template mutable_data();
const auto* indices = Input(INDICES).template data();
const auto* gradIn = Input(GRAD).template data();
const auto* count = counter_halflife_ == -1
? nullptr
: Input(COUNTER).template data();
auto n = Input(INDICES).numel();
if (n == 0) {
return true;
}
auto block_size = Input(GRAD).numel() / n;
// Enforce:
// Input(embedding/momentum) == outputs(embedding/momentum)
CAFFE_ENFORCE_EQ(
Input(PARAM).numel() / block_size,
Input(MOMENT_1).numel(),
"Input Param size: ",
Input(PARAM).numel(),
" Block size: ",
block_size,
" Input Moment size: ",
Input(MOMENT_1).numel());
// input(grad) is compatible with size of indexes
CAFFE_ENFORCE_EQ(
Input(GRAD).numel() % n,
0,
"Incorrect gradient size:",
Input(GRAD).numel(),
" size of indexes:",
n);
#if defined(USE_FBGEMM) && !defined(__NVCC__)
VLOG(1) << "using fbgemm::GenerateSparseAdaGrad in RowWiseSparseAdagradOp";
if (block_size != last_block_size_) {
last_block_size_ = block_size;
if (std::is_same::value) {
kernel_i32_ = fbgemm::GenerateSparseAdaGrad(
block_size,
/*rowwise=*/true,
/*prefetch=*/16,
weight_decay_ != 0.0f);
} else {
CAFFE_ENFORCE((std::is_same::value));
kernel_i64_ = fbgemm::GenerateSparseAdaGrad(
block_size,
/*rowwise=*/true,
/*prefetch=*/16,
weight_decay_ != 0.0f);
}
}
int num_rows_processed;
if (std::is_same::value) {
num_rows_processed = kernel_i32_(
n,
Input(PARAM).numel(),
param,
gradIn,
moment,
reinterpret_cast(indices),
epsilon_,
lr[0],
weight_decay_,
(counter_halflife_ > 0) ? count : nullptr,
counter_halflife_);
} else {
num_rows_processed = kernel_i64_(
n,
Input(PARAM).numel(),
param,
gradIn,
moment,
reinterpret_cast(indices),
epsilon_,
lr[0],
weight_decay_,
(counter_halflife_ > 0) ? count : nullptr,
counter_halflife_);
}
if (num_rows_processed < n) {
// Enforce:
// access within range
CAFFE_ENFORCE_GE(
Input(PARAM).numel(),
(indices[num_rows_processed] + 1) * block_size,
this->debug_def().input(PARAM),
", out of bound, idx:",
indices[num_rows_processed],
" for input i:",
num_rows_processed,
" and block size:",
block_size,
" max size:",
Input(PARAM).numel());
return false;
} else {
return true;
}
#else
VLOG(1) << "using plain adagrad updates in RowWiseSparseAdagradOp";
for (auto i = 0; i < n; ++i) {
auto idx = indices[i];
float freq = (count[idx] > 0 && counter_halflife_ > 0)
? counter_halflife_ / count[idx]
: 1.0;
if (block_size == 1) {
float gi = std::fma(weight_decay_ * freq, param[idx], gradIn[i]);
float hi = moment[idx] = moment[idx] + gi * gi;
param[idx] = param[idx] + lr[0] * gi / (std::sqrt(hi) + epsilon_);
} else {
auto offsetI = i * block_size;
auto offsetIdx = idx * block_size;
#ifndef NDEBUG
CAFFE_ENFORCE_GE(
Input(PARAM).numel(),
block_size + offsetIdx,
this->debug_def().input(PARAM),
", out of bound, idx:",
idx,
" for input i:",
i,
" and block size:",
block_size);
CAFFE_ENFORCE_GE(
Input(GRAD).numel(),
block_size + offsetI,
this->debug_def().input(GRAD),
", out of bound idx, idx:",
idx,
" for input i:",
i);
#endif
float* w = param + offsetIdx;
const float* g = gradIn + offsetI;
float* h = moment + idx;
float hs = 0.;
for (auto j = 0; j < block_size; ++j) {
float gj = std::fma(weight_decay_ * freq, w[j], g[j]);
hs += gj * gj;
}
float hi = h[0] = h[0] + hs / block_size;
float step = lr[0] / (std::sqrt(hi) + epsilon_);
for (auto j = 0; j < block_size; ++j) {
float gj = std::fma(weight_decay_ * freq, w[j], g[j]);
w[j] = w[j] + gj * step;
}
}
}
return true;
#endif // !USE_FBGEMM
}
protected:
float epsilon_;
const float weight_decay_;
const int64_t counter_halflife_;
#if defined(USE_FBGEMM) && !defined(__NVCC__)
fbgemm::SparseAdaGradSignature::Type kernel_i32_;
fbgemm::SparseAdaGradSignature::Type kernel_i64_;
std::int64_t last_block_size_{-1};
#endif
INPUT_TAGS(PARAM, MOMENT_1, INDICES, GRAD, LR, COUNTER);
OUTPUT_TAGS(OUTPUT_PARAM, OUTPUT_MOMENT_1);
};
} // namespace caffe2