normalizer.py 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. # @package optimizer
  2. # Module caffe2.python.normalizer
  3. class Normalizer(object):
  4. def __init__(self):
  5. pass
  6. """
  7. Adds normalization to train_net for given parameter. Its factor ahead of
  8. regularization is given when initialization.
  9. The param should be a BlobReference.
  10. """
  11. def __call__(self, net, param):
  12. return self._run(net, param)
  13. def _run(self, net, param):
  14. raise Exception("Not Impelemented")
  15. class BatchNormalizer(Normalizer):
  16. def __init__(self, momentum, scale_init_value=1.0):
  17. super(BatchNormalizer, self).__init__()
  18. self._momentum = float(momentum)
  19. self._scale_init_value = float(scale_init_value)
  20. def _run(self, layer_model, param):
  21. return layer_model.BatchNormalization(
  22. param, momentum=self._momentum, scale_init_value=self._scale_init_value
  23. )
  24. class LayerNormalizer(Normalizer):
  25. def __init__(self, epsilon, use_layer_norm_op=True, scale_init_value=1.0):
  26. super(LayerNormalizer, self).__init__()
  27. self._epsilon = float(epsilon)
  28. self._use_layer_norm_op = use_layer_norm_op
  29. self._scale_init_value = float(scale_init_value)
  30. def _run(self, layer_model, param):
  31. return layer_model.LayerNormalization(
  32. param, epsilon=self._epsilon, use_layer_norm_op=self._use_layer_norm_op, scale_init_value=self._scale_init_value
  33. )