pytorch

1
# @package regularizer_context
2
# Module caffe2.python.normalizer_context
3

4

5

6

7

8
from caffe2.python import context
9
from caffe2.python.modifier_context import (
10
    ModifierContext, UseModifierBase)
11

12

13
class NormalizerContext(ModifierContext, context.DefaultManaged):
14
    """
15
    provide context to allow param_info to have different normalizers
16
    """
17

18
    def has_normalizer(self, name):
19
        return self._has_modifier(name)
20

21
    def get_normalizer(self, name):
22
        assert self.has_normalizer(name), (
23
            "{} normalizer is not provided!".format(name))
24
        return self._get_modifier(name)
25

26

27
class UseNormalizer(UseModifierBase):
28
    '''
29
    context class to allow setting the current context.
30
    Example usage with layer:
31
        normalizers = {'norm1': norm1, 'norm2': norm2}
32
        with UseNormalizer(normalizers):
33
            norm = NormalizerContext.current().get_normalizer('norm1')
34
            layer(norm=norm)
35
    '''
36
    def _context_class(self):
37
        return NormalizerContext
38