/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/weight_scale_op.h (2552B)
/** * Copyright (c) 2016-present, Facebook, Inc. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ #pragma once #include "caffe2/core/operator.h" #include #include namespace caffe2 { template void weight_scale_update( int N, const T* w, const T scale, int64_t iter, int64_t stepsize, int64_t update_upper_bound, T* nw, Context* context) { const auto w_size = N * sizeof(float); if (iter % stepsize != 0 || iter >= update_upper_bound) { memcpy(nw, w, w_size); return; } // perform the weight scaling caffe2::math::Scale(N, scale, w, nw, context); } template class WeightScaleOp final : public Operator { public: USE_OPERATOR_CONTEXT_FUNCTIONS; WeightScaleOp(const OperatorDef& operator_def, Workspace* ws) : Operator(operator_def, ws), stepsize_(OperatorBase::GetSingleArgument( "stepsize", std::numeric_limits::max())), update_upper_bound_(OperatorBase::GetSingleArgument( "upper_bound_iter", std::numeric_limits::max())), scale_(this->template GetSingleArgument("scale", 1.0f)) {} bool RunOnDevice() override { Output(OUTPUT_WEIGHTS)->ResizeLike(Input(WEIGHTS)); return DispatchHelper>::call(this, Input(WEIGHTS)); } template bool DoRunWithType() { const auto iter = OperatorBase::Input(ITER, CPU).template data()[0] + 1; weight_scale_update( Input(WEIGHTS).size(), Input(WEIGHTS).template data(), scale_, iter, stepsize_, update_upper_bound_, Output(OUTPUT_WEIGHTS)->template mutable_data(), &context_); return true; } protected: int64_t stepsize_; int64_t update_upper_bound_; float scale_; INPUT_TAGS(WEIGHTS, ITER); OUTPUT_TAGS(OUTPUT_WEIGHTS); }; } // namespace caffe2