Caffe2 - Python API
A deep learning, cross platform ML framework
optimizer_context.py
1 ## @package optimizer_context
2 # Module caffe2.python.optimizer_context
3 from __future__ import absolute_import
4 from __future__ import division
5 from __future__ import print_function
6 from __future__ import unicode_literals
7 
8 from caffe2.python import context
10  ModifierContext, UseModifierBase)
11 
12 
13 DEFAULT_OPTIM = 'DEFAULT'
14 
15 
16 @context.define_context(allow_default=True)
17 class OptimizerContext(ModifierContext):
18  """
19  provide context to allow param_info to have different optimizers
20  """
21 
22  def has_optimizer(self, name):
23  return self._has_modifier(name)
24 
25  def get_optimizer(self, name):
26  assert self.has_optimizer(name), (
27  "{} optimizer is not provided!".format(name))
28  return self._get_modifier(name)
29 
30 
31 class UseOptimizer(UseModifierBase):
32  '''
33  context class to allow setting the current context.
34  Example usage with brew:
35  - with UseOptimizer(optim):
36  brew.func
37  - with UseOptimizer({'WEIGHT': weight_optim}):
38  brew.func
39  - with UseOptimizer({'DEFAULT': optim, 'BIAS': bias_optim,
40  'WEIGHT': weight_optim}):
41  brew.func
42  - with UseOptimizer(optim1):
43  brew.func
44  with UseOptimizer(optim2):
45  brew.func
46 
47  Example useage with layer:
48  optimizers = {'optim1': optim1, 'optim2': optim2}
49  with Optimizers(optimizers):
50  optim = OptimizerContext.current().get_optimizer('optim1')
51  layer(optim=optim)
52  '''
53  def _context_class(self):
54  return OptimizerContext