Caffe2 - C++ API
A deep learning, cross platform ML framework
ftrl_op.h
1 
17 #pragma once
18 
19 #include "caffe2/core/operator.h"
20 
21 namespace caffe2 {
22 
23 template <typename T>
24 struct FtrlParams {
25  explicit FtrlParams(OperatorBase* op)
26  : alphaInv(1.0 / op->GetSingleArgument<float>("alpha", 0.005f)),
27  beta(op->GetSingleArgument<float>("beta", 1.0f)),
28  lambda1(op->GetSingleArgument<float>("lambda1", 0.001f)),
29  lambda2(op->GetSingleArgument<float>("lambda2", 0.001f)) {}
30  T alphaInv;
31  T beta;
32  T lambda1;
33  T lambda2;
34 };
35 
36 // TODO(dzhulgakov): implement GPU version if necessary
37 template <typename T, class Context>
38 class FtrlOp final : public Operator<Context> {
39  public:
40  USE_OPERATOR_CONTEXT_FUNCTIONS;
41  FtrlOp(const OperatorDef& operator_def, Workspace* ws)
42  : Operator<Context>(operator_def, ws), params_(this) {
43  CAFFE_ENFORCE(
44  !HasArgument("alpha") || ALPHA >= InputSize(),
45  "Cannot specify alpha by both input and argument");
46  }
47  bool RunOnDevice() override;
48 
49  protected:
50  FtrlParams<T> params_;
51  INPUT_TAGS(VAR, N_Z, GRAD, ALPHA);
52  OUTPUT_TAGS(OUTPUT_VAR, OUTPUT_N_Z);
53 };
54 
55 template <typename T>
56 class SparseFtrlOp final : public Operator<CPUContext> {
57  public:
58  SparseFtrlOp(const OperatorDef& operator_def, Workspace* ws)
59  : Operator<CPUContext>(operator_def, ws), params_(this) {
60  CAFFE_ENFORCE(
61  !HasArgument("alpha") || ALPHA >= InputSize(),
62  "Cannot specify alpha by both input and argument");
63  }
64 
65  bool RunOnDevice() override {
66  // run time learning rate override
67  if (ALPHA < InputSize()) {
68  CAFFE_ENFORCE_EQ(Input(ALPHA).size(), 1, "alpha should be real-valued");
69  params_.alphaInv = 1.0 / *(Input(ALPHA).template data<T>());
70  }
71  // Use run-time polymorphism
72  auto& indices = Input(INDICES);
73  if (indices.template IsType<int32_t>()) {
74  DoRun<int32_t>();
75  } else if (indices.template IsType<int64_t>()) {
76  DoRun<int64_t>();
77  } else {
78  LOG(FATAL) << "Unsupported type of INDICES in SparseFtrlOp: "
79  << indices.meta().name();
80  }
81  return true;
82  }
83 
84  protected:
85  FtrlParams<T> params_;
86  INPUT_TAGS(VAR, N_Z, INDICES, GRAD, ALPHA);
87  OUTPUT_TAGS(OUTPUT_VAR, OUTPUT_N_Z);
88 
89  private:
90  template <typename SIndex>
91  void DoRun();
92 };
93 
94 }
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.