1
## @package optimizer_context
2
# Module caffe2.python.optimizer_context
8
from caffe2.python import context
9
from caffe2.python.modifier_context import (
10
ModifierContext, UseModifierBase)
13
DEFAULT_OPTIM = 'DEFAULT'
16
class OptimizerContext(ModifierContext, context.DefaultManaged):
18
provide context to allow param_info to have different optimizers
21
def has_optimizer(self, name):
22
return self._has_modifier(name)
24
def get_optimizer(self, name):
25
assert self.has_optimizer(name), (
26
"{} optimizer is not provided!".format(name))
27
return self._get_modifier(name)
30
class UseOptimizer(UseModifierBase):
32
context class to allow setting the current context.
33
Example usage with brew:
34
- with UseOptimizer(optim):
36
- with UseOptimizer({'WEIGHT': weight_optim}):
38
- with UseOptimizer({'DEFAULT': optim, 'BIAS': bias_optim,
39
'WEIGHT': weight_optim}):
41
- with UseOptimizer(optim1):
43
with UseOptimizer(optim2):
46
Example usage with layer:
47
optimizers = {'optim1': optim1, 'optim2': optim2}
48
with Optimizers(optimizers):
49
optim = OptimizerContext.current().get_optimizer('optim1')
52
def _context_class(self):
53
return OptimizerContext