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
67
68
69
70
71
72
73
74
75
76
|
// (c) Facebook, Inc. and its affiliates. Confidential and proprietary.
// This is a very basic knob implementation that purely uses command line flags.
// This can be replaced with a more sophisticated implementation for use in
// other production environments.
#include <map>
#include <c10/util/string_view.h>
#include <c10/util/Flags.h>
#include "caffe2/utils/knobs.h"
namespace caffe2 {
namespace detail {
// Get the map of knob names to pointers to their command-line controlled
// boolean value.
std::map<c10::string_view, bool*>& getRegisteredKnobs() {
// It's safe to store the keys as string_view, since DEFINE_KNOB() ensures
// that these views always point to string literals.
static std::map<c10::string_view, bool*> registeredKnobs;
return registeredKnobs;
}
} // namespace detail
bool CheckKnob(c10::string_view name) {
const auto& knobs = detail::getRegisteredKnobs();
auto iter = knobs.find(name);
if (iter == knobs.end()) {
throw std::invalid_argument(
"attempted to check unknown knob \"" + std::string(name) + "\"");
}
return *iter->second;
}
namespace {
class RegisterKnob {
public:
RegisterKnob(c10::string_view name, bool* cmdlineFlag) {
auto ret = caffe2::detail::getRegisteredKnobs().emplace(name, cmdlineFlag);
if (!ret.second) {
throw std::runtime_error("duplicate knob name: " + std::string(name));
}
}
};
} // namespace
} // namespace caffe2
/**
* Define a knob.
*
* This will define a --caffe2_knob_<name> command line flag to control the
* knob.
*
* The knob can be checked in code by calling CheckKnob(name)
* or CheckKnob<check_fn_name>()
*/
#define DEFINE_KNOB(name, check_fn_name, default_value, docstring) \
C10_DEFINE_bool(caffe2_knob_##name, default_value, docstring); \
namespace caffe2 { \
bool CheckKnob##check_fn_name() { \
return FLAGS_caffe2_knob_##name; \
} \
} \
static caffe2::RegisterKnob _knob_##name(#name, &FLAGS_caffe2_knob_##name)
/*
* Definitions of well-known knobs.
*/
DEFINE_KNOB(
example_knob,
ExampleKnob,
false,
"An example knob, mainly intended for use in unit tests");
|
