/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/sgd
NameSizeModeActions
adadelta_op.h58520644editdlrm
adagrad_fused.h166320644editdlrm
adagrad_op.h179930644editdlrm
adam_op.h261650644editdlrm
clip_tensor_op.h18530644editdlrm
decay_adagrad_op.h32150644editdlrm
fp16_momentum_sgd_op.h22920644editdlrm
fp32_momentum_sgd_op.h20150644editdlrm
ftrl_op.h22190644editdlrm
gftrl_op.h10280644editdlrm
iter_op.h33790644editdlrm
lars_op.h24790644editdlrm
learning_rate_adaption_op.h19740644editdlrm
learning_rate_functors.h151350644editdlrm
learning_rate_op.h126960644editdlrm
math_lp.h6550644editdlrm
momentum_sgd_op.h60350644editdlrm
rmsprop_op.h19800644editdlrm
rowwise_adagrad_fused.h295950644editdlrm
rowwise_counter.h20830644editdlrm
storm_op.h59150644editdlrm
weight_scale_op.h25520644editdlrm
wngrad_op.h72720644editdlrm
yellowfin_op.h102840644editdlrm
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/sgd/lars_op.h (2479B)
#ifndef CAFFE2_OPERATORS_LARS_OP_H_ #define CAFFE2_OPERATORS_LARS_OP_H_ #include "caffe2/core/context.h" #include "caffe2/core/logging.h" #include "caffe2/core/operator.h" #include "caffe2/utils/math.h" namespace caffe2 { template class LarsOp final : public Operator { public: USE_OPERATOR_CONTEXT_FUNCTIONS; LarsOp(const OperatorDef& operator_def, Workspace* ws) : Operator(operator_def, ws), offset_(this->template GetSingleArgument("offset", 0.5)), lr_min_(this->template GetSingleArgument("lr_min", 0.02)) {} bool RunOnDevice() override { auto& X = Input(0); auto& dX = Input(1); CAFFE_ENFORCE( dX.numel() == X.numel(), "Gradient size doesn't match parameter size."); CAFFE_ENFORCE_GE(offset_, 0); CAFFE_ENFORCE_GE(lr_min_, 0); auto& wd = Input(2); auto& trust = Input(3); auto& lr_max = Input(4); auto* lr_rescaled = Output(0, vector{1}, at::dtype()); ReinitializeTensor(&X_norm_tensor_, {1}, at::dtype().device(Context::GetDeviceType())); T* X_norm_ = X_norm_tensor_.template mutable_data(); ReinitializeTensor(&dX_norm_tensor_, {1}, at::dtype().device(Context::GetDeviceType())); T* dX_norm_ = dX_norm_tensor_.template mutable_data(); ComputeNorms( dX.numel(), X.template data(), dX.template data(), X_norm_, dX_norm_); ComputeLearningRate( wd.template data(), trust.template data(), lr_max.template data(), offset_, lr_min_, X_norm_, dX_norm_, lr_rescaled->template mutable_data()); return true; } private: // Compute the l2 norm of X_data and dX_data void ComputeNorms( int64_t N, const T* X_data, const T* dX_data, T* X_norm, T* dX_norm) { math::SumSqr(N, X_data, X_norm, &context_); math::Sqrt(1, X_norm, X_norm, &context_); math::SumSqr(N, dX_data, dX_norm, &context_); math::Sqrt(1, dX_norm, dX_norm, &context_); } // Compute the learning rate and apply clipping void ComputeLearningRate( const T* wd, const T* trust, const T* lr_max, T offset, T lr_min, T* X_norm, T* dX_norm, T* lr_rescaled); T offset_; T lr_min_; Tensor X_norm_tensor_; Tensor dX_norm_tensor_; }; } // namespace caffe2 #endif // CAFFE2_OPERATORS_LARS_OP_H_