parameter_info.py 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. from caffe2.python import core
  2. import numpy as np
  3. class ParameterTags(object):
  4. BIAS = 'BIAS'
  5. WEIGHT = 'WEIGHT'
  6. COMPUTED_PARAM = 'COMPUTED_PARAM'
  7. class ParameterInfo(object):
  8. def __init__(
  9. self, param_id, param, key=None, shape=None, length=None,
  10. grad=None, blob_copy=None):
  11. assert isinstance(param, core.BlobReference)
  12. self.param_id = param_id
  13. self.name = str(param)
  14. self.blob = param
  15. self.key = key
  16. self.shape = shape
  17. self.size = None if shape is None else np.prod(shape)
  18. self.length = max(1, length if length is not None else 1)
  19. self.grad = grad
  20. self._cloned_init_net = None
  21. # Optionally store equivalent copies of the blob
  22. # in different precisions (i.e. half and float copies)
  23. # stored as a dict of TensorProto.DataType -> BlobReference
  24. self.blob_copy = blob_copy
  25. # each param_info can have its own optimizer. It can be set within
  26. # OptimizerContext (caffe2/python/optimizer.py)
  27. self._optimizer = None
  28. @property
  29. def parameter(self):
  30. return self.blob
  31. @property
  32. def optimizer(self):
  33. return self._optimizer
  34. @optimizer.setter
  35. def optimizer(self, value):
  36. assert self._optimizer is None, "optimizer has already been set"
  37. self._optimizer = value
  38. def __str__(self):
  39. return self.name