1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
|
#pragma once
#include "caffe2/core/operator.h"
#include "caffe2/core/timer.h"
namespace caffe2 {
template <class Context>
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 <typename T, class Context>
class FP32MomentumSGDUpdateOp final : public Operator<Context> {
public:
USE_OPERATOR_CONTEXT_FUNCTIONS;
FP32MomentumSGDUpdateOp(const OperatorDef& operator_def, Workspace* ws)
: Operator<Context>(operator_def, ws),
momentum_(this->template GetSingleArgument<float>("momentum", 0.0)),
weight_decay_(
this->template GetSingleArgument<float>("weight_decay", 0.0)),
nesterov_(this->template GetSingleArgument<int>("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<Context>(
Input(GRAD).size(),
Input(GRAD).template data<T>(),
Input(MOMENTUM).template data<T>(),
Output(OUTPUT_GRAD)->template mutable_data<T>(),
Output(OUTPUT_MOMENTUM)->template mutable_data<T>(),
Input(LR).template data<float>(),
momentum_,
nesterov_,
weight_decay_,
Output(OUTPUT_PARAM)->template mutable_data<T>(),
&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);
};
}
|
