/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/fp32_momentum_sgd_op.h (2015B)
#pragma once #include "caffe2/core/operator.h" #include "caffe2/core/timer.h" namespace caffe2 { template void fp32_momentum_sgd_update( int N, const float* g, const float* m, float* ng, float* nm, const float* lr, float momentum, bool nesterov, float weight_decay, float* param, Context* /*context*/) {} template class FP32MomentumSGDUpdateOp final : public Operator { public: USE_OPERATOR_CONTEXT_FUNCTIONS; FP32MomentumSGDUpdateOp(const OperatorDef& operator_def, Workspace* ws) : Operator(operator_def, ws), momentum_(this->template GetSingleArgument("momentum", 0.0)), weight_decay_( this->template GetSingleArgument("weight_decay", 0.0)), nesterov_(this->template GetSingleArgument("nesterov", 0)) {} bool RunOnDevice() override { auto device_type = Context::GetDeviceType(); // Iter live on the CPU CAFFE_ENFORCE(OperatorBase::InputIsTensorType(GRAD, device_type)); CAFFE_ENFORCE(OperatorBase::InputIsTensorType(MOMENTUM, device_type)); CAFFE_ENFORCE(Input(LR).size() == 1); CAFFE_ENFORCE(Input(GRAD).size() == Input(MOMENTUM).size()); Output(OUTPUT_GRAD)->ResizeLike(Input(GRAD)); Output(OUTPUT_MOMENTUM)->ResizeLike(Input(MOMENTUM)); fp32_momentum_sgd_update( Input(GRAD).size(), Input(GRAD).template data(), Input(MOMENTUM).template data(), Output(OUTPUT_GRAD)->template mutable_data(), Output(OUTPUT_MOMENTUM)->template mutable_data(), Input(LR).template data(), momentum_, nesterov_, weight_decay_, Output(OUTPUT_PARAM)->template mutable_data(), &context_); return true; } protected: float momentum_{0.9}; float weight_decay_{0.0}; bool nesterov_; INPUT_TAGS(GRAD, MOMENTUM, LR, PARAM); OUTPUT_TAGS(OUTPUT_GRAD, OUTPUT_MOMENTUM, OUTPUT_PARAM); }; }